#!/usr/bin/env python3

import numpy as np
import rasterio
from concurrent.futures import ProcessPoolExecutor, as_completed
import multiprocessing as mp
import time

# -----------------------------
# CONFIG
# -----------------------------
DEM_FILE = "SRTM15+V2.tiff"

STEP_KM = 50
N_STEPS = 200
N_SEEDS = 2000

R_EARTH = 6371.0

LAT_RANGE = (-20, 20)
LON_RANGE = (-180, 180)

LAT_STEP = 5
LON_STEP = 5

# half cores
N_WORKERS = max(1, mp.cpu_count() // 2)

# -----------------------------
# GEOMETRY
# -----------------------------
def latlon_to_cart(lat, lon):
    lat = np.deg2rad(lat)
    lon = np.deg2rad(lon)
    return np.array([
        np.cos(lat)*np.cos(lon),
        np.cos(lat)*np.sin(lon),
        np.sin(lat)
    ])

def cart_to_latlon_batch(v):
    x, y, z = v[:,0], v[:,1], v[:,2]
    lat = np.rad2deg(np.arcsin(z))
    lon = np.rad2deg(np.arctan2(y, x))
    return lat, lon

def normalize_batch(v):
    norms = np.linalg.norm(v, axis=1, keepdims=True) + 1e-12
    return v / norms

def step_on_sphere_batch(r, direction, step_km):
    theta = step_km / R_EARTH

    axis = np.cross(r, direction)
    axis = normalize_batch(axis)

    cos_t = np.cos(theta)
    sin_t = np.sin(theta)

    dot = np.sum(axis * r, axis=1, keepdims=True)

    return normalize_batch(
        r * cos_t
        + np.cross(axis, r) * sin_t
        + axis * dot * (1 - cos_t)
    )

# -----------------------------
# SAMPLING
# -----------------------------
def sample_field_batch(M, valid, lat, lon, lat_grid, lon_grid):
    iy = np.abs(lat_grid[:, None] - lat[None, :]).argmin(axis=0)
    ix = np.abs(lon_grid[:, None] - lon[None, :]).argmin(axis=0)

    vals = M[iy, ix]
    mask = valid[iy, ix]

    vals[~mask] = np.nan
    return vals

# -----------------------------
# WORKER FUNCTION
# -----------------------------
def evaluate_pole(args):
    plat, plon, seeds_cart, M, valid, lat_grid, lon_grid = args

    omega = latlon_to_cart(plat, plon)

    r = seeds_cart.copy()
    all_samples = []

    for _ in range(N_STEPS):
        F = np.cross(np.tile(omega, (r.shape[0],1)), r)
        F = normalize_batch(F)

        lat_s, lon_s = cart_to_latlon_batch(r)
        vals = sample_field_batch(M, valid, lat_s, lon_s, lat_grid, lon_grid)

        print(f"{plat},{plong} : {stds}")

        all_samples.append(vals)

        r = step_on_sphere_batch(r, F, STEP_KM)

    samples = np.stack(all_samples, axis=1)
    stds = np.nanstd(samples, axis=1)

    return (plat, plon, float(np.nanmean(stds)))

# -----------------------------
# MAIN
# -----------------------------
def main():

    print("Loading DEM...")
    with rasterio.open(DEM_FILE) as src:
        h = src.read(1)

    ny, nx = h.shape

    lat_grid = np.linspace(90, -90, ny)
    lon_grid = np.linspace(-180, 180, nx, endpoint=False)

    print("Loading Mach field...")
    M = np.memmap("mach_plus.dat", dtype="float32", mode="r", shape=(ny, nx))
    valid = np.memmap("valid_plus.dat", dtype="bool", mode="r", shape=(ny, nx))

    # -----------------------------
    # OCEAN-ONLY SEEDS
    # -----------------------------
    print("Selecting ocean-only seeds...")

    ocean = (h <= -200)
    valid_ocean = valid & ocean

    idx = np.argwhere(valid_ocean)
    sel = np.random.choice(len(idx), size=N_SEEDS, replace=False)
    seeds_idx = idx[sel]

    lat0 = lat_grid[seeds_idx[:,0]]
    lon0 = lon_grid[seeds_idx[:,1]]

    seeds_cart = np.stack([
        np.cos(np.deg2rad(lat0))*np.cos(np.deg2rad(lon0)),
        np.cos(np.deg2rad(lat0))*np.sin(np.deg2rad(lon0)),
        np.sin(np.deg2rad(lat0))
    ], axis=1)

    # -----------------------------
    # POLE GRID
    # -----------------------------
    lat_vals = np.arange(LAT_RANGE[0], LAT_RANGE[1]+LAT_STEP, LAT_STEP)
    lon_vals = np.arange(LON_RANGE[0], LON_RANGE[1], LON_STEP)

    tasks = [(plat, plon, seeds_cart, M, valid, lat_grid, lon_grid)
             for plat in lat_vals for plon in lon_vals]

    total = len(tasks)
    print(f"Running {total} poles on {N_WORKERS} workers...")

    results = []
    start_time = time.time()

    with ProcessPoolExecutor(max_workers=N_WORKERS) as ex:
        futures = [ex.submit(evaluate_pole, t) for t in tasks]

        for i, f in enumerate(as_completed(futures), 1):
            results.append(f.result())

            # progress reporting
            if i % 10 == 0 or i == total:
                elapsed = time.time() - start_time
                rate = i / elapsed
                remaining = (total - i) / rate if rate > 0 else 0

                print(f"[{i}/{total}] "
                      f"{rate:.2f} poles/sec | "
                      f"ETA {remaining/60:.1f} min")

    results = np.array(results)
    np.save("refined_streamline_sweep_ocean.npy", results)

    print("Saved refined_streamline_sweep_ocean.npy")


# -----------------------------
# ENTRY POINT
# -----------------------------
if __name__ == "__main__":
    mp.set_start_method("spawn", force=True)
    main()