import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter

# -------------------------------
# CONFIG
# -------------------------------

M_FILE = "refined_streamline_sweep.npy"
CSV_FILE = "mach.csv"

OUTPUT_PREFIX = "stage6_mode"
SMOOTH_SIGMA = 1.0
RIDGE_PERCENTILE = 85

# -------------------------------
# LOAD DATA
# -------------------------------

print("Loading data...")

M = np.load(M_FILE)

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

# --- AUTO DETECT COLUMNS ---
cols = {c.lower(): c for c in df.columns}

lat_col = cols.get('lat') or cols.get('latitude')
lon_col = cols.get('lon') or cols.get('longitude')

if lat_col is None or lon_col is None:
    raise ValueError(f"Could not find lat/lon columns. Found: {df.columns}")

lats = df[lat_col].values
lons = df[lon_col].values

# -------------------------------
# GRID RECONSTRUCTION
# -------------------------------

lat_unique = np.unique(lats)
lon_unique = np.unique(lons)

nlat = len(lat_unique)
nlon = len(lon_unique)

print(f"Grid: {nlat} x {nlon}")

# Sort to ensure correct orientation
lat_unique.sort()
lon_unique.sort()

# -------------------------------
# TRANSPORT VECTOR (GRADIENT)
# -------------------------------

print("Computing transport field...")

dM_dlat, dM_dlon = np.gradient(M)

Vx = -dM_dlon
Vy = -dM_dlat

V = np.stack([Vx, Vy], axis=-1)

norm = np.linalg.norm(V, axis=-1, keepdims=True)
norm[norm == 0] = 1
Vn = V / norm

# -------------------------------
# STRUCTURE TENSOR (VECTORISED)
# -------------------------------

print("Computing anisotropy tensor...")

Jxx = gaussian_filter(dM_dlon * dM_dlon, SMOOTH_SIGMA)
Jyy = gaussian_filter(dM_dlat * dM_dlat, SMOOTH_SIGMA)
Jxy = gaussian_filter(dM_dlon * dM_dlat, SMOOTH_SIGMA)

# Stack tensors
J = np.stack([
    np.stack([Jxx, Jxy], axis=-1),
    np.stack([Jxy, Jyy], axis=-1)
], axis=-2)  # shape (lat, lon, 2, 2)

# Eigen decomposition (vectorised)
w, v = np.linalg.eigh(J)

# Principal directions
E1 = v[..., :, 1]  # dominant
E2 = v[..., :, 0]  # orthogonal

# Normalize
E1 /= np.linalg.norm(E1, axis=-1, keepdims=True)
E2 /= np.linalg.norm(E2, axis=-1, keepdims=True)

# -------------------------------
# MODE PROJECTIONS
# -------------------------------

print("Computing mode projections...")

P1 = np.abs(np.sum(Vn * E1, axis=-1))
P2 = np.abs(np.sum(Vn * E2, axis=-1))

delta_P = P1 - P2
competition = 1.0 - np.abs(delta_P)

# -------------------------------
# RIDGE
# -------------------------------

threshold = np.percentile(M, RIDGE_PERCENTILE)
ridge_mask = M > threshold

# -------------------------------
# SAVE ARRAYS
# -------------------------------

np.save(f"{OUTPUT_PREFIX}_P1.npy", P1)
np.save(f"{OUTPUT_PREFIX}_P2.npy", P2)
np.save(f"{OUTPUT_PREFIX}_deltaP.npy", delta_P)
np.save(f"{OUTPUT_PREFIX}_competition.npy", competition)

# -------------------------------
# PLOTS
# -------------------------------

def plot_field(field, title, fname, cmap='coolwarm'):
    plt.figure(figsize=(12,6))
    plt.imshow(field, origin='lower', cmap=cmap)
    plt.colorbar(label=title)
    plt.title(title)
    plt.tight_layout()
    plt.savefig(fname, dpi=200)
    plt.close()

plot_field(delta_P, "Mode Dominance (ΔP)", f"{OUTPUT_PREFIX}_deltaP.png")
plot_field(competition, "Mode Competition", f"{OUTPUT_PREFIX}_competition.png", cmap='viridis')

# -------------------------------
# KEY OVERLAY
# -------------------------------

print("Generating overlay...")

plt.figure(figsize=(12,6))

plt.imshow(M, origin='lower', cmap='plasma')
plt.colorbar(label='M')

plt.contour(ridge_mask, levels=[0.5], colors='white', linewidths=1)
plt.contour(delta_P, levels=[0], colors='cyan', linewidths=1)

plt.title("Ridge (white) vs Mode Boundary ΔP=0 (cyan)")
plt.tight_layout()
plt.savefig(f"{OUTPUT_PREFIX}_ridge_vs_modes.png", dpi=250)
plt.close()

# -------------------------------
# DIAGNOSTICS
# -------------------------------

print("\n--- Diagnostics ---")

print(f"Mean P1: {np.mean(P1):.4f}")
print(f"Mean P2: {np.mean(P2):.4f}")
print(f"Mode1 fraction: {np.mean(delta_P > 0):.3f}")
print(f"Mode2 fraction: {np.mean(delta_P < 0):.3f}")
print(f"Ridge |ΔP|: {np.mean(np.abs(delta_P[ridge_mask])):.4f}")

# -------------------------------
# HISTOGRAM
# -------------------------------

plt.figure(figsize=(6,4))
plt.hist(delta_P.flatten(), bins=100)
plt.title("ΔP Distribution")
plt.tight_layout()
plt.savefig(f"{OUTPUT_PREFIX}_deltaP_hist.png", dpi=200)
plt.close()

print("Done.")