#!/usr/bin/env python3

import numpy as np
import rasterio
from concurrent.futures import ProcessPoolExecutor
import multiprocessing as mp
from multiprocessing import shared_memory
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

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

# -----------------------------
# GLOBAL SHARED HANDLES
# -----------------------------
shared = {}

# -----------------------------
# INIT WORKER
# -----------------------------
def init_worker(meta):
    global shared

    shared["handles"] = {}

    # attach and KEEP HANDLES
    shared["handles"]["M"] = shared_memory.SharedMemory(name=meta["M"])
    shared["handles"]["valid"] = shared_memory.SharedMemory(name=meta["valid"])
    shared["handles"]["lat"] = shared_memory.SharedMemory(name=meta["lat"])
    shared["handles"]["lon"] = shared_memory.SharedMemory(name=meta["lon"])

    ny, nx = meta["shape"]

    shared["ny"] = ny
    shared["nx"] = nx

    # SAFE reconstruction
    shared["M"] = np.frombuffer(
        shared["handles"]["M"].buf, dtype=np.float32
    ).reshape(ny, nx)

    shared["valid"] = np.frombuffer(
        shared["handles"]["valid"].buf, dtype=np.bool_
    ).reshape(ny, nx)

    shared["lat"] = np.frombuffer(
        shared["handles"]["lat"].buf, dtype=np.float32
    )

    shared["lon"] = np.frombuffer(
        shared["handles"]["lon"].buf, dtype=np.float32
    )

    shared["counter"] = meta["counter"]

# -----------------------------
# FAST INDEXING
# -----------------------------
def latlon_to_index(lat, lon):
    ny, nx = shared["ny"], shared["nx"]

    iy = ((90.0 - lat) / 180.0 * (ny - 1)).astype(np.int32)
    ix = ((lon + 180.0) / 360.0 * nx).astype(np.int32)

    iy = np.clip(iy, 0, ny - 1)
    ix = np.clip(ix, 0, nx - 1)

    return iy, ix

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

def cart_to_latlon(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(v):
    return v / (np.linalg.norm(v, axis=1, keepdims=True) + 1e-12)

def step_on_sphere(r, direction):
    theta = STEP_KM / R_EARTH

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

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

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

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

# -----------------------------
# WORKER
# -----------------------------
def worker(task_chunk):
    results = []

    M = shared["M"]
    valid = shared["valid"]

    for plat, plon, seeds_cart in task_chunk:

        omega = latlon_to_cart(plat, plon)

        r = seeds_cart.copy()
        samples = []

        for _ in range(N_STEPS):
            F = np.cross(r, omega)
            F = normalize(F)

            lat_s, lon_s = cart_to_latlon(r)
            iy, ix = latlon_to_index(lat_s, lon_s)

            vals = M[iy, ix].copy()
            mask = valid[iy, ix]
            vals[~mask] = np.nan

            samples.append(vals)
            r = step_on_sphere(r, F)

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

        results.append((plat, plon, float(np.nanmean(stds))))

        # progress update
        with shared["counter"].get_lock():
            shared["counter"].value += 1

    return results

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

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

    ny, nx = h.shape

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

    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))

    # -----------------------------
    # SHARED MEMORY SETUP
    # -----------------------------
    print("Creating shared memory...")

    shm_M = shared_memory.SharedMemory(create=True, size=M.nbytes)
    shm_valid = shared_memory.SharedMemory(create=True, size=valid.nbytes)
    shm_lat = shared_memory.SharedMemory(create=True, size=lat_grid.nbytes)
    shm_lon = shared_memory.SharedMemory(create=True, size=lon_grid.nbytes)

    np.ndarray(M.shape, dtype=M.dtype, buffer=shm_M.buf)[:] = M
    np.ndarray(valid.shape, dtype=valid.dtype, buffer=shm_valid.buf)[:] = valid
    np.ndarray(lat_grid.shape, dtype=lat_grid.dtype, buffer=shm_lat.buf)[:] = lat_grid
    np.ndarray(lon_grid.shape, dtype=lon_grid.dtype, buffer=shm_lon.buf)[:] = lon_grid

    counter = mp.Value('i', 0)

    meta = {
        "shape": (ny, nx),
        "ny": ny,
        "nx": nx,
        "M": shm_M.name,
        "valid": shm_valid.name,
        "lat": shm_lat.name,
        "lon": shm_lon.name,
        "counter": counter
    }

    # -----------------------------
    # OCEAN SEEDS
    # -----------------------------
    ocean = (h <= -200)
    valid_ocean = valid & ocean

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

    lat0 = lat_grid[idx[sel][:,0]]
    lon0 = lon_grid[idx[sel][:,1]]
    seeds_cart = latlon_to_cart(lat0, lon0)

    # -----------------------------
    # TASKS
    # -----------------------------
    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) for plat in lat_vals for plon in lon_vals]

    # chunk tasks
    chunk_size = 8
    task_chunks = [tasks[i:i+chunk_size] for i in range(0, len(tasks), chunk_size)]

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

    start = time.time()

    results = []

    with ProcessPoolExecutor(
        max_workers=N_WORKERS,
        initializer=init_worker,
        initargs=(meta,)
    ) as ex:

        futures = [ex.submit(worker, chunk) for chunk in task_chunks]

        total = len(tasks)

        while any(not f.done() for f in futures):
            time.sleep(3)
            done = counter.value
            rate = done / (time.time() - start + 1e-6)
            eta = (total - done) / (rate + 1e-6)

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

        for f in futures:
            results.extend(f.result())

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

    print("Done. Results saved.")

    # cleanup
    shm_M.close(); shm_M.unlink()
    shm_valid.close(); shm_valid.unlink()
    shm_lat.close(); shm_lat.unlink()
    shm_lon.close(); shm_lon.unlink()

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