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

Optimizations:
- Vectorized rotations
- Fast interpolator (LinearNDInterpolator)
- Multiprocessing (Apple Silicon safe)
- Progress reporting
"""

import numpy as np
import pandas as pd
import time
import multiprocessing as mp
from scipy.interpolate import LinearNDInterpolator

# ============================================================
# 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 = 10000   # scale this up safely now
N_WORKERS = mp.cpu_count()

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

print("Building interpolator...")
interp = LinearNDInterpolator(points, M)

# ============================================================
# SPHERE UTILITIES (VECTOR SAFE)
# ============================================================

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[:,0], v[:,1], v[:,2]
    lat = np.rad2deg(np.arcsin(z))
    lon = np.rad2deg(np.arctan2(y, x))
    return lat, lon

def rotate_batch(v, axis, angles_deg):
    axis = axis / np.linalg.norm(axis)
    angles = np.deg2rad(angles_deg)

    v = v.reshape(1,3)
    axis = axis.reshape(1,3)

    cos = np.cos(angles)[:,None]
    sin = np.sin(angles)[:,None]

    cross = np.cross(axis, v)
    dot = np.sum(axis * v, axis=1)

    return (
        v * cos
        + cross * sin
        + axis * dot * (1 - cos)
    )

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

ANGLES = np.arange(0, 360, ANGLE_STEP)

def path_integral(axis_vec):
    v_S1 = sph_to_cart(S1_LAT, S1_LON)

    v_rot = rotate_batch(v_S1, axis_vec, ANGLES)
    lat, lon = cart_to_sph(v_rot)

    M_vals = interp(lat, lon)

    mask = ~np.isnan(M_vals)
    if np.sum(mask) < 10:
        return np.nan

    return np.trapezoid(M_vals[mask], dx=np.deg2rad(ANGLE_STEP))

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

def random_axes(n):
    u = np.random.uniform(-1, 1, n)
    theta = np.random.uniform(0, 2*np.pi, n)
    z = u
    r = np.sqrt(1 - z*z)
    x = r * np.cos(theta)
    y = r * np.sin(theta)
    return np.column_stack((x,y,z))

# ============================================================
# PARALLEL WORKER
# ============================================================

def worker_batch(axes_batch):
    return [path_integral(a) for a in axes_batch]

# ============================================================
# MAIN
# ============================================================

def main():
    start = time.time()

    print("Computing candidate axis...")
    v_axis = sph_to_cart(AXIS_LAT, AXIS_LON)
    candidate = path_integral(v_axis)

    print(f"Candidate integral: {candidate:.6f}")

    print(f"\nSampling {N_RANDOM} random axes with {N_WORKERS} workers...")

    axes = random_axes(N_RANDOM)
    batch_size = 50
    batches = [axes[i:i+batch_size] for i in range(0, len(axes), batch_size)]

    ctx = mp.get_context("spawn")
    pool = ctx.Pool(N_WORKERS)

    results = []
    completed = 0

    for out in pool.imap(worker_batch, batches):
        results.extend(out)
        completed += len(out)

        elapsed = time.time() - start
        rate = completed / elapsed
        eta = (N_RANDOM - completed) / rate if rate > 0 else 0

        print(f"[{completed}/{N_RANDOM}] "
              f"{rate:.1f} axes/s | ETA: {eta:.1f}s", end="\r")

    pool.close()
    pool.join()

    results = np.array([r for r in results if not np.isnan(r)])

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

    mean = np.mean(results)
    std  = np.std(results)

    z = (candidate - mean) / std
    pct = np.sum(results < candidate) / len(results) * 100

    print("\n\n--- Results ---")
    print(f"Candidate: {candidate:.6f}")
    print(f"Mean:      {mean:.6f}")
    print(f"Std:       {std:.6f}")
    print(f"Z-score:   {z:.2f}")
    print(f"Percentile:{pct:.1f}%")

    pd.DataFrame({"integral": results}).to_csv(
        "stage12_random_integrals.csv", index=False
    )

    print("\nSaved: stage12_random_integrals.csv")

if __name__ == "__main__":
    main()