#!/usr/bin/env python3
"""
Stage 8: Detect preferred Euler states (local maxima)

Outputs:
- ranked maxima
- angular separations
- distance to target Euler
"""

import numpy as np
import pandas as pd
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_N = 20
NEIGHBOR_SIZE = 5   # controls local maxima detection

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

print("Loading 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]

# ============================================================
# LOCAL MAXIMA DETECTION
# ============================================================

print("Detecting local maxima...")

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

# remove flat areas by threshold
threshold = np.percentile(M_grid, 90)
local_max &= (M_grid >= threshold)

coords = np.argwhere(local_max)

# ============================================================
# GREAT-CIRCLE DISTANCE
# ============================================================

def angular_distance(lat1, lon1, lat2, lon2):
    lat1, lon1, lat2, lon2 = map(np.deg2rad, [lat1, lon1, lat2, lon2])
    return np.degrees(np.arccos(
        np.sin(lat1)*np.sin(lat2) +
        np.cos(lat1)*np.cos(lat2)*np.cos(lon1 - lon2)
    ))

# ============================================================
# COLLECT MAXIMA
# ============================================================

records = []

for i, j in coords:
    lat0 = grid_lat[i]
    lon0 = grid_lon[j]
    mval = M_grid[i, j]

    dist_to_target = angular_distance(lat0, lon0, TARGET_LAT, TARGET_LON)

    records.append({
        "lat": lat0,
        "lon": lon0,
        "M": mval,
        "dist_to_target_deg": dist_to_target
    })

# sort by strength
records = sorted(records, key=lambda x: x["M"], reverse=True)

# keep top N
records = records[:TOP_N]

# ============================================================
# PAIRWISE ANGULAR SEPARATION
# ============================================================

print("Computing pairwise separations...")

for k in range(len(records)):
    for l in range(k+1, len(records)):
        d = angular_distance(
            records[k]["lat"], records[k]["lon"],
            records[l]["lat"], records[l]["lon"]
        )
        records[k][f"sep_to_{l}"] = d

# ============================================================
# OUTPUT
# ============================================================

df_out = pd.DataFrame(records)
df_out.to_csv("stage8_euler_maxima.csv", index=False)

print("\nTop Euler maxima:")
print(df_out[["lat", "lon", "M", "dist_to_target_deg"]].head(10))

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