#!/usr/bin/env python3
"""
Global Mach Field Estimation from DEM using ECDO-derived flow

Pipeline:
- Load SRTM15+V2.tiff
- Build ocean/land masks
- Compute spherical gradients (approx in lat/lon)
- Structure tensor (multi-scale)
- Orientation extraction + coherence filtering
- Flow field from Euler pole
- Angle + Mach inversion

NOTE: This is a first-pass implementation (Stage 1–2 from plan)
"""

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

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

# Masks
OCEAN_MAX = -200.0
LAND_MIN = 200.0

# Coherence thresholds
COH_OCEAN = 0.20
COH_LAND = 0.35

# Scale (km) -> sigma (pixels) placeholder
SIGMA_PIXELS = {
    10: 2,
    50: 10,
    200: 40
}

# Euler pole (0N, 59W)
EULER_LAT = 0.0
EULER_LON = -59.0

# Angle bounds
MU_MIN = np.deg2rad(10)
MU_MAX = np.deg2rad(80)

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

    # eigenvector (dominant)
    vx = 2 * Jxy
    vy = Jyy - Jxx + tmp

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

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

    return vx, vy, coherence


def build_flow_field(lat_grid, lon_grid):
    r = latlon_to_cart(lat_grid, lon_grid)

    omega = latlon_to_cart(EULER_LAT, EULER_LON)

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

    return F, -F


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


# -----------------------------
# MAIN PIPELINE
# -----------------------------

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

    ny, nx = h.shape

    # Build lat/lon grids (approx)
    lon = np.linspace(-180, 180, nx)
    lat = np.linspace(90, -90, ny)
    lon_grid, lat_grid = np.meshgrid(lon, lat)

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

    print("Computing gradients...")
    gx, gy = compute_gradients(h, dx, dy)

    print("Structure tensor (50 km scale)...")
    sigma = SIGMA_PIXELS[50]
    Jxx, Jyy, Jxy = structure_tensor(gx, gy, sigma)

    vx, vy, C = eigen_orientation(Jxx, Jyy, Jxy)

    print("Building masks...")
    ocean_mask = (h <= OCEAN_MAX)
    land_mask = (h >= LAND_MIN)

    mask = np.zeros_like(h, dtype=bool)
    mask |= (ocean_mask & (C >= COH_OCEAN))
    mask |= (land_mask & (C >= COH_LAND))

    print("Building flow field...")
    Fp, Fm = build_flow_field(lat_grid, lon_grid)

    # Project flow into tangent plane (approx: ignore z)
    Fx = Fp[..., 0]
    Fy = Fp[..., 1]

    print("Computing angles...")
    mu = angle_between(vx, vy, Fx, Fy)

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

    print("Computing Mach field...")
    M = np.zeros_like(mu)
    M[valid] = 1.0 / np.sin(mu[valid])

    print("Summary:")
    print("Valid fraction:", valid.mean())
    print("Mean Mach:", np.mean(M[valid]))
    print("Std Mach:", np.std(M[valid]))

    np.save("mach_field.npy", M)
    np.save("mask.npy", valid)


if __name__ == "__main__":
    main()
