#!/usr/bin/env python3

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

# ==========================
# CONFIG
# ==========================

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

SPHERE_RES = 350
IMAGE_SIZE = (4800, 4800)

COLORMAP = "coolwarm"
BACKGROUND = "white"

# Known structural axes
EULER_121E = (0.0, 121.0)
EULER_59W = (0.0, -59.0)

LLVP_AFRICA = (-20.0, 25.0)
LLVP_PACIFIC = (-15.0, -155.0)

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

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

def build_design_matrix(coords):
    x,y,z = coords[:,0],coords[:,1],coords[:,2]
    return np.column_stack([
        x,y,z,
        x*x,y*y,z*z,
        x*y,x*z,y*z
    ])

# ==========================
# Load + Fit Degree-2
# ==========================

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

with open(DATAFILE) as f:
    lines = f.readlines()

header_idx = next(i for i,l in enumerate(lines)
                  if l.strip().startswith("EvtDate"))

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

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

vals = df["dT"].values
valid = vals != 0

X = build_design_matrix(coords[valid])
coeffs,_,_,_ = np.linalg.lstsq(X, vals[valid], rcond=None)

# Quadrupole tensor
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)

# Dipole axis (regression vector)
dipole_vec = coeffs[:3]
dipole_vec /= np.linalg.norm(dipole_vec)

# ==========================
# Sphere
# ==========================

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

field = build_design_matrix(sphere.points) @ coeffs
sphere["field"] = field

max_abs = np.max(np.abs(field))

# ==========================
# Plotter
# ==========================

p = pv.Plotter(window_size=IMAGE_SIZE)
p.set_background(BACKGROUND)

# Base sphere
p.add_mesh(sphere,
           scalars="field",
           cmap=COLORMAP,
           clim=[-max_abs,max_abs],
           smooth_shading=True,
           opacity=0.95)

# --------------------------
# Quadrupole axes
# --------------------------

for i in range(3):
    axis = eigvecs[:,i]
    p.add_mesh(pv.Line(-axis, axis),
               color="black",
               line_width=4)

# --------------------------
# Dipole axis
# --------------------------

p.add_mesh(pv.Line(-dipole_vec, dipole_vec),
           color="red",
           line_width=6)

# --------------------------
# Euler points
# --------------------------

for lat,lon in [EULER_121E, EULER_59W]:
    vec = latlon_to_vec(lat,lon)
    p.add_mesh(pv.Sphere(radius=0.05, center=vec),
               color="yellow")

# --------------------------
# LLVP Caps (projected)
# --------------------------

for lat,lon in [LLVP_AFRICA, LLVP_PACIFIC]:
    vec = latlon_to_vec(lat,lon)
    p.add_mesh(pv.Sphere(radius=0.12, center=vec),
               color="purple",
               opacity=0.6)

# --------------------------

p.camera_position = [(2.5,2.5,2.5),(0,0,0),(0,0,1)]

p.add_scalar_bar(title="Degree-2 Harmonic",
                 vertical=False,
                 position_x=0.25,
                 position_y=0.02,
                 width=0.5,
                 height=0.03)

p.show(screenshot="render_structural_schematic.png")

print("Saved: render_structural_schematic.png")
