#!/usr/bin/env python3

import numpy as np
import rasterio

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

STEP_KM = 50
N_STEPS = 200
N_SEEDS = 2000   # slightly reduced for speed

R_EARTH = 6371.0

# refined search window (equatorial band)
LAT_RANGE = (-85, 85)
LON_RANGE = (-180, 180)

LAT_STEP = 5    # 5° resolution
LON_STEP = 5

# -----------------------------
# UTILITIES
# -----------------------------
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(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):
    theta = step_km / R_EARTH
    axis = normalize(np.cross(r, direction))

    if np.linalg.norm(axis) < 1e-10:
        return r

    return normalize(
        r*np.cos(theta)
        + np.cross(axis, r)*np.sin(theta)
        + axis*np.dot(axis, r)*(1-np.cos(theta))
    )

def sample_field(M, valid, lat, lon, lat_grid, lon_grid):
    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 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))

# -----------------------------
# SEEDS
# -----------------------------
print("Sampling seeds...")

valid_idx = np.argwhere(valid)
idx = np.random.choice(len(valid_idx), size=N_SEEDS, replace=False)
seeds = valid_idx[idx]

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

results = []

print("Refined sweep...")

for plat in lat_vals:
    for plon in lon_vals:

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

                F = normalize(np.cross(omega, r))

                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)

                r = step_on_sphere(r, F, STEP_KM)

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

        mean_var = np.mean(variances) if variances else np.nan

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

        print(f"({plat},{plon}) -> {mean_var:.6f}")

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

print("Saved refined_streamline_sweep.npy")


import numpy as np
import matplotlib.pyplot as plt

data = np.load("refined_streamline_sweep.npy")

lat = data[:,0]
lon = data[:,1]
var = data[:,2]

plt.figure(figsize=(12,6))
sc = plt.scatter(lon, lat, c=var, cmap="plasma")
plt.colorbar(sc, label="Streamline Variance")
plt.xlabel("Longitude")
plt.ylabel("Latitude")
plt.title("Refined Euler Pole Sweep (Streamline Variance)")
plt.show()