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

FULL_SCALE = 0.04     # very subtle for full sphere
CLIP_SCALE = 0.12     # stronger for clipped hemisphere

EULER_LON = 121.0
CONTOUR_LEVELS = 12

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

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

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

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

# ==========================
# Create Deformed Versions
# ==========================

def deform_sphere(scale):
    sphere = base_sphere.copy()
    norm_field = field / max_abs
    radius = 1.0 + scale * norm_field
    sphere.points = base_sphere.points * radius[:,None]
    return sphere

full_sphere = deform_sphere(FULL_SCALE)
clip_sphere = deform_sphere(CLIP_SCALE).clip(normal=(1,0,0), origin=(0,0,0))

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

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

# ==========================
# Plotter with 2 Subplots
# ==========================

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

# --------------------------
# LEFT: Full Sphere
# --------------------------

p.subplot(0,0)

p.add_mesh(full_sphere,
           scalars="field",
           cmap=COLORMAP,
           clim=[-max_abs,max_abs],
           smooth_shading=True,
           specular=0.25,
           ambient=0.2,
           show_scalar_bar=False)

p.add_mesh(euler_line,
           color="gray",
           opacity=0.4,
           line_width=4)

p.camera_position = [(2,2,2),(0,0,0),(0,0,1)]
p.add_text("Full Sphere (Subtle Deformation)", font_size=16)

# --------------------------
# RIGHT: Clipped Hemisphere
# --------------------------

p.subplot(0,1)

p.add_mesh(clip_sphere,
           scalars="field",
           cmap=COLORMAP,
           clim=[-max_abs,max_abs],
           smooth_shading=True,
           specular=0.3,
           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.06
           ))

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

p.add_mesh(euler_line,
           color="gray",
           opacity=0.4,
           line_width=4)

p.camera_position = [(2,2,2),(0,0,0),(0,0,1)]
p.add_text("Clipped Hemisphere (Enhanced Deformation)", font_size=16)

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

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

print("Saved: render_dual_panel_degree2.png")
