#!/usr/bin/env python3

import numpy as np
import multiprocessing as mp

CACHE_DIR = "analysis_cache"
PERMUTATIONS = 10000

CLUSTER_PERCENTILE = 85   # top 15%
SHEAR_PERCENTILE = 80     # top 20%


# --------------------------
# Random Euler rotation
# --------------------------

def random_rotation_matrix(seed):
    np.random.seed(seed)

    u = np.random.normal(size=3)
    u /= np.linalg.norm(u)

    theta = np.random.uniform(0, 2*np.pi)

    a = np.cos(theta)
    b = np.sin(theta)

    K = np.array([
        [0, -u[2], u[1]],
        [u[2], 0, -u[0]],
        [-u[1], u[0], 0]
    ])

    I = np.eye(3)

    return a*I + b*K + (1-a)*np.outer(u, u)


def rotate_shear(coords, shear_vals, seed):
    R = random_rotation_matrix(seed)
    rotated_coords = (R @ coords.T).T

    dots = rotated_coords @ coords.T
    idx = np.argmax(dots, axis=1)

    return shear_vals[idx]


def overlap_fraction(cluster_mask, shear_mask):
    intersection = np.sum(cluster_mask & shear_mask)
    return intersection / np.sum(cluster_mask)


def worker(args):
    coords, shear_vals, cluster_mask, seed = args
    rotated_shear = rotate_shear(coords, shear_vals, seed)

    threshold = np.percentile(rotated_shear, SHEAR_PERCENTILE)
    shear_mask = rotated_shear >= threshold

    return overlap_fraction(cluster_mask, shear_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")
    shear_vals = np.load(f"{CACHE_DIR}/shear_vals.npy")

    # --------------------------
    # Define clusters
    # --------------------------

    cluster_threshold = np.percentile(values, CLUSTER_PERCENTILE)
    cluster_mask = values >= cluster_threshold

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

    # --------------------------
    # Observed overlap
    # --------------------------

    shear_threshold = np.percentile(shear_vals, SHEAR_PERCENTILE)
    shear_mask = shear_vals >= shear_threshold

    observed_overlap = overlap_fraction(cluster_mask, shear_mask)

    print("Observed overlap fraction:", observed_overlap)

    # --------------------------
    # Monte Carlo
    # --------------------------

    print("Running rotation null ensemble...")

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

    null_overlap = np.array(null_overlap)

    p_value = np.mean(null_overlap >= observed_overlap)

    print("\nResults:")
    print("Observed overlap:", observed_overlap)
    print("Null mean:", np.mean(null_overlap))
    print("Null std:", np.std(null_overlap))
    print("p-value:", p_value)
