import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

# =========================
# INPUT FILES
# =========================
SCALING_FILE = "p_lambda_scaling_derivative_free.csv"
RAW_FILE = "p_lambda_raw_derivative_free.csv"

P_MAX = 3.0   # same bound used in fitting

# =========================
# LOAD DATA
# =========================
scaling = pd.read_csv(SCALING_FILE)
raw = pd.read_csv(RAW_FILE)

# =========================
# DETECT SATURATION
# =========================
scaling["saturated"] = scaling["p_mean"] > (P_MAX - 1e-3)

print("\n=== SATURATION CHECK ===")
print(scaling[["lambda", "p_mean", "saturated"]])

# =========================
# MODEL (only if valid)
# =========================
def scaling_model(lam, p_inf, p0):
    return p_inf - (p_inf - p0) * np.exp(-lam)

valid = scaling[~scaling["saturated"]]

fit_success = False
if len(valid) >= 3:
    try:
        popt, _ = curve_fit(
            scaling_model,
            valid["lambda"],
            valid["p_mean"],
            p0=[1.2, 0.5]
        )
        p_inf, p0 = popt
        fit_success = True
    except:
        fit_success = False

# =========================
# PLOT 1: SCALING CURVE
# =========================
plt.figure()

plt.errorbar(
    scaling["lambda"],
    scaling["p_mean"],
    yerr=scaling["p_std"],
    marker='o',
    linestyle='',
    label="Observed"
)

# Highlight saturated points
sat = scaling["saturated"]
plt.scatter(
    scaling["lambda"][sat],
    scaling["p_mean"][sat],
    marker='x',
    s=100,
    label="Saturated (invalid)"
)

# Plot fit if valid
if fit_success:
    lam_fit = np.linspace(0, scaling["lambda"].max(), 200)
    p_fit = scaling_model(lam_fit, p_inf, p0)
    plt.plot(lam_fit, p_fit, linestyle='--', label="Fit")

    print("\n=== FIT RESULTS ===")
    print(f"p_inf = {p_inf:.3f}")
    print(f"p0    = {p0:.3f}")
else:
    print("\nNo valid fit (all values saturated).")

plt.xlabel("λ (normalized window size)")
plt.ylabel("p (hypsometric exponent)")
plt.title("Scaling of p with Spatial Scale λ")
plt.legend()
plt.tight_layout()
plt.savefig("p_lambda_scaling_plot.png", dpi=150)

# =========================
# PLOT 2: RAW DISTRIBUTION
# =========================
plt.figure()

plt.scatter(
    raw["lambda"],
    raw["p"],
    alpha=0.3,
    label="Raw windows"
)

plt.axhline(P_MAX, linestyle='--', label="Upper bound (p=3)")

plt.xlabel("λ")
plt.ylabel("p")
plt.title("Raw p Distribution Across Windows")
plt.legend()
plt.tight_layout()
plt.savefig("p_lambda_raw_scatter.png", dpi=150)

# =========================
# PLOT 3: HISTOGRAM
# =========================
plt.figure()

plt.hist(raw["p"], bins=40)

plt.axvline(P_MAX, linestyle='--', label="Upper bound")

plt.xlabel("p")
plt.ylabel("Frequency")
plt.title("Distribution of Fitted p Values")
plt.legend()
plt.tight_layout()
plt.savefig("p_lambda_hist.png", dpi=150)

plt.show()

print("\nSaved:")
print("- p_lambda_scaling_plot.png")
print("- p_lambda_raw_scatter.png")
print("- p_lambda_hist.png")