#!/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 = 400
IMAGE_SIZE = (4200, 4200)

# Muted colormap
COLORMAP = "coolwarm"
BACKGROUND = "white"

# Structural reference points
EULER_121E = (0.0, 121.0)
EULER_59W  = (0.0, -59.0)

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

# Meridian lines to draw (degrees)
STRUCTURAL_MERIDIANS = [121, -59, 0, 90]

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

def build_meridian(lon_deg, radius=1.02):
    lats = np.linspace(-90, 90, 500)
    pts = np.array([latlon_to_vec(lat, lon_deg) for lat in lats])
    return pv.Spline(pts * radius)

# ==============================
# Load + Fit
# ==============================

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_vec = coeffs[:3]
dipole_vec /= np.linalg.norm(dipole_vec)

# ==============================
# Sphere + Field
# ==============================

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)

# Matte rendering
p.enable_anti_aliasing()

p.add_mesh(
    sphere,
    scalars="field",
    cmap=COLORMAP,
    clim=[-max_abs, max_abs],
    smooth_shading=False,         # matte
    specular=0.0,                 # no highlights
    diffuse=1.0,
    ambient=0.6,
    opacity=1.0
)

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

for axis in eigvecs.T:
    p.add_mesh(
        pv.Line(-axis, axis),
        color="black",
        line_width=2
    )

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

p.add_mesh(
    pv.Line(-dipole_vec, dipole_vec),
    color="black",
    line_width=3
)

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

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

# ------------------------------
# LLVP caps (projected)
# ------------------------------

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

# ------------------------------
# Structural Meridians
# ------------------------------

for lon in STRUCTURAL_MERIDIANS:
    p.add_mesh(
        build_meridian(lon),
        color="black",
        line_width=1
    )

# ------------------------------
# Camera: Revelatory Perspectives
# ------------------------------

# Perspective 1: Quadrupole alignment
p.camera_position = [
    (2.4, 1.8, 1.8),
    (0,0,0),
    (0,0,1)
]

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

# Perspective 2: Along Euler meridian
p.camera_position = [
    (2.0*np.cos(np.radians(121)),
     2.0*np.sin(np.radians(121)),
     0.8),
    (0,0,0),
    (0,0,1)
]

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

# Perspective 3: Polar structural
p.camera_position = [
    (0,0,2.5),
    (0,0,0),
    (1,0,0)
]

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

print("Saved minimalist structural plates.")
