import numpy as np
import rasterio
import matplotlib.pyplot as plt
import pandas as pd
from scipy.ndimage import label
from multiprocessing import get_context
import os

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

BATHY_FILE = "SRTM15v2.tiff"
OUT_DIR = "outputs"
os.makedirs(OUT_DIR, exist_ok=True)

# Dogger Bank bounding box
LAT_MIN, LAT_MAX = 52.0, 62.0
LON_MIN, LON_MAX = -3.0, 9.0

# Time grid (ka BP)
TIME_KA = np.linspace(30, 0, 301)

# Global shared bathy (for workers)
GLOBAL_BATHY = None


# ==============================
# RSL MODEL (FIXED + TOTAL)
# ==============================

def rsl_curve(t_ka):
    """
    Robust Dogger Bank RSL model
    Always returns float
    """

    # --- Base eustatic ---
    if t_ka >= 19:
        rsl = -120.0

    elif 14.6 <= t_ka < 19:
        rsl = -120.0 + (19.0 - t_ka) * 4.0

    elif 11.7 <= t_ka < 14.6:
        rsl = -90.0 + (14.6 - t_ka) * 8.0

    elif 8.2 <= t_ka < 11.7:
        rsl = -60.0 + (11.7 - t_ka) * 10.0

    elif 7.0 <= t_ka < 8.2:
        rsl = -25.0 + (8.2 - t_ka) * 15.0

    else:  # t_ka < 7.0
        rsl = -5.0 + (7.0 - t_ka) * 2.0

    # --- Subsidence ---
    subsidence = 10.0 * np.exp(-t_ka / 10.0)

    # --- Meltwater pulses ---
    mwp1a = 10.0 * np.exp(-((t_ka - 14.6)**2) / (2 * 0.5**2))
    event82 = 6.0 * np.exp(-((t_ka - 8.2)**2) / (2 * 0.3**2))

    return float(rsl + subsidence + mwp1a + event82)


# ==============================
# LOAD BATHYMETRY (WINDOWED)
# ==============================

def load_bathymetry():
    with rasterio.open(BATHY_FILE) as src:

        row_min, col_min = src.index(LON_MIN, LAT_MAX)
        row_max, col_max = src.index(LON_MAX, LAT_MIN)

        row_start, row_stop = sorted([row_min, row_max])
        col_start, col_stop = sorted([col_min, col_max])

        window = ((row_start, row_stop), (col_start, col_stop))

        bathy = src.read(1, window=window)

        print(f"Loaded subset shape: {bathy.shape}")

        return bathy


# ==============================
# MULTIPROCESS INITIALIZER (CRITICAL FIX)
# ==============================

def init_worker(bathy):
    global GLOBAL_BATHY
    GLOBAL_BATHY = bathy


# ==============================
# WORKER FUNCTION
# ==============================

def analyze_time_step(t):
    bathy = GLOBAL_BATHY

    if bathy is None:
        raise RuntimeError("GLOBAL_BATHY not initialized")

    rsl = rsl_curve(t)

    if rsl is None or np.isnan(rsl):
        raise ValueError(f"Invalid RSL at t={t}")

    land = (bathy > rsl)

    area = np.sum(land)

    _, num_features = label(land)

    return t, area, num_features


# ==============================
# RUN ANALYSIS
# ==============================

def run_analysis(bathy):

    ctx = get_context("spawn")  # Apple-safe

    with ctx.Pool(initializer=init_worker, initargs=(bathy,)) as pool:
        results = pool.map(analyze_time_step, TIME_KA)

    results.sort(key=lambda x: x[0])

    time = np.array([r[0] for r in results])
    area = np.array([r[1] for r in results])
    fragments = np.array([r[2] for r in results])

    return time, area, fragments


# ==============================
# STEP MODEL (YOUR HALVING MODEL)
# ==============================

def step_model(time):
    A0 = 1.0
    A = np.zeros_like(time)

    for i, t in enumerate(time):
        if t > 10:
            A[i] = A0
        elif t > 8.5:
            A[i] = A0 / 2
        elif t > 7.5:
            A[i] = A0 / 4
        else:
            A[i] = A0 / 8

    return A


# ==============================
# MAIN
# ==============================

def main():

    print("Loading bathymetry...")
    bathy = load_bathymetry()

    print("Running analysis...")
    time, area, fragments = run_analysis(bathy)

    # Normalize
    area_norm = area / np.max(area)

    # Step model
    step = step_model(time)

    # ==============================
    # SAVE DATA
    # ==============================

    pd.DataFrame({
        "time_ka": time,
        "area": area,
        "area_norm": area_norm
    }).to_csv(f"{OUT_DIR}/area_timeseries.csv", index=False)

    pd.DataFrame({
        "time_ka": time,
        "fragments": fragments
    }).to_csv(f"{OUT_DIR}/fragmentation_timeseries.csv", index=False)

    # ==============================
    # PLOTS
    # ==============================

    plt.figure(figsize=(8, 5))
    plt.plot(time, area_norm, label="Observed")
    plt.plot(time, step, linestyle="--", label="Step Model")
    plt.gca().invert_xaxis()
    plt.xlabel("Time (ka BP)")
    plt.ylabel("Normalized Area")
    plt.title("Dogger Bank Area Decay")
    plt.legend()
    plt.savefig(f"{OUT_DIR}/area_vs_time.png", dpi=200)

    plt.figure(figsize=(8, 5))
    plt.plot(time, fragments)
    plt.gca().invert_xaxis()
    plt.xlabel("Time (ka BP)")
    plt.ylabel("Connected Components")
    plt.title("Fragmentation Over Time")
    plt.savefig(f"{OUT_DIR}/fragmentation_vs_time.png", dpi=200)

    print("\nDone.")
    print(f"Outputs saved to: {OUT_DIR}")


# ==============================
# ENTRY POINT (REQUIRED)
# ==============================

if __name__ == "__main__":
    main()