import numpy as np
import rasterio
import matplotlib.pyplot as plt
import pandas as pd
import os

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

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

BATHY_FILE = "SRTM15v2.tiff"

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

# Elevation step (meters)
DELTA_Z = 5.0

# Mask to isolate Dogger plateau
USE_SHALLOW_MASK = True
Z_MIN_MASK = -60
Z_MAX_MASK = 20

# Stability thresholds
A_MIN = 0.05
A_MAX = 0.95

# ==============================
# LOAD BATHYMETRY
# ==============================

def load_bathy():
    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)

    return bathy

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

# ==============================
# PREPARE DATA
# ==============================

z = bathy.flatten()
z = z[~np.isnan(z)]

if USE_SHALLOW_MASK:
    z = z[(z > Z_MIN_MASK) & (z < Z_MAX_MASK)]

N = len(z)
print(f"Valid cells: {N}")

# ==============================
# HYPSOMETRY (FAST + CORRECT)
# ==============================

print("Computing hypsometry...")

# Sort elevations (ascending)
z_sorted = np.sort(z)

z_min = z_sorted[0]
z_max = z_sorted[-1]

# Elevation grid
z_grid = np.arange(z_max, z_min, -1.0)

# Use searchsorted (vectorized thresholding)
# A(z) = fraction of cells > z
idx = np.searchsorted(z_sorted, z_grid, side='right')
A_vals = (N - idx) / N

# ==============================
# COMPUTE ALPHA (CORRECTED)
# ==============================

print("Computing alpha...")

step = int(DELTA_Z)

alpha = []
z_mid = []
A_mid = []

for i in range(len(z_grid) - step):
    A1 = A_vals[i]
    A2 = A_vals[i + step]

    # avoid instability
    if A1 > A_MIN and A1 < A_MAX and A2 > A_MIN:
        alpha.append(A1 / A2)  # <-- CORRECTED (inverted)
        z_mid.append(z_grid[i])
        A_mid.append(A1)

alpha = np.array(alpha)
z_mid = np.array(z_mid)
A_mid = np.array(A_mid)

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

df = pd.DataFrame({
    "elevation_m": z_mid,
    "area_fraction": A_mid,
    "alpha": alpha
})

df.to_csv(f"{OUT_DIR}/alpha_hypsometry_final.csv", index=False)

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

# Hypsometric curve
plt.figure(figsize=(8,5))
plt.plot(z_grid, A_vals)
plt.xlabel("Elevation (m)")
plt.ylabel("Cumulative Area Fraction A(z)")
plt.title("Hypsometric Curve (Dogger Bank)")
plt.gca().invert_xaxis()
plt.savefig(f"{OUT_DIR}/hypsometric_curve_final.png", dpi=200)

# Alpha vs elevation
plt.figure(figsize=(8,5))
plt.plot(z_mid, alpha)
plt.xlabel("Elevation (m)")
plt.ylabel("Alpha (A(z)/A(z+Δz))")
plt.title(f"Alpha vs Elevation (Δz = {DELTA_Z} m)")
plt.gca().invert_xaxis()
plt.savefig(f"{OUT_DIR}/alpha_vs_elevation_final.png", dpi=200)

# Alpha vs area (important)
plt.figure(figsize=(8,5))
plt.scatter(A_mid, alpha, s=5)
plt.xlabel("Area Fraction A")
plt.ylabel("Alpha")
plt.title("Alpha as Function of Area")
plt.savefig(f"{OUT_DIR}/alpha_vs_area.png", dpi=200)

# Histogram
plt.figure(figsize=(8,5))
plt.hist(alpha, bins=50)
plt.xlabel("Alpha")
plt.ylabel("Frequency")
plt.title("Alpha Distribution (Hypsometric)")
plt.savefig(f"{OUT_DIR}/alpha_hist_final.png", dpi=200)

# ==============================
# SUMMARY
# ==============================

print("\n=== FINAL HYPSOMETRIC ALPHA ===")
print(f"Mean alpha: {np.mean(alpha):.3f}")
print(f"Median alpha: {np.median(alpha):.3f}")
print(f"Std alpha: {np.std(alpha):.3f}")

print("\nSaved outputs:")
print("- alpha_hypsometry_final.csv")
print("- hypsometric_curve_final.png")
print("- alpha_vs_elevation_final.png")
print("- alpha_vs_area.png")
print("- alpha_hist_final.png")

from scipy.optimize import curve_fit

def model(A, k, p):
    return 1 - k * (1 - A)**p

params, _ = curve_fit(model, A_mid, alpha)

print("k, p =", params)
