#!/usr/bin/env python3

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

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

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: 10,
    200: 40
}

CONSIST_RAD = np.deg2rad(20)

EULER_LAT = 0.0
EULER_LON = -59.0

CHUNK_ROWS = 1500
PAD = 60   # must exceed largest sigma

# -----------------------------
# 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.stack([x, y, z], axis=-1)

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)
    l1 = 0.5 * (Jxx + Jyy + tmp)
    l2 = 0.5 * (Jxx + Jyy - tmp)

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

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

    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)

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

    ny, nx = h.shape

    # Optional downsample (recommended first run)
    # h = h[::2, ::2]
    # 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)

    # Output arrays (on disk)
    M_plus = np.memmap("mach_plus.dat", dtype="float32", mode="w+", shape=(ny, nx))
    M_minus = np.memmap("mach_minus.dat", dtype="float32", mode="w+", shape=(ny, nx))

    valid_plus = np.memmap("valid_plus.dat", dtype="bool", mode="w+", shape=(ny, nx))
    valid_minus = np.memmap("valid_minus.dat", dtype="bool", mode="w+", shape=(ny, nx))

    # Euler vector
    omega = latlon_to_cart(EULER_LAT, EULER_LON)

    for y0 in range(0, ny, CHUNK_ROWS):
        y1 = min(y0 + CHUNK_ROWS, ny)

        print(f"Processing rows {y0}:{y1}")

        # padded window
        y0p = max(0, y0 - PAD)
        y1p = min(ny, y1 + PAD)

        h_pad = h[y0p:y1p, :]

        lat_chunk = lat[y0p:y1p]
        lon_grid, lat_grid = np.meshgrid(lon, lat_chunk)

        # gradients
        gx, gy = compute_gradients(h_pad, dx, dy)
        gx = gx.astype(np.float32)
        gy = gy.astype(np.float32)

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

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

        # crop padding
        sl = slice(y0 - y0p, y1 - y0p)

        vx50 = vx50[sl]
        vy50 = vy50[sl]
        vx200 = vx200[sl]
        vy200 = vy200[sl]
        C50 = C50[sl]
        grad_mag = grad_mag[sl]
        h_chunk = h[y0:y1]

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

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

        # coherence
        coh_ok = (ocean & (C50 >= COH_OCEAN)) | (land & (C50 >= COH_LAND))

        # gradient threshold
        if np.any(ocean):
            thr_ocean = np.percentile(grad_mag[ocean], 20)
        else:
            thr_ocean = 0

        if np.any(land):
            thr_land = np.percentile(grad_mag[land], 20)
        else:
            thr_land = 0

        grad_ok = (ocean & (grad_mag >= thr_ocean)) | (land & (grad_mag >= thr_land))

        base_mask = scale_ok & coh_ok & grad_ok

        # flow field
        r = latlon_to_cart(lat_grid[sl], lon_grid[sl])
        F = np.cross(omega, r)
        norm = np.linalg.norm(F, axis=-1, keepdims=True) + 1e-12
        F /= norm

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

        # angles
        mu_p = angle_between(vx50, vy50, Fx, Fy)
        mu_m = angle_between(vx50, vy50, -Fx, -Fy)

        valid_p = base_mask & (mu_p >= MU_MIN) & (mu_p <= MU_MAX)
        valid_m = base_mask & (mu_m >= MU_MIN) & (mu_m <= MU_MAX)

        M_p = np.zeros_like(mu_p, dtype=np.float32)
        M_m = np.zeros_like(mu_m, dtype=np.float32)

        M_p[valid_p] = 1.0 / np.sin(mu_p[valid_p])
        M_m[valid_m] = 1.0 / np.sin(mu_m[valid_m])

        # write to disk
        M_plus[y0:y1] = M_p
        M_minus[y0:y1] = M_m
        valid_plus[y0:y1] = valid_p
        valid_minus[y0:y1] = valid_m

    print("Done.")

if __name__ == "__main__":
    main()