#!/usr/bin/env python3
"""
Stage 13: Low-compatibility avoidance test

Metrics:
1. Fraction of path below threshold
2. Integrated penalty below threshold
3. Worst-case minimum M along path

Goal:
→ Does the candidate axis minimize exposure to low-M regions?
"""

import numpy as np
import pandas as pd
import multiprocessing as mp
import time
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 = 5000

# Threshold definition (adaptive)
LOW_PERCENTILE = 20  # bottom 20% = "bad region"

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

# Threshold
threshold = np.percentile(M, LOW_PERCENTILE)
print(f"Low-M threshold ({LOW_PERCENTILE}th percentile): {threshold:.4f}")

# ============================================================
# GEOMETRY
# ============================================================

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):
    axis = axis / np.linalg.norm(axis)
    angles = np.deg2rad(angles)

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

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

# ============================================================
# METRICS
# ============================================================

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

    v_rot = rotate_batch(v_S1, axis_vec, ANGLES)
    lat_q, lon_q = cart_to_sph(v_rot)

    M_vals = interp(lat_q, lon_q)

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

    M_vals = M_vals[mask]

    # --- Metrics ---
    frac_low = np.mean(M_vals < threshold)

    penalty = np.sum(
        (threshold - M_vals[M_vals < threshold])
    )

    min_M = np.min(M_vals)

    return frac_low, penalty, min_M

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

def worker(batch):
    out = []
    for a in batch:
        res = evaluate_axis(a)
        if res:
            out.append(res)
    return out

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

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

    # Candidate
    print("\nEvaluating candidate axis...")
    v_axis = sph_to_cart(AXIS_LAT, AXIS_LON)
    cand = evaluate_axis(v_axis)

    print(f"\nCandidate:")
    print(f"  frac_low: {cand[0]:.4f}")
    print(f"  penalty : {cand[1]:.4f}")
    print(f"  min_M   : {cand[2]:.4f}")

    # Random
    print(f"\nSampling {N_RANDOM} random axes...")

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

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

    results = []
    done = 0

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

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

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

    pool.close()
    pool.join()

    results = np.array(results)

    frac_rand = results[:,0]
    pen_rand  = results[:,1]
    min_rand  = results[:,2]

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

    def summarize(name, candidate, random_vals, lower_is_better=True):
        mean = np.mean(random_vals)
        std  = np.std(random_vals)

        if lower_is_better:
            z = (mean - candidate) / std
            pct = np.sum(random_vals > candidate)/len(random_vals)*100
        else:
            z = (candidate - mean) / std
            pct = np.sum(random_vals < candidate)/len(random_vals)*100

        print(f"\n{name}:")
        print(f"  candidate: {candidate:.4f}")
        print(f"  mean:      {mean:.4f}")
        print(f"  std:       {std:.4f}")
        print(f"  Z-score:   {z:.2f}")
        print(f"  percentile:{pct:.1f}%")

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

    summarize("Low-fraction (↓ better)", cand[0], frac_rand, True)
    summarize("Penalty (↓ better)",      cand[1], pen_rand,  True)
    summarize("Min M (↑ better)",        cand[2], min_rand,  False)

    pd.DataFrame(results, columns=["frac_low","penalty","min_M"])\
        .to_csv("stage13_lowM_metrics.csv", index=False)

    print("\nSaved: stage13_lowM_metrics.csv")

if __name__ == "__main__":
    main()