#!/usr/bin/env python3
"""
Stage 6: Ridge Coupling Metrics

Quantifies:
- Secondary mode participation (P2) inside vs outside ridge
- ΔP contrast
- Angular misalignment relative to dominant eigenmode

Outputs:
- printed diagnostics
- CSV summary
"""

import numpy as np
import pandas as pd
from scipy.interpolate import griddata

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

CSV_FILE = "mach.csv"
GRID_RES = 1.0
RIDGE_PERCENTILE = 95  # top X% gradient = ridge

EPS = 1e-12

# ============================================================
# LOAD DATA
# ============================================================

print("Loading M field...")

df = pd.read_csv(CSV_FILE, sep=";")
df.columns = [c.strip().lower() for c in df.columns]

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

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

mask = np.isnan(M_grid)
M_grid[mask] = griddata(points, M, (LAT, LON), method="nearest")[mask]

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

# ============================================================
# SHEAR-NET EIGENVECTORS (same as previous script)
# ============================================================

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

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)

delta_V = np.sin(psi1)**2 - np.sin(psi0)**2

dV_dlat, dV_dlon = np.gradient(delta_V)

U1 = dV_dlon
V1 = dV_dlat

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

U2 = -V1
V2 = U1

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

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

total = P1 + P2 + EPS
P1 /= total
P2 /= total

deltaP = P1 - P2

# ============================================================
# RIDGE DETECTION
# ============================================================

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

ridge_mask = grad_mag >= threshold
background_mask = ~ridge_mask

# ============================================================
# ANGULAR MISALIGNMENT
# ============================================================

dot = np.abs(Vx * U1 + Vy * V1)
dot = np.clip(dot, -1.0, 1.0)

theta = np.arccos(dot)  # radians
theta_deg = np.degrees(theta)

# ============================================================
# METRICS
# ============================================================

def stats(field, mask):
    vals = field[mask]
    return {
        "mean": np.mean(vals),
        "std": np.std(vals),
        "median": np.median(vals)
    }

print("\n--- Ridge Coupling Metrics ---")

P2_ridge = stats(P2, ridge_mask)
P2_bg    = stats(P2, background_mask)

dP_ridge = stats(np.abs(deltaP), ridge_mask)
dP_bg    = stats(np.abs(deltaP), background_mask)

theta_ridge = stats(theta_deg, ridge_mask)
theta_bg    = stats(theta_deg, background_mask)

# Print results
print("\nSecondary Mode (P2):")
print(f"  Ridge     mean={P2_ridge['mean']:.4f}, std={P2_ridge['std']:.4f}")
print(f"  Background mean={P2_bg['mean']:.4f}, std={P2_bg['std']:.4f}")

print("\n|ΔP| (dominance):")
print(f"  Ridge     mean={dP_ridge['mean']:.4f}")
print(f"  Background mean={dP_bg['mean']:.4f}")

print("\nAngular misalignment θ (deg):")
print(f"  Ridge     mean={theta_ridge['mean']:.2f}")
print(f"  Background mean={theta_bg['mean']:.2f}")

# ============================================================
# CSV OUTPUT (for paper)
# ============================================================

out = pd.DataFrame({
    "metric": [
        "P2_mean",
        "P2_std",
        "|dP|_mean",
        "theta_mean_deg"
    ],
    "ridge": [
        P2_ridge["mean"],
        P2_ridge["std"],
        dP_ridge["mean"],
        theta_ridge["mean"]
    ],
    "background": [
        P2_bg["mean"],
        P2_bg["std"],
        dP_bg["mean"],
        theta_bg["mean"]
    ]
})

out.to_csv("stage6_ridge_coupling_metrics.csv", index=False)

print("\nSaved: stage6_ridge_coupling_metrics.csv")
print("Done.")