#!/usr/bin/env python3

import numpy as np
import pandas as pd
from math import radians, cos, sin
import os

# Apple Silicon multiprocessing safety (future use)
if __name__ == "__main__":

    np.random.seed(42)

    DATAFILE = "frost2021.dat"
    OUTDIR = "analysis_cache"

    os.makedirs(OUTDIR, exist_ok=True)

    # -----------------------------
    # Load Frost dataset
    # -----------------------------
    with open(DATAFILE) as f:
        lines = f.readlines()

    header_idx = None
    for i, line in enumerate(lines):
        if line.strip().startswith("EvtDate"):
            header_idx = i
            break

    df = pd.read_csv(
        DATAFILE,
        sep=r"\s+",
        engine="python",
        skiprows=header_idx
    )

    # Resolve dT exactly as in Frost scripts
    df["dT"] = df["dTabdf"].copy()
    mask = df["dT"] == 0
    df.loc[mask, "dT"] = df.loc[mask, "dTbcdf"]
    df["abs_dT"] = df["dT"].abs()

    # -----------------------------
    # Convert to spherical vectors
    # -----------------------------
    def latlon_to_vec(lat, lon):
        lat = np.radians(lat)
        lon = np.radians(lon)
        x = np.cos(lat) * np.cos(lon)
        y = np.cos(lat) * np.sin(lon)
        z = np.sin(lat)
        return np.vstack((x, y, z)).T

    coords = latlon_to_vec(df["TurnLa"].values,
                           df["TurnLo"].values)

    values = df["abs_dT"].values.astype(np.float64)

    # -----------------------------
    # Precompute great-circle cosine matrix
    # -----------------------------
    print("Computing spherical distance matrix...")

    dot_matrix = np.clip(coords @ coords.T, -1.0, 1.0)
    dist_matrix = np.arccos(dot_matrix)  # radians

    print("Saving cache...")

    np.save(f"{OUTDIR}/coords.npy", coords)
    np.save(f"{OUTDIR}/values.npy", values)
    np.save(f"{OUTDIR}/dist_matrix.npy", dist_matrix)

    print("Done.")
    print("N =", len(values))
