#!/usr/bin/env python3

import numpy as np
import multiprocessing as mp

CACHE_DIR = "analysis_cache"
PERMUTATIONS = 10000

PLANE_LON = 170.0  # excursion plane longitude


def plane_normal_from_lon(lon_deg):
    """
    Great circle plane oriented N-S at given longitude.
    Plane normal points 90 degrees east of longitude.
    """
    lon = np.radians(lon_deg + 90.0)

    # equatorial vector
    return np.array([
        np.cos(lon),
        np.sin(lon),
        0.0
    ])


def distance_to_plane(coords, normal):
    dots = np.abs(coords @ normal)
    dots = np.clip(dots, -1.0, 1.0)
    return np.arcsin(dots)


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


def worker(args):
    coords, cluster_mask, seed = args
    n = random_plane_normal(seed)
    d = distance_to_plane(coords, n)
    return np.mean(d[cluster_mask])


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")

    # define clusters (same threshold as before)
    cluster_threshold = np.percentile(values, 85)
    cluster_mask = values >= cluster_threshold

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

    # observed plane
    normal = plane_normal_from_lon(PLANE_LON)
    d_obs = distance_to_plane(coords, normal)

    observed_mean_distance = np.mean(d_obs[cluster_mask])

    print("Observed mean cluster distance (radians):",
          observed_mean_distance)

    # Monte Carlo
    print("Running random plane null ensemble...")

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

    null_dist = np.array(null_dist)

    p_value = np.mean(null_dist <= observed_mean_distance)

    print("\nResults:")
    print("Observed mean distance:", observed_mean_distance)
    print("Null mean:", np.mean(null_dist))
    print("Null std:", np.std(null_dist))
    print("p-value:", p_value)
