#!/usr/bin/env python3

import numpy as np

CACHE_DIR = "analysis_cache"


def latlon_to_vec(lat_deg, lon_deg):
    lat = np.radians(lat_deg)
    lon = np.radians(lon_deg)
    return np.array([
        np.cos(lat) * np.cos(lon),
        np.cos(lat) * np.sin(lon),
        np.sin(lat)
    ])


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 angular_sep(v1, v2):
    dot = np.clip(np.dot(v1, v2), -1.0, 1.0)
    return np.degrees(np.arccos(dot))


def build_design_matrix(coords):
    x = coords[:,0]
    y = coords[:,1]
    z = coords[:,2]

    return np.column_stack([
        x, y, z,
        x*x, y*y, z*z,
        x*y, x*z, y*z
    ])


if __name__ == "__main__":

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

    # Reload signed dT
    import pandas as pd
    DATAFILE = "frost2021.dat"

    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
    )

    df["dT"] = df["dTabdf"].copy()
    mask = df["dT"] == 0
    df.loc[mask, "dT"] = df.loc[mask, "dTbcdf"]

    signed_vals = df["dT"].values

    valid = signed_vals != 0
    coords = coords[valid]
    signed_vals = signed_vals[valid]

    X = build_design_matrix(coords)

    coeffs, _, _, _ = np.linalg.lstsq(X, signed_vals, rcond=None)

    # Quadratic coefficients
    b_xx, b_yy, b_zz, b_xy, b_xz, b_yz = coeffs[3:]

    # Build symmetric tensor Q
    Q = np.array([
        [b_xx, b_xy/2, b_xz/2],
        [b_xy/2, b_yy, b_yz/2],
        [b_xz/2, b_yz/2, b_zz]
    ])

    # Eigen decomposition
    eigvals, eigvecs = np.linalg.eigh(Q)
    idx = np.argsort(eigvals)
    eigvals = eigvals[idx]
    eigvecs = eigvecs[:, idx]

    print("Quadrupole eigenvalues:", eigvals)

    print("\nPrincipal axes:")

    axes = []
    for i in range(3):
        lat, lon = vec_to_latlon(eigvecs[:, i])
        axes.append((lat, lon))
        print(f"Axis {i+1}: lat {lat:.3f}, lon {lon:.3f}")

    # Compare to Euler candidate
    euler_vec = latlon_to_vec(0.0, -59.0)

    print("\nAngular separation from 59°W Euler axis:")
    for i, (lat, lon) in enumerate(axes):
        v = latlon_to_vec(lat, lon)
        angle = angular_sep(v, euler_vec)
        print(f"Axis {i+1}: {angle:.2f}°")
