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

ELEV_LEVELS = [-30, -20, -10]

# Thresholds
A_C = 0.609
A_ISO = 0.864

# RSL table
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],
])

# =========================
# LOAD REGION
# =========================
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)]
    z = np.linspace(-120, 10, 200)
    A = np.array([(flat <= zi).sum() for zi in z], dtype=float)
    A /= A.max()
    return z, A

# =========================
# CONNECTIVITY
# =========================
def connectivity(grid, zl):
    mask = grid > zl
    labeled, num = label(mask)

    if num == 0:
        return 0, 0

    sizes = np.bincount(labeled.ravel())[1:]
    return num, sizes.max() / sizes.sum()

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

    grid = load_region()
    flat = grid[np.isfinite(grid)]

    z, A = compute_hypsometry(flat)
    A_of_z = interp1d(z, A, bounds_error=False, fill_value=(0,1))

    # -------------------------
    # TIME SERIES
    # -------------------------
    t = RSL_TABLE[:,0]
    rsl = RSL_TABLE[:,1]

    t_dense = np.linspace(t.max(), t.min(), 400)
    z_t = interp1d(t, rsl)(t_dense)

    A_t = A_of_z(z_t)

    frag, largest = [], []
    for zl in z_t:
        n, lf = connectivity(grid, zl)
        frag.append(n)
        largest.append(lf)

    frag = np.array(frag)
    largest = np.array(largest)

    # -------------------------
    # FIGURE LAYOUT
    # -------------------------
    fig = plt.figure(figsize=(12,10))
    gs = fig.add_gridspec(2,2, height_ratios=[1,1])


    import geopandas as gpd
    import os
    import zipfile
    import urllib.request

    # =========================
    # NATURAL EARTH LOADER (FIXED)
    # =========================
    def load_coastlines():
        url = "https://naturalearth.s3.amazonaws.com/110m_physical/ne_110m_coastline.zip"
        zip_path = "ne_coastline.zip"
        extract_dir = "ne_coastline"

        if not os.path.exists(extract_dir):
            print("Downloading Natural Earth coastline...")
            urllib.request.urlretrieve(url, zip_path)

            with zipfile.ZipFile(zip_path, 'r') as z:
                z.extractall(extract_dir)

        shp_path = os.path.join(extract_dir, "ne_110m_coastline.shp")
        return gpd.read_file(shp_path)


    # =========================
    # PANEL A — CARTOGRAPHIC SNAPSHOTS (CORRECTED)
    # =========================
    axA = fig.add_subplot(gs[0,0])

    # --- Load coastline ---
    coast = load_coastlines()

    # Clip to region
    coast = coast.cx[LON_MIN:LON_MAX, LAT_MIN:LAT_MAX]

    # Plot present coastline
    coast.plot(ax=axA, color='black', linewidth=1.0)

    # --- Coordinate grid ---
    ny, nx = grid.shape
    lons = np.linspace(LON_MIN, LON_MAX, nx)
    lats = np.linspace(LAT_MIN, LAT_MAX, ny)

    # --- Paleo coastlines ---
    colors = ['#1f77b4', '#ff7f0e', '#2ca02c']

    for zl, c in zip(ELEV_LEVELS, colors):
        mask = (grid > zl).astype(int)

        axA.contour(
            lons, lats, mask,
            levels=[0.5],
            colors=[c],
            linewidths=1.5
        )

    # --- Formatting ---
    axA.set_xlim(LON_MIN, LON_MAX)
    axA.set_ylim(LAT_MIN, LAT_MAX)

    axA.set_xlabel("Longitude")
    axA.set_ylabel("Latitude")

    axA.set_title("A: Paleogeography vs Present Coastline")

    # --- Legend ---
    from matplotlib.lines import Line2D
    legend_lines = [
        Line2D([0],[0], color='black', lw=1, label='Present coastline'),
        Line2D([0],[0], color=colors[0], lw=1.5, label='-30 m'),
        Line2D([0],[0], color=colors[1], lw=1.5, label='-20 m'),
        Line2D([0],[0], color=colors[2], lw=1.5, label='-10 m'),
    ]
    axA.legend(handles=legend_lines, loc='lower left')



    # =========================
    # PANEL B — TIME EVOLUTION
    # =========================
    axB = fig.add_subplot(gs[0,1])

    axB.plot(t_dense, A_t, label="A(t)", linewidth=2)
    axB.plot(t_dense, largest, linestyle='--', label="Largest cluster")
    axB.plot(t_dense, frag/frag.max(), linestyle=':', label="Fragmentation")

    axB.axhline(A_C, linestyle='--', alpha=0.5)
    axB.axhline(A_ISO, linestyle='--', alpha=0.5)

    axB.set_title("B: Collapse Dynamics")
    axB.invert_xaxis()
    axB.legend()

    # =========================
    # PANEL C — SCALING
    # =========================
    axC = fig.add_subplot(gs[1,0])

    df = pd.read_csv("p_lambda_scaling_derivative_free.csv")

    axC.errorbar(df["lambda"], df["p_mean"], yerr=df["p_std"], fmt='o')
    axC.set_xlabel("λ")
    axC.set_ylabel("p")
    axC.set_title("C: Scaling of p(λ)")

    # =========================
    # PANEL D — DATA COLLAPSE
    # =========================
    axD = fig.add_subplot(gs[1,1])

    for zl in np.linspace(-100, 0, 10):
        A_slice = A_of_z(z + zl)
        z_scaled = (z - z.mean()) / z.std()
        axD.plot(z_scaled, A_slice, alpha=0.5)

    axD.set_title("D: Hypsometric Collapse")
    axD.set_xlabel("z'")
    axD.set_ylabel("A")

    # =========================
    plt.tight_layout()
    plt.savefig("dogger_nature_multiplot.png", dpi=300)
    plt.show()


if __name__ == "__main__":
    main()