import numpy as np
import rasterio
import matplotlib.pyplot as plt
import multiprocessing as mp
from scipy.ndimage import label

# =========================
# 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 = 2.0

OUTPUT_PREFIX = "dual_threshold_final"

# =========================
# 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)))

# =========================
# WINDOW EXTRACTION
# =========================
def extract_window(ds, lat0, lon0):
    dlat = km_to_deg_lat(WINDOW_KM)
    dlon = km_to_deg_lon(WINDOW_KM, lat0)

    lat1, lat2 = lat0, lat0 + dlat
    lon1, lon2 = lon0, lon0 + dlon

    row1, col1 = ds.index(lon1, lat1)
    row2, col2 = ds.index(lon2, lat2)

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

    if (row_max - row_min) < 20 or (col_max - col_min) < 20:
        return None

    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)

    if np.isnan(grid).all():
        return None

    return grid

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

    if len(flat) < 1000:
        return None

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

    if A.max() == 0:
        return None

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

# =========================
# CONNECTIVITY (LAND)
# =========================
def compute_connectivity(grid, sea_level):
    mask = grid > sea_level  # exposed land

    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

# =========================
# PROCESS WINDOW
# =========================
def process_window(args):
    dataset_path, lat0, lon0 = args

    try:
        with rasterio.open(dataset_path) as ds:
            grid = extract_window(ds, lat0, lon0)

        if grid is None:
            return None

        flat = grid[np.isfinite(grid)]
        relief = np.std(flat)

        hyps = compute_hypsometry(flat)
        if hyps is None:
            return None

        z, A = hyps

        frag = []
        largest = []

        for zl in z:
            n, lf = compute_connectivity(grid, zl)
            frag.append(n)
            largest.append(lf)

        return {
            "z": z,
            "A": A,
            "frag": np.array(frag),
            "largest": np.array(largest),
            "relief": relief
        }

    except Exception:
        return None

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

    # -----------------------
    # Build window grid
    # -----------------------
    tasks = []

    dlat = km_to_deg_lat(STRIDE_KM)
    lat_vals = np.arange(LAT_MIN, LAT_MAX, dlat)

    for lat0 in lat_vals:
        dlon = km_to_deg_lon(STRIDE_KM, lat0)
        lon_vals = np.arange(LON_MIN, LON_MAX, dlon)

        for lon0 in lon_vals:
            tasks.append((BATHY_FILE, lat0, lon0))

    print(f"Windows: {len(tasks)}")

    # -----------------------
    # Parallel execution
    # -----------------------
    ctx = mp.get_context("spawn")
    with ctx.Pool(mp.cpu_count()) as pool:
        results = list(pool.imap_unordered(process_window, tasks, chunksize=10))

    results = [r for r in results if r is not None]

    if len(results) == 0:
        print("No valid windows.")
        return

    # =========================
    # THRESHOLD EXTRACTION
    # =========================
    A_perc, z_perc = [], []
    A_iso, z_iso = [], []
    reliefs = []

    for r in results:
        z = r["z"]
        A = r["A"]
        frag = r["frag"]
        largest = r["largest"]

        # --------------------------
        # PEROCOLATION (MAIN COLLAPSE)
        # --------------------------
        L = largest / (largest.max() + 1e-6)

        # smooth
        kernel = np.ones(5) / 5
        L_smooth = np.convolve(L, kernel, mode='same')

        # collapse rate
        dL = -np.gradient(L_smooth)

        # ignore edges (avoid artifacts)
        valid = slice(5, -5)

        idx_p = np.argmax(dL[valid]) + 5

        A_perc.append(A[idx_p])
        z_perc.append(z[idx_p])

        # --------------------------
        # ISOLATION (FRAGMENTATION)
        # --------------------------
        score = frag * (1 - largest)
        idx_i = np.argmax(score)

        A_iso.append(A[idx_i])
        z_iso.append(z[idx_i])

        reliefs.append(r["relief"])

    A_perc = np.array(A_perc)
    A_iso = np.array(A_iso)
    z_perc = np.array(z_perc)
    z_iso = np.array(z_iso)
    reliefs = np.array(reliefs)

    # =========================
    # PRINT RESULTS
    # =========================
    print("\n=== FINAL DUAL THRESHOLDS ===")
    print(f"Percolation A_c = {A_perc.mean():.3f} ± {A_perc.std():.3f}")
    print(f"Isolation  A_iso = {A_iso.mean():.3f} ± {A_iso.std():.3f}")

    # =========================
    # HISTOGRAM
    # =========================
    plt.figure()
    plt.hist(A_perc, bins=20, alpha=0.7, label="Percolation")
    plt.hist(A_iso, bins=20, alpha=0.7, label="Isolation")
    plt.xlabel("Critical area fraction")
    plt.ylabel("Frequency")
    plt.legend()
    plt.title("Dual Threshold Distribution")
    plt.tight_layout()
    plt.savefig(f"{OUTPUT_PREFIX}_hist.png", dpi=150)

    # =========================
    # RELIEF SCALING
    # =========================
    plt.figure()
    plt.scatter(reliefs, A_perc, label="Percolation", alpha=0.7)
    plt.scatter(reliefs, A_iso, label="Isolation", alpha=0.7)
    plt.xlabel("Relief σ_z")
    plt.ylabel("Critical area")
    plt.legend()
    plt.title("Thresholds vs Relief")
    plt.tight_layout()
    plt.savefig(f"{OUTPUT_PREFIX}_relief.png", dpi=150)

    # =========================
    # EXAMPLE WINDOW
    # =========================
    r0 = results[0]

    plt.figure()
    plt.plot(r0["z"], r0["largest"], label="Largest cluster")
    plt.plot(r0["z"], r0["frag"], label="Fragment count")
    plt.xlabel("Elevation (m)")
    plt.ylabel("Connectivity")
    plt.legend()
    plt.title("Connectivity Transition")
    plt.tight_layout()
    plt.savefig(f"{OUTPUT_PREFIX}_example.png", dpi=150)

    plt.show()


if __name__ == "__main__":
    main()