#!/usr/bin/env python3
"""
Centrifugal Viscoelastic Return Sequence After 104° True Polar Wander

Updates in this version:
- Logarithmic viscoelastic relaxation normalized to reach
  99% by 1.8 ka BP (10.2 ka since start)
- TPW angular plateaus unchanged
- Continuous relaxation during plateaus
"""

import json
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from pathlib import Path

# ============================================================
# CONFIGURATION
# ============================================================

GEBCO_NC = "GEBCO_2025_sub_ice.nc"
OUTPUT_NC = "tpw_return_sequence_piecewise_angles.nc"
PNG_DIR = Path("tpw_return_png_piecewise")
PNG_DIR.mkdir(exist_ok=True)

HOMO_GEOJSON = "early_homo_sites.geojson"
CIV_GEOJSON  = "early_civilizations.geojson"

# Physical constants
R = 6371000.0
OMEGA = 7.2921159e-5
G = 9.81

# TPW geometry
TPW_MERIDIAN_DEG = 31.0

# Temporal / relaxation model
N_STEPS = 12
T_YEARS = 12_000.0
TAU_YEARS = 500.0

# Relaxation constraint:
# 99% achieved by 1.8 ka BP → 10.2 ka since start
RELAX_FINAL = 0.99
RELAX_T99_YEARS = 10_200.0

# Numerical
DOWNSAMPLE = 60

# Fixed color scale (meters)
VMIN = -15000.0
VMAX =  15000.0
LEVELS = np.linspace(VMIN, VMAX, 31)

# ============================================================
# LOAD HUMAN DATA
# ============================================================

def load_points(fname):
    lons, lats = [], []
    with open(fname) as f:
        data = json.load(f)
    for feat in data["features"]:
        lon, lat = feat["geometry"]["coordinates"]
        lons.append(lon)
        lats.append(lat)
    return np.array(lons), np.array(lats)

homo_lons, homo_lats = load_points(HOMO_GEOJSON)
civ_lons,  civ_lats  = load_points(CIV_GEOJSON)

# ============================================================
# VISCOELASTIC RELAXATION (LOGARITHMIC, RENORMALIZED)
# ============================================================

def log_viscoelastic_relaxation(t_years, tau, T99):
    """
    Logarithmic relaxation normalized so that:
    alpha(T99) = 1.0
    """
    alpha = np.log1p(t_years / tau) / np.log1p(T99 / tau)
    return np.clip(alpha, 0.0, 1.0)

# ============================================================
# TPW ANGULAR HISTORY
# ============================================================

def tpw_angle_from_time(t_ka):
    if t_ka >= 11.1:
        return 104.0
    elif t_ka >= 7.5:
        return 34.6
    elif t_ka >= 3.8:
        return 11.5
    else:
        return 0.0

# ============================================================
# LOAD & DOWNSAMPLE GEBCO
# ============================================================

ds = xr.open_dataset(GEBCO_NC, engine="netcdf4")

for v in ds.data_vars:
    if "elev" in v.lower() or v.lower() in ("z", "bedrock", "height"):
        elev_var = v
        break
else:
    raise RuntimeError("Elevation variable not found")

Z = ds[elev_var].isel(
    lat=slice(None, None, DOWNSAMPLE),
    lon=slice(None, None, DOWNSAMPLE)
)

lat = Z.lat.values
lon = Z.lon.values
lon2d, lat2d = np.meshgrid(lon, lat)

# ============================================================
# GEOMETRY & PHYSICS
# ============================================================

def colatitude(lon, lat, pole_lon, pole_colat):
    lon = np.deg2rad(lon)
    lat = np.deg2rad(lat)
    pole_lon = np.deg2rad(pole_lon)
    pole_lat = np.pi/2 - np.deg2rad(pole_colat)
    return np.arccos(
        np.sin(lat) * np.sin(pole_lat)
        + np.cos(lat) * np.cos(pole_lat) * np.cos(lon - pole_lon)
    )

def centrifugal_potential(theta):
    return 0.5 * OMEGA**2 * R**2 * np.sin(theta)**2

# ============================================================
# BUILD SEQUENCE (CORRECT PHYSICS)
# ============================================================

steps = np.arange(N_STEPS + 1)
times_years = np.linspace(0.0, T_YEARS, N_STEPS + 1)
times_ka = T_YEARS / 1000.0 - times_years / 1000.0

