#!/usr/bin/env python3
"""
Stage 9: Euler Geometry Figure

Outputs:
- Global M map
- Dominant basin centroid (S1)
- Candidate Euler
- Great-circle arc between them
"""

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

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

CSV_FILE = "mach.csv"
GRID_RES = 1.0

TARGET_LAT = 0.0
TARGET_LON = -59.0

TOP_K = 5   # use top maxima to estimate centroid

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

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]

# ============================================================
# FIND LOCAL MAXIMA
# ============================================================

footprint = np.ones((5, 5))
local_max = (M_grid == maximum_filter(M_grid, footprint=footprint))

threshold = np.percentile(M_grid, 90)
local_max &= (M_grid >= threshold)

coords = np.argwhere(local_max)

# collect maxima
maxima = []
for i, j in coords:
    maxima.append((grid_lat[i], grid_lon[j], M_grid[i, j]))

# sort
maxima = sorted(maxima, key=lambda x: x[2], reverse=True)

# take top K
top = maxima[:TOP_K]

# ============================================================
# COMPUTE CENTROID (on sphere)
# ============================================================

def sph_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.array([x, y, z])

vecs = np.array([sph_to_cart(lat, lon) for lat, lon, _ in top])
centroid = vecs.mean(axis=0)
centroid /= np.linalg.norm(centroid)

# back to lat/lon
lat_c = np.rad2deg(np.arcsin(centroid[2]))
lon_c = np.rad2deg(np.arctan2(centroid[1], centroid[0]))

print(f"S1 centroid: ({lat_c:.2f}, {lon_c:.2f})")

# ============================================================
# GREAT-CIRCLE ARC
# ============================================================

def great_circle(lat1, lon1, lat2, lon2, n=100):
    lat1, lon1, lat2, lon2 = map(np.deg2rad, [lat1, lon1, lat2, lon2])
    
    # convert to vectors
    A = sph_to_cart(np.rad2deg(lat1), np.rad2deg(lon1))
    B = sph_to_cart(np.rad2deg(lat2), np.rad2deg(lon2))
    
    omega = np.arccos(np.clip(np.dot(A, B), -1, 1))
    
    points = []
    for t in np.linspace(0, 1, n):
        sin_omega = np.sin(omega)
        if sin_omega == 0:
            p = A
        else:
            p = (np.sin((1-t)*omega)/sin_omega)*A + (np.sin(t*omega)/sin_omega)*B
        p /= np.linalg.norm(p)
        lat = np.rad2deg(np.arcsin(p[2]))
        lon = np.rad2deg(np.arctan2(p[1], p[0]))
        points.append((lat, lon))
    
    return np.array(points)

arc = great_circle(lat_c, lon_c, TARGET_LAT, TARGET_LON)

# ============================================================
# PLOT
# ============================================================

plt.figure(figsize=(14, 7))

plt.imshow(
    M_grid,
    extent=[-180, 180, -90, 90],
    origin="lower",
    cmap="viridis",
    aspect="auto"
)

plt.colorbar(label="Compatibility (M)")

# plot centroid
plt.scatter(lon_c, lat_c, c="red", s=100, label="Dominant Basin (S1)")

# plot target Euler
plt.scatter(TARGET_LON, TARGET_LAT, c="white", s=100, edgecolor="black", label="Candidate Euler")

# plot arc
plt.plot(arc[:,1], arc[:,0], 'w--', linewidth=2, label="Angular Separation")

plt.xlabel("Longitude")
plt.ylabel("Latitude")
plt.title("Euler Compatibility Field with Dominant Basin and Candidate Axis")

plt.legend()
plt.tight_layout()

plt.savefig("stage9_euler_geometry.png", dpi=300)
plt.close()

print("Saved: stage9_euler_geometry.png")
print("Done.")