#!/usr/bin/env python3

import numpy as np
import multiprocessing as mp

CACHE_DIR = "analysis_cache"
PERMUTATIONS = 10000


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 fit_model(X, values):
    coeffs, _, _, _ = np.linalg.lstsq(X, values, rcond=None)
    predicted = X @ coeffs

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

    r2 = 1 - ss_res / ss_tot

    return coeffs, r2


def worker(args):
    X, values, seed = args
    np.random.seed(seed)
    perm = np.random.permutation(values)
    _, r2 = fit_model(X, 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))

    X = build_design_matrix(coords)

    # Fit observed model
    coeffs_obs, r2_obs = fit_model(X, signed_vals)

    print("\nObserved degree-2 model:")
    print("R^2:", r2_obs)

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

    with mp.Pool(mp.cpu_count()) as pool:
        null_r2 = pool.map(
            worker,
            [(X, 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)
