#!/usr/bin/env python3

import numpy as np
import pandas as pd
import pyvista as pv

DATAFILE = "frost2021.dat"
CACHE_DIR = "analysis_cache"

SPHERE_RES = 300      # angular resolution
RADIAL_SCALE = 0.15   # exaggeration factor


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


def evaluate_degree2(coeffs, coords):
    X = build_design_matrix(coords)
    return X @ coeffs


# -------------------------------------------------
# Load dataset and fit degree-2 model
# -------------------------------------------------

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

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_fit = coords[valid]
signed_vals = signed_vals[valid]

X_fit = build_design_matrix(coords_fit)
coeffs, _, _, _ = np.linalg.lstsq(X_fit, signed_vals, rcond=None)

# -------------------------------------------------
# Build evaluation sphere
# -------------------------------------------------

phi = np.linspace(-90, 90, SPHERE_RES)
theta = np.linspace(-180, 180, SPHERE_RES)

phi_grid, theta_grid = np.meshgrid(phi, theta)

lat = np.radians(phi_grid.flatten())
lon = np.radians(theta_grid.flatten())

sphere_coords = np.vstack((
    np.cos(lat) * np.cos(lon),
    np.cos(lat) * np.sin(lon),
    np.sin(lat)
)).T

predicted = evaluate_degree2(coeffs, sphere_coords)

# -------------------------------------------------
# Create PyVista mesh
# -------------------------------------------------

radius = 1.0 + RADIAL_SCALE * predicted

points = sphere_coords * radius[:, None]

mesh = pv.PolyData(points)

mesh["degree2_field"] = predicted

mesh.save("degree2_harmonic_sphere.vtp")

print("Saved: degree2_harmonic_sphere.vtp")


import pyvista as pv
import numpy as np

axis = latlon_to_vec(-6.3, 111.7)

line = pv.Line(-axis, axis)
line.save("quadrupole_axis.vtp")
