#!/usr/bin/env python3

import numpy as np
import multiprocessing as mp

CACHE_DIR = "analysis_cache"
PERMUTATIONS = 10000


def latlon_to_vec(lat_deg, lon_deg):
    lat = np.radians(lat_deg)
    lon = np.radians(lon_deg)
    return np.array([
        np.cos(lat) * np.cos(lon),
        np.cos(lat) * np.sin(lon),
        np.sin(lat)
    ])


def vec_to_latlon(v):
    v = v / np.linalg.norm(v)
    lat = np.degrees(np.arcsin(v[2]))
    lon = np.degrees(np.arctan2(v[1], v[0]))
    return lat, lon


def angular_sep(v1, v2):
    dot = np.clip(np.dot(v1, v2), -1.0, 1.0)
    return np.degrees(np.arccos(dot))


def build_design_matrix(coords):
    x = coords[:,0]
    y = coords[:,1]
    z = coords[:,2]

    return np.column_stack([
        x, y, z,
        x*x, y*y, z*z,
        x*y, x*z, y*z
    ])


def dominant_quadrupole_axis(coeffs):
    b_xx, b_yy, b_zz, b_xy, b_xz, b_yz = coeffs[3:]

    Q = np.array([
        [b_xx, b_xy/2, b_xz/2],
        [b_xy/2, b_yy, b_yz/2],
        [b_xz/2, b_yz/2, b_zz]
    ])

    eigvals, eigvecs = np.linalg.eigh(Q)
    idx = np.argmax(eigvals)
    return eigvecs[:, idx]


def worker(args):
    X, coords, values, euler_vec, seed = args
    np.random.seed(seed)
    perm = np.random.permutation(values)
    coeffs, _, _, _ = np.linalg.lstsq(X, perm, rcond=None)
    axis = dominant_quadrupole_axis(coeffs)
    return angular_sep(axis, euler_vec)


if __name__ == "__main__":

    mp.set_start_method("spawn", force=True)

    coords = np.load(f"{CACHE_DIR}/coords.npy")

    # Reload signed dT
    import pandas as pd
    DATAFILE = "frost2021.dat"

    with open(DATAFILE) as f:
        lines = f.readlines()

    header_idx = None
    for i, line in enumerate(lines):
        if line.strip().startswith("EvtDate"):
            header_idx = i
            break

    df = pd.read_csv(
        DATAFILE,
        sep=r"\s+",
        engine="python",
        skiprows=header_idx
    )

    df["dT"] = df["dTabdf"].copy()
    mask = df["dT"] == 0
    df.loc[mask, "dT"] = df.loc[mask, "dTbcdf"]

    signed_vals = df["dT"].values

    valid = signed_vals != 0
    coords = coords[valid]
    signed_vals = signed_vals[valid]

    X = build_design_matrix(coords)

    # Observed fit
    coeffs_obs, _, _, _ = np.linalg.lstsq(X, signed_vals, rcond=None)
    axis_obs = dominant_quadrupole_axis(coeffs_obs)

    euler_vec = latlon_to_vec(0.0, 121.0)

    observed_angle = angular_sep(axis_obs, euler_vec)

    print("Observed quadrupole axis angle to 121°E:", observed_angle)

    print("\nRunning permutation Monte Carlo...")

    with mp.Pool(mp.cpu_count()) as pool:
        null_angles = pool.map(
            worker,
            [(X, coords, signed_vals, euler_vec, i)
             for i in range(PERMUTATIONS)]
        )

    null_angles = np.array(null_angles)

    p_value = np.mean(null_angles <= observed_angle)

    print("\nResults:")
    print("Observed angle:", observed_angle)
    print("Null mean angle:", np.mean(null_angles))
    print("Null std:", np.std(null_angles))
    print("p-value:", p_value)
