#!/usr/bin/env python3

import numpy as np

CACHE_DIR = "analysis_cache"
PERCENTILES = [80, 85, 90, 95]


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=None):
    M = np.zeros((3,3))
    if weights is None:
        for x in coords:
            M += np.outer(x, x)
    else:
        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


if __name__ == "__main__":

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

    print("Threshold Robustness Sweep")
    print("---------------------------------------")

    for pct in PERCENTILES:

        threshold = np.percentile(values, pct)
        mask = values >= threshold

        coords_sub = coords[mask]
        weights_sub = values[mask]

        print(f"\nPercentile: {pct}")
        print("Cluster size:", len(coords_sub))

        # Weighted PCA
        M_weighted = compute_tensor(coords_sub, weights_sub)
        eigvals_w, eigvecs_w = eigen_analysis(M_weighted)

        lat_w, lon_w = vec_to_latlon(eigvecs_w[:, -1])
        spread_w = eigvals_w[-1] - eigvals_w[0]

        print("Weighted axis:")
        print(" Lat:", lat_w, "Lon:", lon_w)
        print(" Spread:", spread_w)

        # Binary PCA
        M_binary = compute_tensor(coords_sub)
        eigvals_b, eigvecs_b = eigen_analysis(M_binary)

        lat_b, lon_b = vec_to_latlon(eigvecs_b[:, -1])
        spread_b = eigvals_b[-1] - eigvals_b[0]

        print("Binary axis:")
        print(" Lat:", lat_b, "Lon:", lon_b)
        print(" Spread:", spread_b)
