#!/usr/bin/env python3

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

# ==========================
# CONFIGURATION
# ==========================

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

SPHERE_RES = 400           # Sphere resolution (higher = smoother)
RADIAL_SCALE = 0.18        # Radial exaggeration
OUTPUT_DPI = 300
IMAGE_SIZE = (3840, 2160)  # 4K

BACKGROUND = "white"
COLORMAP = "coolwarm"

# ==========================
# Utilities
# ==========================

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 dominant_quadrupole_axis(coeffs):
    b_xx, b_yy, b_zz, b_xy, b_xz, b_yz = coeffs[3:]

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

    eigvals, eigvecs = np.linalg.eigh(Q)
    idx = np.argmax(eigvals)
    return eigvecs[:, idx]


# ==========================
# 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

if header_idx is None:
    raise RuntimeError("Could not find Frost header row.")

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 = build_design_matrix(coords_fit)
coeffs, _, _, _ = np.linalg.lstsq(X, signed_vals, rcond=None)

quad_axis = dominant_quadrupole_axis(coeffs)

# ==========================
# Build Sphere Mesh
# ==========================

sphere = pv.Sphere(theta_resolution=SPHERE_RES,
                   phi_resolution=SPHERE_RES,
                   radius=1.0)

points = sphere.points
coords_eval = points.copy()

X_eval = build_design_matrix(coords_eval)
predicted = X_eval @ coeffs

# Radial deformation
radius = 1.0 + RADIAL_SCALE * predicted
sphere.points = coords_eval * radius[:, None]

sphere["degree2"] = predicted

# ==========================
# Create Plotter
# ==========================

def render_view(camera_position, filename):

    p = pv.Plotter(off_screen=True, window_size=IMAGE_SIZE)
    p.set_background(BACKGROUND)

    # Main surface
    p.add_mesh(
        sphere,
        scalars="degree2",
        cmap=COLORMAP,
        smooth_shading=True,
        lighting=True,
        specular=0.2,
        ambient=0.2,
    )

    # Quadrupole axis
    axis_line = pv.Line(-quad_axis, quad_axis)
    p.add_mesh(axis_line, color="black", line_width=6)

    # Euler meridian (121°E great circle)
    lats = np.linspace(-90, 90, 400)
    euler_points = np.array([latlon_to_vec(lat, 121.0) for lat in lats])
    euler_line = pv.Spline(euler_points, 400)
    p.add_mesh(euler_line, color="gray", line_width=4)

    p.camera_position = camera_position
    p.show(screenshot=filename)

    print("Saved:", filename)


# ==========================
# Render Views
# ==========================

render_view("iso", "render_degree2_isometric.png")
render_view("xy", "render_degree2_equatorial.png")
render_view("yz", "render_degree2_polar.png")
