#!/usr/bin/env python3

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

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

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

RES = 200

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

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)

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

u = np.linspace(0, 2*np.pi, RES)
v = np.linspace(0, np.pi, RES)

x = np.outer(np.cos(u), np.sin(v))
y = np.outer(np.sin(u), np.sin(v))
z = np.outer(np.ones_like(u), np.cos(v))

points = np.stack([x.flatten(), y.flatten(), z.flatten()], axis=1)

field = build_design_matrix(points) @ coeffs
field = field.reshape(x.shape)

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

# ==============================
# Plot
# ==============================

fig = plt.figure(figsize=(14, 14))

views = [
    (90, 0, "Top (Z-axis)"),
    (0, 0, "Front (Y-axis)"),
    (0, 90, "Side (X-axis)"),
    (30, 45, "Isometric")
]

for i,(elev,azim,title) in enumerate(views,1):

    ax = fig.add_subplot(2,2,i, projection='3d')

    ax.plot_surface(
        x,y,z,
        facecolors=plt.cm.coolwarm((field + max_abs)/(2*max_abs)),
        linewidth=0,
        antialiased=False,
        shade=False
    )

    # Contours projected onto surface
    ax.contour(
        x,y,z,
        field,
        levels=12,
        colors='black',
        linewidths=0.6
    )

    # Axis lines
    ax.plot([-1,1],[0,0],[0,0], color='black', linewidth=1)
    ax.plot([0,0],[-1,1],[0,0], color='black', linewidth=1)
    ax.plot([0,0],[0,0],[-1,1], color='black', linewidth=1)

    ax.set_box_aspect([1,1,1])
    ax.view_init(elev=elev, azim=azim)

    ax.set_axis_off()
    ax.set_title(title, fontsize=12)

plt.tight_layout()
plt.savefig("geometric_four_panel.png", dpi=400)
plt.show()
