#!/usr/bin/env python3

import numpy as np
import multiprocessing as mp

CACHE_DIR = "analysis_cache"
PERMUTATIONS = 10000
CLUSTER_COUNT = 772  # from previous result


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_tensor(coords):
    M = np.zeros((3,3))
    for x in coords:
        M += np.outer(x, x)
    return M


def eigen_analysis(M):
    eigvals, eigvecs = np.linalg.eigh(M)
    idx = np.argsort(eigvals)
    eigvals = eigvals[idx]
    eigvecs = eigvecs[:, idx]
    return eigvals, eigvecs


def worker(args):
    coords, seed = args
    np.random.seed(seed)
    idx = np.random.choice(len(coords), CLUSTER_COUNT, replace=False)
    M = compute_tensor(coords[idx])
    eigvals, _ = eigen_analysis(M)
    return eigvals[-1] - eigvals[0]


if __name__ == "__main__":

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

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

    print("Total turning points:", len(coords))

    # Full dataset PCA
    M_full = compute_tensor(coords)
    eigvals_full, eigvecs_full = eigen_analysis(M_full)

    print("\nFull sampling eigenvalues:", eigvals_full)

    lat_full, lon_full = vec_to_latlon(eigvecs_full[:, -1])

    print("\nFull sampling dominant axis:")
    print("Lat:", lat_full, "Lon:", lon_full)

    full_spread = eigvals_full[-1] - eigvals_full[0]
    print("Full sampling spread:", full_spread)

    # Monte Carlo random subsets
    print("\nRunning random subset null ensemble...")

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

    null_spread = np.array(null_spread)

    print("\nRandom subset statistics:")
    print("Mean spread:", np.mean(null_spread))
    print("Std spread:", np.std(null_spread))
