#!/usr/bin/env python3
"""
Stage 11: Full rotational sweep of S1 about candidate Euler axis

Produces:
- M vs rotation angle
- trajectory map
- CSV of results
"""

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

# ============================================================
# INPUTS
# ============================================================

S1_LAT = -44.93
S1_LON = 39.71

AXIS_LAT = 0.0
AXIS_LON = -59.0

CSV_FILE = "mach.csv"

ANGLE_STEP = 2.0  # degrees

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

points = np.column_stack((lat, lon))

# ============================================================
# SPHERE UTILITIES
# ============================================================

def sph_to_cart(lat, lon):
    lat = np.deg2rad(lat)
    lon = np.deg2rad(lon)
    return np.array([
        np.cos(lat) * np.cos(lon),
        np.cos(lat) * np.sin(lon),
        np.sin(lat)
    ])

def cart_to_sph(v):
    x, y, z = v
    lat = np.rad2deg(np.arcsin(z))
    lon = np.rad2deg(np.arctan2(y, x))
    return lat, lon

def rotate(v, axis, angle_deg):
    angle = np.deg2rad(angle_deg)
    axis = axis / np.linalg.norm(axis)
    return (
        v * np.cos(angle)
        + np.cross(axis, v) * np.sin(angle)
        + axis * np.dot(axis, v) * (1 - np.cos(angle))
    )

def interp_M(lat_q, lon_q):
    return griddata(points, M, [(lat_q, lon_q)], method="linear")[0]

# ============================================================
# INITIAL VECTORS
# ============================================================

v_S1 = sph_to_cart(S1_LAT, S1_LON)
v_axis = sph_to_cart(AXIS_LAT, AXIS_LON)

# ============================================================
# SWEEP
# ============================================================

angles = np.arange(0, 360 + ANGLE_STEP, ANGLE_STEP)

traj_lat = []
traj_lon = []
traj_M   = []

print("Running rotational sweep...")

for a in angles:
    v = rotate(v_S1, v_axis, a)
    lat_a, lon_a = cart_to_sph(v)
    M_a = interp_M(lat_a, lon_a)

    traj_lat.append(lat_a)
    traj_lon.append(lon_a)
    traj_M.append(M_a)

traj_lat = np.array(traj_lat)
traj_lon = np.array(traj_lon)
traj_M   = np.array(traj_M)

# ============================================================
# SAVE CSV
# ============================================================

out_df = pd.DataFrame({
    "angle_deg": angles,
    "lat": traj_lat,
    "lon": traj_lon,
    "M": traj_M
})

out_df.to_csv("stage11_rotational_sweep.csv", index=False)

# ============================================================
# PLOT 1: M vs angle
# ============================================================

plt.figure(figsize=(10, 4))
plt.plot(angles, traj_M)
plt.xlabel("Rotation angle (deg)")
plt.ylabel("Compatibility (M)")
plt.title("Rotational Sweep: M vs Angle")
plt.grid(True)
plt.savefig("stage11_M_vs_angle.png", dpi=200)
plt.close()

# ============================================================
# PLOT 2: trajectory map
# ============================================================

plt.figure(figsize=(12, 6))
plt.scatter(traj_lon, traj_lat, c=traj_M, s=10)
plt.xlabel("Longitude")
plt.ylabel("Latitude")
plt.title("Rotational Trajectory of S1 about Euler Axis")
plt.colorbar(label="M")
plt.savefig("stage11_trajectory_map.png", dpi=200)
plt.close()

# ============================================================
# DIAGNOSTICS
# ============================================================

print("\n--- Sweep Diagnostics ---")
print(f"Max M: {np.nanmax(traj_M):.4f}")
print(f"Min M: {np.nanmin(traj_M):.4f}")
print(f"Mean M: {np.nanmean(traj_M):.4f}")

print("\nSaved:")
print("  stage11_rotational_sweep.csv")
print("  stage11_M_vs_angle.png")
print("  stage11_trajectory_map.png")

print("\nDone.")