#!/usr/bin/env python3

import numpy as np
import multiprocessing as mp

EARTH_RADIUS_KM = 6371.0
PERMUTATIONS = 10000
SCALES_KM = [250, 500, 1000, 2000, 3000, 4000]

CACHE_DIR = "analysis_cache"


def morans_I(values, weight_matrix):
    """
    Compute Moran's I given weight matrix
    """

    N = len(values)
    x = values
    x_bar = np.mean(x)

    W = weight_matrix.sum()

    num = np.sum(weight_matrix * np.outer(x - x_bar, x - x_bar))
    den = np.sum((x - x_bar) ** 2)

    return (N / W) * (num / den)


def permutation_worker(args):
    values, weight_matrix, seed = args

    np.random.seed(seed)
    perm = np.random.permutation(values)
    return morans_I(perm, weight_matrix)


if __name__ == "__main__":

    mp.set_start_method("spawn", force=True)

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

    print("Dataset size:", len(values))

    results = {}

    for scale_km in SCALES_KM:

        print(f"\nScale: {scale_km} km")

        theta = scale_km / EARTH_RADIUS_KM

        # Binary weight matrix
        weight_matrix = (dist_matrix < theta).astype(np.float64)

        np.fill_diagonal(weight_matrix, 0.0)

        observed_I = morans_I(values, weight_matrix)

        print("Observed I:", observed_I)

        # Monte Carlo permutations
        print("Running permutations...")

        with mp.Pool(mp.cpu_count()) as pool:
            perm_I = pool.map(
                permutation_worker,
                [(values, weight_matrix, i) for i in range(PERMUTATIONS)]
            )

        perm_I = np.array(perm_I)

        p_value = np.mean(perm_I >= observed_I)

        print("p-value:", p_value)

        results[scale_km] = {
            "observed_I": observed_I,
            "null_mean": np.mean(perm_I),
            "null_std": np.std(perm_I),
            "p_value": p_value
        }

    print("\nFinal Results:")
    for k in results:
        print(k, results[k])
