#!/usr/bin/env python3
"""
Bootstrap + Sensitivity Analysis for p(λ)

- Computes bootstrap confidence intervals for scaling exponent p
- Evaluates stability across λ (window scale)
- Includes optional spatial block bootstrap

Input:
    outputs/p_lambda_raw_derivative_free.csv

Output:
    outputs/p_bootstrap_summary.csv
    outputs/p_bootstrap_distributions.png
"""

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

from collections import defaultdict

# -----------------------------
# CONFIG
# -----------------------------
INPUT_FILE = "outputs/p_lambda_raw_derivative_free.csv"
N_BOOT = 5000
BLOCK_BOOTSTRAP = True      # toggle spatial block bootstrap
BLOCK_SIZE = 2.0           # degrees (lat/lon binning)

np.random.seed(42)

# -----------------------------
# LOAD DATA
# -----------------------------
df = pd.read_csv(INPUT_FILE)

# Filter pathological values (optional but recommended)
df = df[np.isfinite(df["p"])]
df = df[(df["p"] > 0) & (df["p"] < 5)]

# -----------------------------
# SPATIAL BLOCKING (optional)
# -----------------------------
if BLOCK_BOOTSTRAP:
    df["lat_bin"] = (df["lat"] // BLOCK_SIZE) * BLOCK_SIZE
    df["lon_bin"] = (df["lon"] // BLOCK_SIZE) * BLOCK_SIZE
    df["block_id"] = df["lat_bin"].astype(str) + "_" + df["lon_bin"].astype(str)

    blocks = list(df.groupby("block_id"))
else:
    blocks = None

# -----------------------------
# BOOTSTRAP FUNCTION
# -----------------------------
def bootstrap_mean(values, n_boot=1000):
    """Standard bootstrap of mean."""
    n = len(values)
    samples = np.random.choice(values, size=(n_boot, n), replace=True)
    means = samples.mean(axis=1)
    return means

def block_bootstrap_mean(df, n_boot=1000):
    """Spatial block bootstrap."""
    block_ids = df["block_id"].unique()
    means = []

    for _ in range(n_boot):
        sampled_blocks = np.random.choice(block_ids, size=len(block_ids), replace=True)
        sampled = pd.concat([df[df["block_id"] == b] for b in sampled_blocks])
        means.append(sampled["p"].mean())

    return np.array(means)

# -----------------------------
# MAIN ANALYSIS
# -----------------------------
results = []
distributions = {}

for lam, group in df.groupby("lambda"):
    p_vals = group["p"].values

    if len(p_vals) < 5:
        continue

    if BLOCK_BOOTSTRAP:
        boot = block_bootstrap_mean(group, N_BOOT)
    else:
        boot = bootstrap_mean(p_vals, N_BOOT)

    p_mean = np.mean(p_vals)
    ci_low, ci_high = np.percentile(boot, [2.5, 97.5])
    p_std = np.std(boot)

    results.append({
        "lambda": lam,
        "p_mean": p_mean,
        "p_boot_mean": boot.mean(),
        "p_std": p_std,
        "ci_low": ci_low,
        "ci_high": ci_high,
        "n": len(p_vals)
    })

    distributions[lam] = boot

# -----------------------------
# SAVE SUMMARY
# -----------------------------
summary_df = pd.DataFrame(results).sort_values("lambda")
summary_df.to_csv("outputs/p_bootstrap_summary.csv", index=False)

print("\nBootstrap Summary:\n")
print(summary_df)

# -----------------------------
# PLOT DISTRIBUTIONS
# -----------------------------
plt.figure(figsize=(8, 5))

for lam, boot in distributions.items():
    plt.hist(boot, bins=50, alpha=0.4, label=f"λ={lam}")

plt.xlabel("Bootstrap p")
plt.ylabel("Frequency")
plt.title("Bootstrap Distributions of p(λ)")
plt.legend()

plt.tight_layout()
plt.savefig("outputs/p_bootstrap_distributions.png", dpi=300)
plt.close()

# -----------------------------
# SENSITIVITY METRIC
# -----------------------------
print("\nSensitivity Analysis:\n")

p_range = summary_df["p_boot_mean"].max() - summary_df["p_boot_mean"].min()
p_std_mean = summary_df["p_std"].mean()

print(f"Range of p across λ: {p_range:.3f}")
print(f"Mean bootstrap std: {p_std_mean:.3f}")

# Optional: stability criterion
if p_range < 0.5:
    print("Result: p is stable across scales (strong evidence)")
else:
    print("Result: p shows scale sensitivity (interpret cautiously)")
    