#!/usr/bin/env python3

import numpy as np
import rasterio
from scipy.ndimage import gaussian_filter

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

# Downsample factor (CRITICAL for speed)
DOWNSAMPLE = 4   # use 2–8 depending on speed

CHUNK_ROWS = 1000
PAD = 50

OCEAN_MAX = -200.0
LAND_MIN = 200.0

COH_OCEAN = 0.20
COH_LAND = 0.35

MU_MIN = np.deg2rad(10)
MU_MAX = np.deg2rad(80)

SIGMA_PIXELS = {
    50: 6,
    200: 20
}

CONSIST_RAD = np.deg2rad(20)

# Pole sampling resolution
LAT_STEPS = 18   # ~10° resolution
LON_STEPS = 36   # ~10° resolution

# -----------------------------
# 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 compute_gradients(h, dx, dy):
    gy, gx = np.gradient(h, dy, dx)
    return gx, gy

def structure_tensor(gx, gy, sigma):
    Jxx = gaussian_filter(gx * gx, sigma)
    Jyy = gaussian_filter(gy * gy, sigma)
    Jxy = gaussian_filter(gx * gy, sigma)
    return Jxx, Jyy, Jxy

def eigen_orientation(Jxx, Jyy, Jxy):
    tmp = np.sqrt((Jxx - Jyy)**2 + 4 * Jxy**2)

    vx = 2 * Jxy
    vy = Jyy - Jxx + tmp

    norm = np.sqrt(vx**2 + vy**2) + 1e-12
    vx /= norm
    vy /= norm

    l1 = 0.5 * (Jxx + Jyy + tmp)
    l2 = 0.5 * (Jxx + Jyy - tmp)

    C = (l1 - l2) / (l1 + l2 + 1e-12)

    return vx, vy, C

def angle_between(vx, vy, Fx, Fy):
    dot = np.abs(vx * Fx + vy * Fy)
    dot = np.clip(dot, -1, 1)
    return np.arccos(dot)

def angular_diff(vx1, vy1, vx2, vy2):
    dot = np.abs(vx1 * vx2 + vy1 * vy2)
    dot = np.clip(dot, -1, 1)
    return np.arccos(dot)

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

# downsample
h = h[::DOWNSAMPLE, ::DOWNSAMPLE]
ny, nx = h.shape

lon = np.linspace(-180, 180, nx, endpoint=False)
lat = np.linspace(90, -90, ny)

dx = np.deg2rad(360 / nx)
dy = np.deg2rad(180 / ny)

# -----------------------------
# PRECOMPUTE STRUCTURE (ONCE)
# -----------------------------
print("Precomputing structure tensors...")

gx, gy = compute_gradients(h, dx, dy)

# 50 km
Jxx, Jyy, Jxy = structure_tensor(gx, gy, SIGMA_PIXELS[50])
vx50, vy50, C50 = eigen_orientation(Jxx, Jyy, Jxy)

# 200 km
Jxx, Jyy, Jxy = structure_tensor(gx, gy, SIGMA_PIXELS[200])
vx200, vy200, _ = eigen_orientation(Jxx, Jyy, Jxy)

# masks
ocean = (h <= OCEAN_MAX)
land = (h >= LAND_MIN)

grad_mag = np.sqrt(gx**2 + gy**2)

# gradient thresholds
thr_ocean = np.percentile(grad_mag[ocean], 20) if np.any(ocean) else 0
thr_land  = np.percentile(grad_mag[land], 20) if np.any(land) else 0

grad_ok = (ocean & (grad_mag >= thr_ocean)) | (land & (grad_mag >= thr_land))
coh_ok = (ocean & (C50 >= COH_OCEAN)) | (land & (C50 >= COH_LAND))

dtheta = angular_diff(vx50, vy50, vx200, vy200)
scale_ok = dtheta <= CONSIST_RAD

base_mask = grad_ok & coh_ok & scale_ok

# -----------------------------
# SWEEP EULER POLES
# -----------------------------
results = []

lat_vals = np.linspace(-90, 90, LAT_STEPS)
lon_vals = np.linspace(-180, 180, LON_STEPS)

print("Sweeping Euler poles...")

for plat in lat_vals:
    for plon in lon_vals:

        omega = latlon_to_cart(plat, plon)

        # flow field
        lon_grid, lat_grid = np.meshgrid(lon, lat)
        r = np.stack([
            np.cos(np.deg2rad(lat_grid)) * np.cos(np.deg2rad(lon_grid)),
            np.cos(np.deg2rad(lat_grid)) * np.sin(np.deg2rad(lon_grid)),
            np.sin(np.deg2rad(lat_grid))
        ], axis=-1)

        F = np.cross(omega, r)
        norm = np.linalg.norm(F, axis=-1, keepdims=True) + 1e-12
        F /= norm

        Fx = F[..., 0]
        Fy = F[..., 1]

        mu = angle_between(vx50, vy50, Fx, Fy)

        valid = base_mask & (mu >= MU_MIN) & (mu <= MU_MAX)

        if np.any(valid):
            M = 1.0 / np.sin(mu[valid])
            results.append([
                plat,
                plon,
                M.mean(),
                M.std(),
                valid.mean()
            ])
        else:
            results.append([plat, plon, np.nan, np.nan, 0])

        print(f"Pole ({plat:.1f}, {plon:.1f}) done")

# -----------------------------
# SAVE RESULTS
# -----------------------------
results = np.array(results)

np.save("euler_sweep_results.npy", results)

print("Saved euler_sweep_results.npy")


import numpy as np
import matplotlib.pyplot as plt

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

lat = data[:,0]
lon = data[:,1]
meanM = data[:,2]
stdM = data[:,3]

plt.scatter(lon, lat, c=meanM, cmap="viridis")
plt.colorbar(label="Mean Mach")
plt.title("Euler Pole Sweep (Mean Mach)")
plt.show()

plt.scatter(lon, lat, c=stdM, cmap="plasma")
plt.colorbar(label="Std Mach")
plt.title("Euler Pole Sweep (Std Dev)")
plt.show()