#!/usr/bin/env python3
"""
Stage 12: Path-integrated compatibility test

Does the candidate Euler axis maximize ∫ M(θ) dθ ?
"""

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

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

S1_LAT = -44.93
S1_LON = 39.71

AXIS_LAT = 0.0
AXIS_LON = -59.0

CSV_FILE = "mach.csv"

ANGLE_STEP = 3.0
N_RANDOM = 500  # increase to 2000+ if desired

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

# ============================================================
# CORE FUNCTION
# ============================================================

def path_integral(axis_vec, v_S1):
    angles = np.arange(0, 360, ANGLE_STEP)
    values = []

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

        if np.isnan(M_a):
            continue

        values.append(M_a)

    if len(values) == 0:
        return np.nan

    return np.trapezoid(values, dx=np.deg2rad(ANGLE_STEP))


# ============================================================
# PREP
# ============================================================

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

print("Computing candidate axis integral...")
candidate_integral = path_integral(v_axis, v_S1)

# ============================================================
# RANDOM AXES
# ============================================================

def random_axis():
    u = np.random.uniform(-1, 1)
    theta = np.random.uniform(0, 2*np.pi)
    z = u
    r = np.sqrt(1 - z*z)
    x = r * np.cos(theta)
    y = r * np.sin(theta)
    return np.array([x, y, z])

print(f"Sampling {N_RANDOM} random axes...")

random_integrals = []

for i in range(N_RANDOM):
    axis = random_axis()
    val = path_integral(axis, v_S1)

    if not np.isnan(val):
        random_integrals.append(val)

random_integrals = np.array(random_integrals)

# ============================================================
# STATS
# ============================================================

mean_rand = np.mean(random_integrals)
std_rand  = np.std(random_integrals)

z_score = (candidate_integral - mean_rand) / std_rand

percentile = np.sum(random_integrals < candidate_integral) / len(random_integrals) * 100

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

print("\n--- Path Integral Results ---")

print(f"Candidate axis integral: {candidate_integral:.6f}")
print(f"Random mean:             {mean_rand:.6f}")
print(f"Random std:              {std_rand:.6f}")
print(f"Z-score:                 {z_score:.2f}")
print(f"Percentile:              {percentile:.1f}%")

# ============================================================
# SAVE
# ============================================================

out = pd.DataFrame({
    "random_integral": random_integrals
})
out.to_csv("stage12_random_integrals.csv", index=False)

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