#!/usr/bin/env python3
"""
Steps 1–3: Spatiotemporal Stability Diagnostics

Implements:
1. Time-integrated gradient magnitude of effective elevation
2. Sampling of this metric at Early Homo and Early Civilization sites
3. Time-averaged distance of sites to the equilibrium (zero) contour

Outputs:
- stability_gradient.nc
- homo_site_metrics.csv
- civilization_site_metrics.csv
"""

import json
import numpy as np
import xarray as xr
import pandas as pd
from scipy.interpolate import RegularGridInterpolator
from scipy.spatial import cKDTree
from shapely.geometry import LineString, Point
from skimage import measure

# ============================================================
# INPUT FILES
# ============================================================

MODEL_NC = "tpw_return_sequence_piecewise_angles.nc"
HOMO_GEOJSON = "early_homo_sites.geojson"
CIV_GEOJSON  = "early_civilizations.geojson"

# ============================================================
# LOAD MODEL OUTPUT
# ============================================================

ds = xr.open_dataset(MODEL_NC)
Zeff = ds["effective_elevation"].values   # (step, lat, lon)
lat = ds["lat"].values
lon = ds["lon"].values
steps = ds["step"].values

# Grid spacing (approx, meters not required for relative test)
dlat = np.abs(lat[1] - lat[0])
dlon = np.abs(lon[1] - lon[0])

# ============================================================
# STEP 1 — TIME-INTEGRATED GRADIENT MAGNITUDE
# ============================================================

grad_mag = np.zeros_like(Zeff)

for i in range(len(steps)):
    dz_dlat, dz_dlon = np.gradient(Zeff[i], dlat, dlon)
    grad_mag[i] = np.sqrt(dz_dlat**2 + dz_dlon**2)

# Time-integrated (mean) gradient magnitude
G = grad_mag.mean(axis=0)

# Save as NetCDF for inspection
xr.Dataset(
    data_vars=dict(
        stability_gradient=(("lat","lon"), G)
    ),
    coords=dict(lat=lat, lon=lon),
    attrs=dict(
        description="Time-integrated magnitude of effective elevation gradient",
        method="Mean over all relaxation steps"
    )
).to_netcdf("stability_gradient.nc")

# ============================================================
# LOAD SITE DATA
# ============================================================

def load_sites(fname):
    with open(fname) as f:
        data = json.load(f)
    sites = []
    for feat in data["features"]:
        lon_s, lat_s = feat["geometry"]["coordinates"]
        sites.append((lat_s, lon_s))
    return np.array(sites)

homo_sites = load_sites(HOMO_GEOJSON)
civ_sites  = load_sites(CIV_GEOJSON)

# ============================================================
# STEP 2 — SAMPLE GRADIENT AT SITE LOCATIONS
# ============================================================

interp_G = RegularGridInterpolator((lat, lon), G, bounds_error=False, fill_value=np.nan)

homo_G = interp_G(homo_sites)
civ_G  = interp_G(civ_sites)

# ============================================================
# STEP 3 — TIME-AVERAGED DISTANCE TO ZERO CONTOUR
# ============================================================

def extract_zero_contour(z, lat, lon):
    """
    Extract zero contour as lat/lon points.
    """
    contours = measure.find_contours(z, level=0.0)
    lines = []
    for c in contours:
        lats = np.interp(c[:,0], np.arange(len(lat)), lat)
        lons = np.interp(c[:,1], np.arange(len(lon)), lon)
        lines.append(np.column_stack([lats, lons]))
    return lines

def mean_distance_to_contour(sites, contours):
    """
    Mean distance (in degrees) from sites to nearest contour, averaged over time.
    """
    distances = np.zeros((len(contours), len(sites)))

    for i, lines in enumerate(contours):
        pts = np.vstack(lines)
        tree = cKDTree(pts)
        distances[i], _ = tree.query(sites)

    return distances.mean(axis=0), distances.var(axis=0)

# Extract zero contours for all timesteps
all_contours = [
    extract_zero_contour(Zeff[i], lat, lon)
    for i in range(len(steps))
]

# Compute distances
homo_dist_mean, homo_dist_var = mean_distance_to_contour(homo_sites, all_contours)
civ_dist_mean,  civ_dist_var  = mean_distance_to_contour(civ_sites, all_contours)

# ============================================================
# WRITE RESULTS
# ============================================================

homo_df = pd.DataFrame({
    "lat": homo_sites[:,0],
    "lon": homo_sites[:,1],
    "stability_gradient": homo_G,
    "mean_zero_contour_distance_deg": homo_dist_mean,
    "var_zero_contour_distance_deg2": homo_dist_var
})

civ_df = pd.DataFrame({
    "lat": civ_sites[:,0],
    "lon": civ_sites[:,1],
    "stability_gradient": civ_G,
    "mean_zero_contour_distance_deg": civ_dist_mean,
    "var_zero_contour_distance_deg2": civ_dist_var
})

homo_df.to_csv("homo_site_metrics.csv", index=False)
civ_df.to_csv("civilization_site_metrics.csv", index=False)

print("✓ Stability diagnostics complete")
print("  - stability_gradient.nc")
print("  - homo_site_metrics.csv")
print("  - civilization_site_metrics.csv")
