"""Claim 5 (Theorem 2 / Definition 4.1): a lambda-parameterised algorithm
interpolates between trusting the prediction and being robust, with
    c(ALG) <= min{ (1+lambda)(c(OPT_p)+B),  (1+1/lambda) c(OPT_p) }.

Previously: "the logbook explicitly excludes supplemental parameterized
algorithms; no implementation or testing of the lambda-parameterized algorithm."

The Clamp Policy, pinned from source (Definition 4.1):
    t~ := min{ max{ t*_phat , ceil(lambda b) } , floor(b/lambda) },   lambda in (0,1)
where t*_phat minimises the expected cost under the PREDICTED distribution.
The paper proves robustness  ECR <= 1 + 1/lambda - 1/b  for ANY prediction.

Ski rental: rent costs 1/day, buying costs b. Buying on day t against a true
demand X costs (t-1) + b if X >= t, else X.
"""
import json, numpy as np

RES = {}


def cost(t, X, b):
    return np.where(X >= t, (t-1)+b, X)


def expected_cost(t, xs, ps, b):
    return float(np.sum(ps*cost(t, xs, b)))


def opt_threshold(xs, ps, b, T):
    ts = np.arange(1, T+1)
    return int(ts[int(np.argmin([expected_cost(t, xs, ps, b) for t in ts]))])


def offline_opt(xs, ps, b):
    return float(np.sum(ps*np.minimum(xs, b)))


def make_dist(kind, T, rng, b):
    xs = np.arange(1, T+1)
    if kind == "geometric":
        q = 1.0/(0.8*b); p = q*(1-q)**(xs-1)
    elif kind == "uniform":
        p = np.ones(T)
    elif kind == "short":                 # demand almost always below b
        p = np.exp(-xs/(0.3*b))
    elif kind == "long":                  # demand almost always above b
        p = np.exp(-(xs-2.0*b)**2/(0.3*b)**2)+1e-12
    elif kind == "bimodal":
        p = np.exp(-(xs-0.2*b)**2/(0.1*b)**2)+np.exp(-(xs-2.5*b)**2/(0.2*b)**2)+1e-12
    return xs, p/p.sum()


def main():
    b, T = 100, 600
    rng = np.random.default_rng(0)
    kinds = ["geometric", "uniform", "short", "long", "bimodal"]
    dists = {k: make_dist(k, T, rng, b) for k in kinds}
    rows = []
    for lam in (0.1, 0.2, 0.3, 0.5, 0.7, 0.9):
        bound = 1+1/lam-1/b
        worst = 0.0; worst_cell = None; consist = []
        for true_k in kinds:
            xs, pt = dists[true_k]
            opt = offline_opt(xs, pt, b)
            for pred_k in kinds:                       # every prediction, incl. adversarially wrong
                _, ph = dists[pred_k]
                tstar = opt_threshold(xs, ph, b, T)    # optimal under the PREDICTION
                tclamp = int(min(max(tstar, int(np.ceil(lam*b))), int(np.floor(b/lam))))
                ecr = expected_cost(tclamp, xs, pt, b)/opt
                if ecr > worst: worst, worst_cell = ecr, (true_k, pred_k, tstar, tclamp)
                if true_k == pred_k: consist.append(ecr)
        rows.append({"lambda": lam, "robustness_bound_1+1/lam-1/b": round(bound, 4),
                     "worst_ECR_over_all_25_pairs": round(worst, 4),
                     "bound_holds": bool(worst <= bound+1e-9),
                     "worst_cell": {"true": worst_cell[0], "pred": worst_cell[1],
                                    "t_star_pred": worst_cell[2], "t_clamped": worst_cell[3]},
                     "mean_ECR_perfect_prediction": round(float(np.mean(consist)), 4),
                     "max_ECR_perfect_prediction": round(float(np.max(consist)), 4)})
        print("  lam=%.1f  bound=%.3f  worst ECR (25 pairs)=%.4f  holds=%s | perfect-prediction ECR mean=%.4f max=%.4f"
              % (lam, bound, worst, rows[-1]["bound_holds"],
                 rows[-1]["mean_ECR_perfect_prediction"], rows[-1]["max_ECR_perfect_prediction"]), flush=True)
    RES["claim5_clamp_policy"] = {
        "b": b, "horizon": T, "distributions": kinds,
        "pairs_per_lambda": len(kinds)**2, "rows": rows,
        "bound_holds_everywhere": all(r["bound_holds"] for r in rows),
        "consistency_improves_as_lambda_rises": bool(
            rows[-1]["mean_ECR_perfect_prediction"] <= rows[0]["mean_ECR_perfect_prediction"]+1e-9)}
    json.dump(RES, open("ski_results.json", "w"), indent=1)


if __name__ == "__main__":
    main(); print("DONE")