#!/usr/bin/env python3
"""
Stage 6: Mode decomposition using TRUE shear-net eigenvectors

- Uses Net1 / Net2 directions from shear geometry
- Computes projection dominance of transport field
- Outputs diagnostic maps + distributions

Author: Craig Stone / ChatGPT assist
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.interpolate import griddata

# ============================================================
# CONFIG
# ============================================================

CSV_FILE = "mach.csv"   # semicolon-delimited
GRID_RES = 1.0          # degrees

EPS = 1e-12

# ============================================================
# LOAD M FIELD
# ============================================================

print("Loading M field...")

df = pd.read_csv(CSV_FILE, sep=";")

# Normalize column names
df.columns = [c.strip().lower() for c in df.columns]

lat = df["latitude"].values
lon = df["longitude"].values
M   = df["correlation"].values

# ============================================================
# BUILD REGULAR GRID
# ============================================================

print("Building grid...")

grid_lat = np.arange(-90, 91, GRID_RES)
grid_lon = np.arange(-180, 181, GRID_RES)

LON, LAT = np.meshgrid(grid_lon, grid_lat)

points = np.column_stack((lat, lon))
M_grid = griddata(points, M, (LAT, LON), method="linear")

# Fill NaNs (simple nearest fallback)
mask = np.isnan(M_grid)
M_grid[mask] = griddata(points, M, (LAT, LON), method="nearest")[mask]

# ============================================================
# COMPUTE TRANSPORT FIELD (gradient of M)
# ============================================================

print("Computing transport field...")

dM_dlat, dM_dlon = np.gradient(M_grid)

Vx = dM_dlon
Vy = dM_dlat

mag = np.sqrt(Vx**2 + Vy**2) + EPS

Vx /= mag
Vy /= mag

# ============================================================
# LOAD / RECONSTRUCT SHEAR-NET EIGENVECTORS
# ============================================================

print("Constructing shear-net eigenbasis...")

# --- IMPORTANT ---
# Replace this with your exact shear-net generator if desired.
# This is a CLEAN reconstruction based on your paper geometry.

lat_r = np.deg2rad(LAT)
lon_r = np.deg2rad(LON)

# TPW geometry (from your paper)
tpw_lon = np.deg2rad(31.0)
tpw_angle = np.deg2rad(104.0)

lat_p0, lon_p0 = np.pi / 2, 0.0
lat_p1, lon_p1 = lat_p0 - tpw_angle, tpw_lon

def angular_distance(lat_p, lon_p):
    return np.arccos(
        np.sin(lat_r) * np.sin(lat_p) +
        np.cos(lat_r) * np.cos(lat_p) * np.cos(lon_r - lon_p)
    )

psi0 = angular_distance(lat_p0, lon_p0)
psi1 = angular_distance(lat_p1, lon_p1)

# Differential potential (drives shear geometry)
delta_V = np.sin(psi1)**2 - np.sin(psi0)**2

# Gradient of potential → flow proxy
dV_dlat, dV_dlon = np.gradient(delta_V)

# Net 1 direction
U1 = dV_dlon
V1 = dV_dlat

norm1 = np.sqrt(U1**2 + V1**2) + EPS
U1 /= norm1
V1 /= norm1

# Net 2 = orthogonal field
U2 = -V1
V2 = U1

# ============================================================
# MODE PROJECTIONS
# ============================================================

print("Projecting onto modes...")

# dot products
P1 = np.abs(Vx * U1 + Vy * V1)
P2 = np.abs(Vx * U2 + Vy * V2)

# normalize
total = P1 + P2 + EPS
P1 /= total
P2 /= total

deltaP = P1 - P2
competition = 1.0 - np.abs(deltaP)
ratio = P1 / (P2 + EPS)

# ============================================================
# OUTPUT PLOTS
# ============================================================

print("Plotting...")

def save_map(data, title, fname, cmap="viridis", vmin=None, vmax=None):
    plt.figure(figsize=(12, 6))
    plt.imshow(
        data,
        extent=[-180, 180, -90, 90],
        origin="lower",
        cmap=cmap,
        vmin=vmin,
        vmax=vmax,
        aspect="auto"
    )
    plt.colorbar(label=title)
    plt.title(title)
    plt.xlabel("Longitude")
    plt.ylabel("Latitude")
    plt.tight_layout()
    plt.savefig(fname, dpi=150)
    plt.close()

# Core outputs
save_map(deltaP, "Delta P (P1 - P2)", "stage6_deltaP.png", vmin=-1, vmax=1)
save_map(competition, "Mode Competition", "stage6_competition.png", vmin=0, vmax=1)
save_map(np.log10(ratio), "log10(P1/P2)", "stage6_ratio.png")

# Histogram
plt.figure(figsize=(6, 4))
plt.hist(deltaP.flatten(), bins=100)
plt.title("Delta P Distribution")
plt.xlabel("Delta P")
plt.ylabel("Count")
plt.tight_layout()
plt.savefig("stage6_deltaP_hist.png", dpi=150)
plt.close()

# ============================================================
# RIDGE OVERLAY (optional but useful)
# ============================================================

print("Computing ridge overlay...")

grad_mag = np.sqrt(dM_dlat**2 + dM_dlon**2)
threshold = np.percentile(grad_mag, 95)

ridge = grad_mag > threshold

plt.figure(figsize=(12, 6))
plt.imshow(deltaP, extent=[-180, 180, -90, 90],
           origin="lower", cmap="coolwarm", vmin=-1, vmax=1)

# overlay ridge
ys, xs = np.where(ridge)
plt.scatter(
    grid_lon[xs],
    grid_lat[ys],
    s=1,
    c="black",
    alpha=0.3,
    label="High-variance ridge"
)

plt.legend()
plt.title("DeltaP with Ridge Overlay")
plt.tight_layout()
plt.savefig("stage6_deltaP_ridge_overlay.png", dpi=150)
plt.close()

print("Done.")