"""Synthetic experiments for the JRPM manuscript.

No proprietary airline data are used. The script reproduces the cabin-floating
decision experiment and creates summary CSVs. Forecasting-score results are
loaded from the previously generated synthetic experiment when available.
"""
import numpy as np
import pandas as pd
from scipy.optimize import minimize
from scipy.special import gammaln
from scipy.stats import poisson, nbinom

SEED = 20261001


def fit_gamma_eb_agg(Yi, Si):
    Yi = np.asarray(Yi, dtype=float)
    Si = np.asarray(Si, dtype=float)
    r = Yi / Si
    m = max(np.mean(r), 1e-6)
    v = np.var(r, ddof=1)
    noise = np.mean(np.maximum(r, 1e-6) / Si)
    latent = max(v - noise, m * m / 100)
    a0 = max(m * m / latent, 0.1)
    b0 = max(a0 / m, 0.1)

    def nll(z):
        a, b = np.exp(z)
        ll = (
            gammaln(a + Yi) - gammaln(a) - gammaln(Yi + 1)
            + a * np.log(b) + Yi * np.log(Si)
            - (a + Yi) * np.log(b + Si)
        )
        return -ll.sum()

    opt = minimize(
        nll,
        np.log([a0, b0]),
        method="L-BFGS-B",
        bounds=[(-8, 8), (-8, 12)],
    )
    return tuple(np.exp(opt.x))


def protection_from_survival(survival, price_ratio):
    """Largest y such that p_P P(D_P >= y) >= p_E for abundant economy demand."""
    return (price_ratio * survival >= 1.0).sum(axis=1).astype(int)


def run_decision_experiment(
    n_routes=600,
    n_replications=100,
    gamma_shape=1.2,
    mean_rate=0.4,
    future_exposure=8,
    premium_capacity=12,
    economy_fare=200.0,
    premium_economy_ratios=(1.1, 1.5, 2.0, 3.0, 5.0),
    training_exposure_low=2,
    training_exposure_high=14,
    seed=SEED,
):
    rng = np.random.default_rng(seed)
    rows = []
    seats = np.arange(1, premium_capacity + 1)

    for rep in range(n_replications):
        lam = rng.gamma(gamma_shape, mean_rate / gamma_shape, size=n_routes)
        Si = rng.integers(
            training_exposure_low, training_exposure_high + 1, size=n_routes
        ).astype(float)
        Yi = rng.poisson(lam * Si)
        future_demand = rng.poisson(lam * future_exposure)

        a, b = fit_gamma_eb_agg(Yi, Si)
        pooled = Yi.sum() / Si.sum()

        surv_pool = np.tile(
            poisson.sf(seats - 1, pooled * future_exposure), (n_routes, 1)
        )
        mu_mle = (Yi / Si) * future_exposure
        surv_mle = poisson.sf(seats[None, :] - 1, mu_mle[:, None])

        Aw = 0.5 + Yi
        Bw = 0.5 + Si
        pw = Bw / (Bw + future_exposure)
        surv_weak = nbinom.sf(seats[None, :] - 1, Aw[:, None], pw[:, None])

        A = a + Yi
        B = b + Si
        phepg = B / (B + future_exposure)
        surv_hepg = nbinom.sf(seats[None, :] - 1, A[:, None], phepg[:, None])

        surv_oracle = poisson.sf(
            seats[None, :] - 1, (lam * future_exposure)[:, None]
        )
        hepg_mean = (A / B) * future_exposure
        y_mean_only = np.clip(np.rint(hepg_mean).astype(int), 0, premium_capacity)

        predictive = {
            "Pooled Poisson": surv_pool,
            "Route Poisson MLE": surv_mle,
            "Weak-prior PG": surv_weak,
            "HEPG": surv_hepg,
            "Oracle": surv_oracle,
        }

        for ratio in premium_economy_ratios:
            premium_fare = economy_fare * ratio
            decisions = {
                name: protection_from_survival(surv, ratio)
                for name, surv in predictive.items()
            }
            decisions["HEPG mean-only"] = y_mean_only
            decisions["No sharing"] = np.full(n_routes, premium_capacity, dtype=int)

            for method, protected in decisions.items():
                floated = premium_capacity - protected
                relevant_revenue = (
                    premium_fare * np.minimum(future_demand, protected)
                    + economy_fare * floated
                )
                spill = np.maximum(future_demand - protected, 0)
                rows.append(
                    {
                        "replication": rep,
                        "fare_ratio": ratio,
                        "method": method,
                        "avg_relevant_revenue": relevant_revenue.mean(),
                        "avg_protected": protected.mean(),
                        "avg_floated": floated.mean(),
                        "avg_premium_spill": spill.mean(),
                        "premium_spill_probability": (spill > 0).mean(),
                    }
                )

    rep_df = pd.DataFrame(rows)
    summary = (
        rep_df.groupby(["fare_ratio", "method"], as_index=False)
        .agg(
            revenue=("avg_relevant_revenue", "mean"),
            revenue_se=("avg_relevant_revenue", "sem"),
            protected=("avg_protected", "mean"),
            floated=("avg_floated", "mean"),
            spill=("avg_premium_spill", "mean"),
            spill_probability=("premium_spill_probability", "mean"),
        )
    )
    oracle = summary.loc[summary.method == "Oracle", ["fare_ratio", "revenue"]].rename(
        columns={"revenue": "oracle_revenue"}
    )
    summary = summary.merge(oracle, on="fare_ratio", how="left")
    summary["regret"] = summary["oracle_revenue"] - summary["revenue"]
    summary["revenue_pct_oracle"] = 100 * summary["revenue"] / summary["oracle_revenue"]
    return rep_df, summary


if __name__ == "__main__":
    rep, summary = run_decision_experiment()
    rep.to_csv("decision_replications.csv", index=False)
    summary.to_csv("decision_summary.csv", index=False)
    print(summary.round(4).to_string(index=False))