#!/usr/bin/env python3

import numpy as np
import multiprocessing as mp
from scipy.stats import spearmanr

CACHE_DIR = "analysis_cache"
PERMUTATIONS = 10000

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

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

    # Random unit vector
    u = np.random.normal(size=3)
    u /= np.linalg.norm(u)

    # Random angle
    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_field(coords, shear_vals, seed):
    R = random_rotation_matrix(seed)
    rotated_coords = (R @ coords.T).T

    # Re-evaluate shear field at rotated positions
    # Instead of recomputing full shear, use nearest-neighbour mapping
    # since shear_vals already computed on sphere

    # Compute dot products to original coords
    dots = rotated_coords @ coords.T
    idx = np.argmax(dots, axis=1)

    return shear_vals[idx]


def worker(args):
    coords, shear_vals, values, seed = args
    rotated_shear = rotate_field(coords, shear_vals, seed)
    rho, _ = spearmanr(values, rotated_shear)
    return rho


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

    print("Computing observed correlation...")

    observed_rho, _ = spearmanr(values, shear_vals)
    print("Observed Spearman rho:", observed_rho)

    print("Running Euler-rotation Monte Carlo...")

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

    null_rho = np.array(null_rho)

    p_value = np.mean(null_rho >= observed_rho)

    print("\nResults:")
    print("Observed rho:", observed_rho)
    print("Null mean:", np.mean(null_rho))
    print("Null std:", np.std(null_rho))
    print("p-value:", p_value)