# Continuous relaxation reaching 99% by 1.8 ka BP
relax_frac = RELAX_FINAL * log_viscoelastic_relaxation(
    times_years, TAU_YEARS, RELAX_T99_YEARS
)

tpw_angles = np.array([tpw_angle_from_time(t) for t in times_ka])

theta_S1 = colatitude(lon2d, lat2d, 0.0, 0.0)
Phi_S1 = centrifugal_potential(theta_S1)

theta_init = colatitude(lon2d, lat2d, TPW_MERIDIAN_DEG, 104.0)
Phi_init = centrifugal_potential(theta_init)

Zeff = np.zeros((len(steps), lat.size, lon.size), dtype=np.float32)

for i in steps:
    theta_target = colatitude(
        lon2d, lat2d,
        TPW_MERIDIAN_DEG,
        tpw_angles[i]
    )
    Phi_target = centrifugal_potential(theta_target)

    Phi_eff = (
        (1.0 - relax_frac[i]) * Phi_init
        + relax_frac[i] * Phi_target
    )

    Zeff[i] = Z.values + (Phi_S1 - Phi_eff) / G

# ============================================================
# WRITE NETCDF
# ============================================================

xr.Dataset(
    data_vars=dict(effective_elevation=(("step","lat","lon"), Zeff)),
    coords=dict(step=steps, time_ka=("step", times_ka), lat=lat, lon=lon),
    attrs=dict(
        method="Piecewise TPW with logarithmic relaxation (99% by 1.8 ka BP)",
        relaxation_tau_years=TAU_YEARS,
        relaxation_T99_years=RELAX_T99_YEARS
    )
).to_netcdf(
    OUTPUT_NC,
    encoding={"effective_elevation": {"zlib": True, "complevel": 4}}
)

# ============================================================
# DIAGNOSTIC PLOT
# ============================================================

plt.figure(figsize=(7,5))
plt.plot(tpw_angles, relax_frac * 100.0, marker="o")
plt.xlabel("TPW reorientation angle (degrees)")
plt.ylabel("Viscoelastic relaxation (%)")
plt.title("Viscoelastic Relaxation vs TPW Reorientation Angle")
plt.grid(alpha=0.3)
plt.savefig("viscoelastic_relaxation_vs_tpw_angle.png",
            dpi=200, bbox_inches="tight")
plt.close()

# ============================================================
# PNG MAP OUTPUT
# ============================================================

for i in steps:
    fig = plt.figure(figsize=(14,7))
    ax = plt.axes(projection=ccrs.Robinson())
    ax.set_global()

    cf = ax.contourf(
        lon, lat, Zeff[i],
        levels=LEVELS,
        cmap="RdBu_r",
        vmin=VMIN, vmax=VMAX,
        transform=ccrs.PlateCarree(),
        extend="both"
    )

    ax.contour(
        lon, lat, Zeff[i],
        levels=[0.0],
        colors="black",
        linewidths=1.2,
        transform=ccrs.PlateCarree()
    )

    ax.scatter(homo_lons, homo_lats, s=35, color="black",
               edgecolor="white", linewidth=0.5,
               transform=ccrs.PlateCarree(), zorder=5,
               label="Early Homo")

    ax.scatter(civ_lons, civ_lats, s=40, color="gold",
               edgecolor="black", linewidth=0.5,
               transform=ccrs.PlateCarree(), zorder=6,
               label="Early Civilizations")

    ax.add_feature(cfeature.COASTLINE, linewidth=0.6)
    ax.add_feature(cfeature.BORDERS, linewidth=0.3)

    ax.set_title(
        f"{times_ka[i]:.1f} ka | "
        f"TPW target {tpw_angles[i]:.1f}° | "
        f"{relax_frac[i]*100:.1f}% relaxation"
    )

    cbar = plt.colorbar(cf, orientation="horizontal", pad=0.05)
    cbar.set_label("Effective elevation relative to equilibrium sea level (m)")
    cbar.set_ticks(np.linspace(VMIN, VMAX, 7))

    ax.legend(loc="lower left", framealpha=0.9)

    plt.savefig(PNG_DIR / f"tpw_piecewise_step_{i:02d}.png",
                dpi=200, bbox_inches="tight")
    plt.close()
