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

COLORMAP = "coolwarm"
BACKGROUND = "white"

GRATICULE_STEP = 10  # degrees

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

def latlon_to_vec(lat_deg, lon_deg):
    lat = np.radians(lat_deg)
    lon = np.radians(lon_deg)
    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 dominant_quadrupole_axis(coeffs):
    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)
    return eigvecs[:, np.argmax(eigvals)]

# ==========================
# Fit Model
# ==========================

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

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

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

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

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

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

coords_fit = coords[valid]
signed_vals = signed_vals[valid]

X = build_design_matrix(coords_fit)
coeffs, _, _, _ = np.linalg.lstsq(X, signed_vals, rcond=None)

quad_axis = dominant_quadrupole_axis(coeffs)

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

# Symmetric scaling
max_abs = np.max(np.abs(predicted))

radius = 1.0 + RADIAL_SCALE * predicted
sphere.points = points * radius[:, None]
sphere["degree2"] = predicted

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

# ==========================
# Graticule
# ==========================

def make_graticule():
    lines = []

    # latitude circles
    for lat in np.arange(-80, 90, GRATICULE_STEP):
        pts = [latlon_to_vec(lat, lon) for lon in np.linspace(-180, 180, 400)]
        lines.append(pv.Spline(np.array(pts), 400))

    # longitude circles
    for lon in np.arange(-180, 180, GRATICULE_STEP):
        pts = [latlon_to_vec(lat, lon) for lat in np.linspace(-90, 90, 400)]
        lines.append(pv.Spline(np.array(pts), 400))

    return lines

graticule_lines = make_graticule()

# ==========================
# Render Function
# ==========================

def render(camera_position, filename):

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

    # faint reference sphere
    p.add_mesh(ref_sphere,
               color="lightgray",
               opacity=0.15,
               smooth_shading=True)

    # main harmonic sphere
    p.add_mesh(
        sphere,
        scalars="degree2",
        cmap=COLORMAP,
        clim=[-max_abs, max_abs],
        smooth_shading=True,
        specular=0.2,
        ambient=0.2,
        scalar_bar_args=dict(
            title="Degree-2 Field",
            vertical=False,
            position_x=0.25,
            position_y=0.05,
            width=0.5,
            height=0.08
        )
    )

    # graticule
    for line in graticule_lines:
        p.add_mesh(line, color="gray", line_width=1)

    # equator emphasized
    equator_pts = [latlon_to_vec(0, lon) for lon in np.linspace(-180, 180, 400)]
    p.add_mesh(pv.Spline(np.array(equator_pts), 400),
               color="black",
               line_width=3)

    # principal axis
    p.add_mesh(pv.Line(-quad_axis, quad_axis),
               color="black",
               line_width=5)

    p.camera_position = camera_position
    p.show(screenshot=filename)
    print("Saved:", filename)


# ==========================
# Camera Views
# ==========================

render([(2,2,2),(0,0,0),(0,0,1)],
       "render_degree2_iso.png")

render([(0,3,0),(0,0,0),(0,0,1)],
       "render_degree2_equatorial.png")

render([(0,0,3),(0,0,0),(0,1,0)],
       "render_degree2_polar.png")
