#!/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_tensor(coords, weights):
    M = np.zeros((3,3))
    for x, w in zip(coords, weights):
        M += w * 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, values, threshold, seed = args
    np.random.seed(seed)
    perm = np.random.permutation(values)
    mask = perm >= threshold
    M = compute_tensor(coords[mask], perm[mask])
    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")
    values = np.load(f"{CACHE_DIR}/values.npy")

    # Cluster definition
    threshold = np.percentile(values, 85)
    mask = values >= threshold

    print("Cluster points:", np.sum(mask))

    weights = values[mask]

    M = compute_tensor(coords[mask], weights)
    eigvals, eigvecs = eigen_analysis(M)

    print("\nEigenvalues:", eigvals)

    # Largest axis
    lat_max, lon_max = vec_to_latlon(eigvecs[:, -1])
    # Smallest axis
    lat_min, lon_min = vec_to_latlon(eigvecs[:, 0])

    print("\nDominant axis (max eigenvector):")
    print("Lat:", lat_max, "Lon:", lon_max)

    print("\nPlane normal candidate (min eigenvector):")
    print("Lat:", lat_min, "Lon:", lon_min)

    observed_spread = eigvals[-1] - eigvals[0]
    print("\nEigenvalue spread:", observed_spread)

    # Monte Carlo
    print("\nRunning permutation null ensemble...")

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

    null_spread = np.array(null_spread)

    p_value = np.mean(null_spread >= observed_spread)

    print("\nResults:")
    print("Observed spread:", observed_spread)
    print("Null mean:", np.mean(null_spread))
    print("Null std:", np.std(null_spread))
    print("p-value:", p_value)
