#!/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 = (3840, 2160)

RADIAL_SCALE = 0.08   # reduced exaggeration
EULER_LON = 121.0
CLIP_HEMISPHERE = True
CONTOUR_LEVELS = 10

COLORMAP = "coolwarm"
BACKGROUND = "white"

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

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

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

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

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)

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

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

points = sphere.points
predicted = build_design_matrix(points) @ coeffs

# RMS-normalized exaggeration
rms = np.sqrt(np.mean(predicted**2))
radius = 1.0 + RADIAL_SCALE * (predicted / rms)

sphere.points = points * radius[:,None]
sphere["field"] = predicted

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

# ==========================
# Euler Meridian
# ==========================

lats = np.linspace(-90,90,500)
euler_pts = np.array([latlon_to_vec(lat,EULER_LON) for lat in lats])
euler_line = pv.Spline(euler_pts,500)

# ==========================
# Contours
# ==========================

contours = sphere.contour(isosurfaces=CONTOUR_LEVELS,
                          scalars="field")

# ==========================
# Optional Hemisphere Clip
# ==========================

if CLIP_HEMISPHERE:
    clipped = sphere.clip(normal=(1,0,0), origin=(0,0,0))
else:
    clipped = sphere

# ==========================
# Render
# ==========================

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

# main surface
p.add_mesh(
    clipped,
    scalars="field",
    cmap=COLORMAP,
    clim=[-max_abs,max_abs],
    smooth_shading=True,
    specular=0.25,
    ambient=0.2,
    scalar_bar_args=dict(
        title="Degree-2 Harmonic",
        vertical=False,
        position_x=0.25,
        position_y=0.05,
        width=0.5,
        height=0.07
    )
)

# contours
p.add_mesh(contours,
           color="black",
           line_width=2)

# Euler meridian
p.add_mesh(euler_line,
           color="gray",
           opacity=0.4,
           line_width=5)

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

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

print("Saved: render_degree2_advanced.png")
