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


def classification_accuracy(coords, signs, normal):
    preds = np.sign(coords @ normal)
    return np.mean(preds == signs)


def random_normal(seed):
    np.random.seed(seed)
    v = np.random.normal(size=3)
    return v / np.linalg.norm(v)


def worker(args):
    coords, signs, seed = args
    n = random_normal(seed)
    return classification_accuracy(coords, signs, n)


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
    signs = np.sign(signed_vals)

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

    print("Valid signed points:", len(signs))

    # Best-fit normal = dipole direction
    D = compute_dipole(coords, signed_vals)
    n_best = D / np.linalg.norm(D)

    observed_acc = classification_accuracy(coords, signs, n_best)

    print("Observed classification accuracy:", observed_acc)

    print("\nRunning random plane null...")

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

    null_acc = np.array(null_acc)

    p_value = np.mean(null_acc >= observed_acc)

    print("\nResults:")
    print("Observed accuracy:", observed_acc)
    print("Null mean:", np.mean(null_acc))
    print("Null std:", np.std(null_acc))
    print("p-value:", p_value)
