import numpy as np
import rasterio
import pandas as pd
from scipy.optimize import curve_fit
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_SIZES_KM = [100, 200, 400, 800]
STRIDE_KM = 100

ELEV_MIN = -120
ELEV_MAX = 10
DZ = 2.0

MIN_POINTS = 40
MIN_AREA_STD = 0.005

OUTPUT_CSV = "p_lambda_scaling_derivative_free.csv"

# =========================
# MODEL: Integrated solution
# =========================
def A_model(z, k, p, z0):
    """
    Solution to:
    dA/dz = -k A (1-A)^p

    Approximated as generalized logistic curve
    """
    # generalized logistic form
    return 1.0 / (1.0 + np.exp(k * (z - z0)))**(1/p)

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

    if len(elev) < 1000:
        return None

    z_bins = np.arange(ELEV_MIN, ELEV_MAX, DZ)

    counts = np.array([(elev <= z).sum() for z in z_bins], dtype=float)

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

    A = counts / counts.max()

    if np.std(A) < MIN_AREA_STD:
        return None

    return z_bins, A

# =========================
# FIT (DERIVATIVE-FREE)
# =========================
def fit_p(z, A):

    # Smooth slightly
    A_smooth = np.convolve(A, np.ones(5)/5, mode='same')

    mask = (A_smooth > 0.02) & (A_smooth < 0.98)

    if mask.sum() < MIN_POINTS:
        return None

    z_fit = z[mask]
    A_fit = A_smooth[mask]

    try:
        popt, _ = curve_fit(
            A_model,
            z_fit,
            A_fit,
            p0=[0.05, 0.8, -40],   # k, p, z0
            bounds=([0.001, 0.1, -200],
                    [1.0,   2.0,  50]),
            maxfev=20000
        )

        k, p, z0 = popt

        return k, p

    except:
        return None

# =========================
# 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 PROCESSOR
# =========================
def process_window(args):
    dataset_path, lat0, lon0, win_km = args

    try:
        with rasterio.open(dataset_path) as ds:

            dlat = km_to_deg_lat(win_km)
            dlon = km_to_deg_lon(win_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) < 10 or (col_max - col_min) < 10:
                return None

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

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

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

        z, A = hyps

        fit = fit_p(z, A)
        if fit is None:
            return None

        k, p = fit

        return {
            "lat": lat0,
            "lon": lon0,
            "window_km": win_km,
            "lambda": win_km,
            "k": k,
            "p": p
        }

    except:
        return None

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

    tasks = []

    print("Building window grid...")

    for win_km in WINDOW_SIZES_KM:

        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, win_km))

    print(f"Total windows: {len(tasks)}")

    # =========================
    # PARALLEL
    # =========================
    ctx = mp.get_context("spawn")
    with ctx.Pool(processes=mp.cpu_count()) as pool:
        results = list(pool.imap_unordered(process_window, tasks, chunksize=20))

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

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

    df = pd.DataFrame(results)

    # Normalize λ
    Lc = min(WINDOW_SIZES_KM)
    df["lambda"] = df["window_km"] / Lc

    # Aggregate
    summary = df.groupby("lambda").agg({
        "p": ["mean", "std", "count"],
        "k": ["mean", "std"]
    }).reset_index()

    summary.columns = ["lambda", "p_mean", "p_std", "n", "k_mean", "k_std"]

    print("\n=== DERIVATIVE-FREE SCALING ===")
    print(summary)

    # Save
    df.to_csv("p_lambda_raw_derivative_free.csv", index=False)
    summary.to_csv(OUTPUT_CSV, index=False)

    print("\nSaved:")
    print("- p_lambda_raw_derivative_free.csv")
    print(f"- {OUTPUT_CSV}")

# =========================
# ENTRY
# =========================
if __name__ == "__main__":
    main()