#!/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 compute_dipole(coords, values):
    return np.sum(coords * values[:, None], axis=0)


def worker(args):
    coords, values, seed = args
    np.random.seed(seed)
    perm = np.random.permutation(values)
    D = compute_dipole(coords, perm)
    return np.linalg.norm(D)


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

    print("Computing observed dipole...")

    D_obs = compute_dipole(coords, signed_vals)
    D_mag = np.linalg.norm(D_obs)

    lat_obs, lon_obs = vec_to_latlon(D_obs)

    print("Observed dipole magnitude:", D_mag)
    print("Dipole axis lat:", lat_obs, "lon:", lon_obs)

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

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

    null_mag = np.array(null_mag)

    p_value = np.mean(null_mag >= D_mag)

    print("\nResults:")
    print("Observed magnitude:", D_mag)
    print("Null mean:", np.mean(null_mag))
    print("Null std:", np.std(null_mag))
    print("p-value:", p_value)
