#!/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)

FULL_SCALE = 0.03
CLIP_SCALE = 0.10

CONTOUR_LEVELS = 14
EULER_LON = 121.0

COLORMAP = "coolwarm"

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

# ==========================
# Base Sphere
# ==========================

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

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

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

# Deformation helper
def deform(scale):
    s = base.copy()
    s.points = base.points * (1 + scale * norm)[:,None]
    return s

full = deform(FULL_SCALE)
clip = deform(CLIP_SCALE).clip(normal=(1,0,0))

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

# Euler meridian
lats = np.linspace(-90,90,400)
euler_pts = np.array([latlon_to_vec(lat,EULER_LON)
                      for lat in lats])
euler = pv.Spline(euler_pts,400)

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

p = pv.Plotter(shape=(2,2),
               window_size=IMAGE_SIZE)

# ---- (1) Full Sphere ----
p.subplot(0,0)
p.add_mesh(full,
           scalars="field",
           cmap=COLORMAP,
           clim=[-max_abs,max_abs],
           smooth_shading=True,
           show_scalar_bar=False)
p.add_text("Full Sphere", font_size=18)
p.camera_position = [(2,2,2),(0,0,0),(0,0,1)]

# ---- (2) Clipped ----
p.subplot(0,1)
p.add_mesh(clip,
           scalars="field",
           cmap=COLORMAP,
           clim=[-max_abs,max_abs],
           smooth_shading=True,
           show_scalar_bar=False)
p.add_mesh(euler, color="gray", opacity=0.4)
p.add_text("Clipped Hemisphere", font_size=18)
p.camera_position = [(2,2,2),(0,0,0),(0,0,1)]

# ---- (3) Undeformed Field ----
p.subplot(1,0)
p.add_mesh(base,
           scalars="field",
           cmap=COLORMAP,
           clim=[-max_abs,max_abs],
           smooth_shading=False,
           show_scalar_bar=False)
p.add_text("Field on Undeformed Sphere", font_size=18)
p.camera_position = [(2,2,2),(0,0,0),(0,0,1)]

# ---- (4) Contours Only ----
p.subplot(1,1)
p.add_mesh(base,
           color="white",
           opacity=0.15)
p.add_mesh(contours,
           color="black",
           line_width=2)
p.add_text("Harmonic Contours", font_size=18)
p.camera_position = [(2,2,2),(0,0,0),(0,0,1)]

# Global scalar bar
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_four_panel_degree2.png")

print("Saved: render_four_panel_degree2.png")
