import numpy as np
import pandas as pd
import rasterio
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
from scipy.ndimage import label
import multiprocessing as mp

# =========================
# CONFIG
# =========================
BATHY_FILE = "SRTM15v2.tiff"

LAT_MIN, LAT_MAX = 52.0, 62.0
LON_MIN, LON_MAX = -3.0, 9.0

WINDOW_KM = 400
STRIDE_KM = 200

ELEV_MIN = -120
ELEV_MAX = 10
DZ = 1.0

# Thresholds from your result
A_C = 0.609
A_ISO = 0.864

# =========================
# RSL TABLE (Dogger)
# =========================
RSL_TABLE = np.array([
    [30.0, -115],
    [26.0, -120],
    [21.0, -120],
    [19.0, -115],
    [16.0, -100],
    [14.6, -90],
    [13.5, -70],
    [12.8, -65],
    [11.7, -60],
    [10.5, -50],
    [9.5, -40],
    [9.0, -35],
    [8.7, -30],
    [8.2, -25],
    [7.8, -15],
    [7.5, -5],
    [7.0, 0],
    [6.0, 1],
    [0.0, 0],
])

# =========================
# KM → DEG
# =========================
def km_to_deg_lat(km):
    return km / 111.0

def km_to_deg_lon(km, lat):
    return km / (111.0 * np.cos(np.radians(lat)))

# =========================
# EXTRACT REGION (FULL)
# =========================
def load_region():
    with rasterio.open(BATHY_FILE) as ds:

        row1, col1 = ds.index(LON_MIN, LAT_MIN)
        row2, col2 = ds.index(LON_MAX, LAT_MAX)

        row_min, row_max = sorted([row1, row2])
        col_min, col_max = sorted([col1, col2])

        window = rasterio.windows.Window(
            col_off=col_min,
            row_off=row_min,
            width=col_max - col_min,
            height=row_max - row_min
        )

        grid = ds.read(1, window=window)

    return grid

# =========================
# HYPSOMETRY
# =========================
def compute_hypsometry(flat):
    flat = flat[np.isfinite(flat)]
    flat = flat[(flat >= ELEV_MIN) & (flat <= ELEV_MAX)]

    z = np.arange(ELEV_MIN, ELEV_MAX, DZ)
    A = np.array([(flat <= zi).sum() for zi in z], dtype=float)

    A /= A.max()
    return z, A

# =========================
# CONNECTIVITY
# =========================
def connectivity_metrics(grid, sea_level):
    mask = grid > sea_level

    labeled, num = label(mask)

    if num == 0:
        return 0, 0

    sizes = np.bincount(labeled.ravel())[1:]
    largest = sizes.max()
    total = sizes.sum()

    return num, largest / total

# =========================
# MAIN
# =========================
def main():

    print("Loading bathymetry...")
    grid = load_region()
    flat = grid[np.isfinite(grid)]

    print("Computing hypsometry...")
    z, A = compute_hypsometry(flat)

    # Interpolators
    A_of_z = interp1d(z, A, bounds_error=False, fill_value=(0,1))
    z_of_A = interp1d(A, z, bounds_error=False)

    # RSL interpolation
    t = RSL_TABLE[:,0]
    rsl = RSL_TABLE[:,1]

    t_dense = np.linspace(t.max(), t.min(), 500)
    rsl_interp = interp1d(t, rsl, kind='linear')
    z_t = rsl_interp(t_dense)

    # Convert to area
    A_t = A_of_z(z_t)

    print("Computing connectivity over time...")

    frag_t = []
    largest_t = []

    for zl in z_t:
        n, lf = connectivity_metrics(grid, zl)
        frag_t.append(n)
        largest_t.append(lf)

    frag_t = np.array(frag_t)
    largest_t = np.array(largest_t)

    # =========================
    # FIND TRANSITIONS
    # =========================
    idx_c = np.argmin(np.abs(A_t - A_C))
    idx_iso = np.argmin(np.abs(A_t - A_ISO))

    t_c = t_dense[idx_c]
    t_iso = t_dense[idx_iso]

    print("\n=== TIMING ===")
    print(f"Connectivity collapse ~ {t_c:.2f} ka BP")
    print(f"Isolation peak      ~ {t_iso:.2f} ka BP")

    # =========================
    # FINAL FIGURE
    # =========================
    fig, ax1 = plt.subplots(figsize=(10,6))

    # Area
    ax1.plot(t_dense, A_t, label="Area fraction A(t)", linewidth=2)
    ax1.set_xlabel("Time (ka BP)")
    ax1.set_ylabel("Area fraction A")
    ax1.invert_xaxis()

    # Connectivity
    ax2 = ax1.twinx()
    ax2.plot(t_dense, largest_t, linestyle='--', label="Largest cluster")
    ax2.plot(t_dense, frag_t / frag_t.max(), linestyle=':', label="Fragmentation (norm)")

    # Threshold lines
    ax1.axhline(A_C, linestyle='--', alpha=0.5)
    ax1.axhline(A_ISO, linestyle='--', alpha=0.5)

    ax1.axvline(t_c, linestyle='--', alpha=0.5)
    ax1.axvline(t_iso, linestyle='--', alpha=0.5)

    # Labels
    ax1.text(t_c, A_C, "  Connectivity collapse", verticalalignment='bottom')
    ax1.text(t_iso, A_ISO, "  Isolation peak", verticalalignment='bottom')

    # Legend
    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()

    ax1.legend(lines1 + lines2, labels1 + labels2, loc="best")

    plt.title("Dogger Bank Collapse: Geometry-Driven Phase Transition")
    plt.tight_layout()
    plt.savefig("dogger_final_figure.png", dpi=300)

    plt.show()


if __name__ == "__main__":
    main()