#!/usr/bin/env python3

import numpy as np
import multiprocessing as mp

CACHE_DIR = "analysis_cache"
PERMUTATIONS = 10000


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 fit_dipole(coords, values):
    # Solve least squares: coords @ a = values
    a, _, _, _ = np.linalg.lstsq(coords, values, rcond=None)
    predicted = coords @ a

    ss_res = np.sum((values - predicted)**2)
    ss_tot = np.sum((values - np.mean(values))**2)

    r2 = 1 - ss_res / ss_tot

    return a, r2


def worker(args):
    coords, values, seed = args
    np.random.seed(seed)
    perm = np.random.permutation(values)
    _, r2 = fit_dipole(coords, perm)
    return r2


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

    # Remove zeros
    valid = signed_vals != 0
    coords = coords[valid]
    signed_vals = signed_vals[valid]

    print("Valid points:", len(signed_vals))

    # Fit observed dipole
    a_obs, r2_obs = fit_dipole(coords, signed_vals)

    lat_obs, lon_obs = vec_to_latlon(a_obs)

    print("\nObserved dipole regression:")
    print("Axis lat:", lat_obs, "lon:", lon_obs)
    print("R^2:", r2_obs)

    print("\nRunning permutation null...")

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

    null_r2 = np.array(null_r2)

    p_value = np.mean(null_r2 >= r2_obs)

    print("\nResults:")
    print("Observed R^2:", r2_obs)
    print("Null mean:", np.mean(null_r2))
    print("Null std:", np.std(null_r2))
    print("p-value:", p_value)
