#!/usr/bin/env python3

import numpy as np
import rasterio

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

# streamline parameters
STEP_KM = 50
N_STEPS = 200
N_SEEDS = 2000

# Earth radius
R_EARTH = 6371.0

# candidate poles (add more if needed)
POLES = [
    (0.0, -59.0),   # your ECDO axis
    (0.0, 121.0),   # antipode
    (0.0, 0.0),
    (45.0, 0.0),
    (-45.0, 90.0),
    (60.0, -120.0),
]

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

def cart_to_latlon(v):
    x, y, z = v
    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) + 1e-12)

def step_on_sphere(r, direction, step_km):
    # convert step to radians
    theta = step_km / R_EARTH

    # rotate r along direction
    axis = normalize(np.cross(r, direction))
    if np.linalg.norm(axis) < 1e-10:
        return r

    # Rodrigues rotation
    r_new = (
        r * np.cos(theta)
        + np.cross(axis, r) * np.sin(theta)
        + axis * np.dot(axis, r) * (1 - np.cos(theta))
    )

    return normalize(r_new)

def sample_field(M, valid, lat, lon, lat_grid, lon_grid):
    # nearest neighbor lookup (fast, sufficient)
    iy = np.argmin(np.abs(lat_grid - lat))
    ix = np.argmin(np.abs(lon_grid - lon))

    if valid[iy, ix]:
        return M[iy, ix]
    return np.nan

# -----------------------------
# LOAD DATA
# -----------------------------
print("Loading DEM shape...")

with rasterio.open(DEM_FILE) as src:
    ny, nx = src.shape

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

print("Loading Stage 3 outputs...")

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

# -----------------------------
# SEED POINTS
# -----------------------------
print("Generating seed points...")

valid_indices = np.argwhere(valid)

# random subset
idx = np.random.choice(len(valid_indices), size=N_SEEDS, replace=False)
seeds = valid_indices[idx]

# -----------------------------
# STREAMLINE ANALYSIS
# -----------------------------
results = []

for plat, plon in POLES:

    print(f"Testing pole: ({plat}, {plon})")

    omega = latlon_to_cart(plat, plon)

    variances = []

    for iy, ix in seeds:

        lat0 = lat_grid[iy]
        lon0 = lon_grid[ix]

        r = latlon_to_cart(lat0, lon0)

        samples = []

        for _ in range(N_STEPS):

            # flow direction
            F = np.cross(omega, r)
            F = normalize(F)

            # sample Mach
            lat_s, lon_s = cart_to_latlon(r)
            val = sample_field(M, valid, lat_s, lon_s, lat_grid, lon_grid)

            if not np.isnan(val):
                samples.append(val)

            # advance
            r = step_on_sphere(r, F, STEP_KM)

        if len(samples) > 10:
            variances.append(np.std(samples))

    if len(variances) > 0:
        mean_var = np.mean(variances)
    else:
        mean_var = np.nan

    results.append([plat, plon, mean_var])

    print(f"  Mean streamline variance: {mean_var:.6f}")

# -----------------------------
# SAVE RESULTS
# -----------------------------
results = np.array(results)
np.save("streamline_results.npy", results)

print("Saved streamline_results.npy")