#!/usr/bin/env python3

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature

from matplotlib.colors import TwoSlopeNorm

# =========================
# CONFIG
# =========================

DATAFILE = "frost2021.dat"
CACHE_DIR = "analysis_cache"

OUTPUT_PREFIX = "fig_"
DPI = 300


# =========================
# Utility Functions
# =========================

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 vec_to_latlon(v):
    v = v / np.linalg.norm(v)
    lat = np.degrees(np.arcsin(v[2]))
    lon = np.degrees(np.arctan2(v[1], v[0]))
    return lat, lon


def build_design_matrix(coords):
    x = coords[:,0]
    y = coords[:,1]
    z = 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)
    idx = np.argmax(eigvals)
    return eigvecs[:, idx]


def plot_meridian(ax, lon, color="black", lw=2, label=None):
    lats = np.linspace(-90, 90, 361)
    lons = np.full_like(lats, lon)
    ax.plot(lons, lats, transform=ccrs.PlateCarree(),
            color=color, linewidth=lw, label=label)


def plot_axis_line(ax, vec, color="yellow", lw=2, label=None):
    lat, lon = vec_to_latlon(vec)
    plot_meridian(ax, lon, color=color, lw=lw, label=label)


# =========================
# Load Data
# =========================

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 = coords[valid]
signed_vals = signed_vals[valid]
df = df.iloc[valid]

lat = df["TurnLa"].values
lon = df["TurnLo"].values

# =========================
# Fit Degree-2 Model
# =========================

X = build_design_matrix(coords)
coeffs, _, _, _ = np.linalg.lstsq(X, signed_vals, rcond=None)
predicted = X @ coeffs
residual = signed_vals - predicted

quad_axis = dominant_quadrupole_axis(coeffs)

# Known axes for overlay
euler_axis = latlon_to_vec(0.0, 121.0)
amp_axis = latlon_to_vec(-6.8, 105.0)
sampling_axis = latlon_to_vec(24.4, -89.1)

# =========================
# Common Map Setup
# =========================

def base_map(title):
    fig = plt.figure(figsize=(10,5))
    ax = plt.axes(projection=ccrs.Robinson())
    ax.set_global()
    ax.coastlines(linewidth=0.5)
    ax.add_feature(cfeature.BORDERS, linewidth=0.3)
    ax.set_title(title)
    return fig, ax


norm = TwoSlopeNorm(vcenter=0)

# =========================
# FIGURE 1 — Observed Field
# =========================

fig, ax = base_map("Observed Signed dT")

sc = ax.scatter(
    lon, lat,
    c=signed_vals,
    cmap="RdBu_r",
    norm=norm,
    s=6,
    transform=ccrs.PlateCarree()
)

plot_axis_line(ax, quad_axis, "yellow", 2, "Quadrupole Axis")
plot_axis_line(ax, euler_axis, "black", 1.5, "121°E Euler")

plt.colorbar(sc, orientation="horizontal", pad=0.05)
plt.legend()
plt.savefig(f"{OUTPUT_PREFIX}01_observed.png", dpi=DPI)
plt.close()

# =========================
# FIGURE 2 — Predicted Field
# =========================

fig, ax = base_map("Degree-2 Predicted Field")

sc = ax.scatter(
    lon, lat,
    c=predicted,
    cmap="RdBu_r",
    norm=norm,
    s=6,
    transform=ccrs.PlateCarree()
)

plot_axis_line(ax, quad_axis, "yellow", 2)
plot_axis_line(ax, euler_axis, "black", 1.5)

plt.colorbar(sc, orientation="horizontal", pad=0.05)
plt.savefig(f"{OUTPUT_PREFIX}02_predicted.png", dpi=DPI)
plt.close()

# =========================
# FIGURE 3 — Residual Field
# =========================

fig, ax = base_map("Residual (Observed - Degree-2)")

sc = ax.scatter(
    lon, lat,
    c=residual,
    cmap="RdBu_r",
    norm=norm,
    s=6,
    transform=ccrs.PlateCarree()
)

plot_axis_line(ax, quad_axis, "yellow", 2)
plt.colorbar(sc, orientation="horizontal", pad=0.05)
plt.savefig(f"{OUTPUT_PREFIX}03_residual.png", dpi=DPI)
plt.close()

# =========================
# FIGURE 4 — Harmonic Axis Comparison
# =========================

fig, ax = base_map("Principal Axes Comparison")

plot_axis_line(ax, quad_axis, "yellow", 2, "Quadrupole Axis")
plot_axis_line(ax, euler_axis, "black", 2, "121°E Euler")
plot_axis_line(ax, amp_axis, "cyan", 2, "Amplitude PCA Axis")
plot_axis_line(ax, sampling_axis, "red", 2, "Sampling Axis")

plt.legend()
plt.savefig(f"{OUTPUT_PREFIX}04_axis_comparison.png", dpi=DPI)
plt.close()

print("All figures generated.")
