"""Put it together: split, tune on the calibration half, certify on the held-out half."""
from __future__ import annotations

import math
from typing import Optional

import numpy as np

from . import bounds
from .data import EvalLog
from .policies import (SKIP, enumerate_cascades, evaluate_cascade, evaluate_router, evaluate_single,
                       grid_size, learn_confidence, oracle, pareto, pick_best, quantile_grid, refine,
                       split_indices)

G_MAX = 24
G_FINE = 40


def certify(res: dict, n: int, delta: float, binary: bool, target: float) -> dict:
    """Attach the held-out guarantees to an evaluation."""
    lb = bounds.quality_lower_bound(res["qualitySum"], n, delta, binary)
    out = dict(res)
    out["qualityLB"] = lb
    out["costUB"] = res["cost"] + bounds.hoeffding_radius(n, delta, res["maxCost"])
    out["certified"] = lb >= target - 1e-12
    out["rowsNeeded"] = None if out["certified"] else bounds.rows_to_certify(res["quality"], target, delta, binary)
    return out


def sweep_targets(lo: float, hi: float, user: float, step: float = 0.01):
    """lo, lo+step, ... up to hi, plus the user's own target."""
    k0 = int(math.floor(lo / step + 1e-9))
    ts = []
    k = k0
    while (k0 + (k - k0)) * step <= hi + 1e-9:
        ts.append(round(k * step, 6))
        k += 1
    u = round(user, 6)
    if u not in ts:
        ts.append(u)
    return sorted(set(ts))


def aim_for(target: float, n_hold: int, n_cal: int, delta: float, binary: bool, safety: float):
    """The calibration quality to tune for: the rate that certifies the target on n_hold rows, plus
    `safety` standard errors against the winner's curse of picking the best of thousands."""
    a = bounds.required_rate(target, n_hold, delta, binary)
    if a is None:
        return None
    a = a + safety * math.sqrt(a * (1.0 - a) / n_cal)
    return a if a <= 1.0 else None


def plan(log: EvalLog, target: float = 0.9, delta: float = 0.05, cal_frac: float = 0.5, seed: int = 1,
         signal: str = "auto", step: float = 0.01, safety: float = 1.0) -> dict:
    n, m = log.n, log.m
    if n < 4:
        raise ValueError("Need at least 4 prompts.")
    cal, hold = split_indices(n, cal_frac, seed)
    binary = log.binary
    scale = float(log.cost.mean(axis=0).max())

    # --- which signal gates the cascade
    pred = learn_confidence(log, cal) if log.prompts is not None else None
    if signal == "auto":
        signal = "confidence" if log.confidence is not None else ("learned" if pred is not None else "none")
    if signal == "confidence" and log.confidence is None:
        signal = "learned" if pred is not None else "none"
    if signal == "learned" and pred is None:
        signal = "confidence" if log.confidence is not None else "none"
    conf = log.confidence if signal == "confidence" else pred

    report = {
        "aim": aim_for(target, len(hold), len(cal), delta, binary, safety), "safety": safety,
        "requiredRate": bounds.required_rate(target, len(hold), delta, binary),
        "n": n, "nCal": int(len(cal)), "nHold": int(len(hold)), "models": list(log.models),
        "target": target, "delta": delta, "calFrac": cal_frac, "seed": seed, "signal": signal,
        "binary": binary, "hasLatency": log.latency is not None, "hasConfidence": log.confidence is not None,
        "hasPrompts": log.prompts is not None, "warnings": list(log.warnings), "dropped": log.dropped,
        "meanCost": [float(x) for x in log.cost.mean(axis=0)],
    }

    # --- single models
    singles = []
    for i in range(m):
        c = evaluate_single(log, cal, i)
        h = certify(evaluate_single(log, hold, i), len(hold), delta, binary, target)
        singles.append({"kind": "model", "name": log.models[i], "index": i, "cal": c, "hold": h})
    report["singles"] = singles
    report["frontier"] = pareto([s["hold"] for s in singles])

    # --- oracle
    orc_hold = certify(oracle(log, hold), len(hold), delta, binary, target)
    report["oracle"] = {"kind": "oracle", "name": "Oracle", "cal": oracle(log, cal), "hold": orc_hold}

    # --- cascade
    report["cascade"] = None
    report["router"] = None
    report["sensitivity"] = []
    lo_q = min(s["cal"]["quality"] for s in singles)
    if conf is not None:
        g = grid_size(m, g_max=G_MAX)
        grids = [quantile_grid(conf[cal, i], g) for i in range(m)]
        fine = [quantile_grid(conf[cal, i], G_FINE) for i in range(m)]
        thrs, costs, quals = enumerate_cascades(log, cal, conf, grids)
        report["enumerated"] = int(len(thrs))
        report["gridSize"] = g
        hi_q = float(quals.max())
        targets = sweep_targets(lo_q, max(hi_q, target), target, step)
        plans = {}
        for t in targets:
            aim = aim_for(t, len(hold), len(cal), delta, binary, safety)
            a = aim if aim is not None else 1.0
            j, feasible = pick_best(costs, quals, a, scale)
            thr, res = list(thrs[j]), None
            if g < G_MAX:
                thr, res = refine(log, cal, conf, thr, a, fine, scale)
            if res is None:
                res = evaluate_cascade(log, cal, thr, conf)
            plans[t] = {"thresholds": res["thresholds"], "cal": res, "aim": aim,
                        "feasible": aim is not None and res["quality"] >= a - 1e-12}
        # a plan for a higher target also meets any lower one: never pay more for less
        prev = None
        for t in sorted(targets, reverse=True):
            p = plans[t]
            if prev is not None and prev["feasible"] and prev["cal"]["cost"] < p["cal"]["cost"] - 1e-15:
                plans[t] = {"thresholds": prev["thresholds"], "cal": prev["cal"], "aim": p["aim"], "feasible": True}
            prev = plans[t]
        for t in targets:
            p = plans[t]
            p["hold"] = certify(evaluate_cascade(log, hold, p["thresholds"], conf), len(hold), delta, binary, t)
        cp = plans[round(target, 6)]
        report["cascade"] = {"kind": "cascade", "name": "Cascade", "thresholds": cp["thresholds"], "aim": cp["aim"],
                             "feasible": cp["feasible"], "cal": cp["cal"], "hold": cp["hold"],
                             "bestCalQuality": hi_q}
        report["sensitivity"] = [{"target": t, "aim": plans[t]["aim"], "cascade": plans[t]["hold"],
                                  "cascadeFeasible": plans[t]["feasible"], "thresholds": plans[t]["thresholds"]} for t in targets]

    # --- router (predict first, then call one model)
    if pred is not None:
        taus = quantile_grid(pred[cal].reshape(-1), 64)
        evals = [evaluate_router(log, cal, tau, pred) for tau in taus]
        rc = np.array([e["cost"] for e in evals])
        rq = np.array([e["quality"] for e in evals])
        targets = [s["target"] for s in report["sensitivity"]] or sweep_targets(lo_q, max(float(rq.max()), target), target, step)
        rplans = {}
        for t in targets:
            aim = aim_for(t, len(hold), len(cal), delta, binary, safety)
            j, feasible = pick_best(rc, rq, aim if aim is not None else 1.0, scale)
            rplans[t] = {"tau": taus[j], "cal": evals[j], "aim": aim, "feasible": aim is not None and feasible}
        prev = None
        for t in sorted(targets, reverse=True):
            p = rplans[t]
            if prev is not None and prev["feasible"] and prev["cal"]["cost"] < p["cal"]["cost"] - 1e-15:
                rplans[t] = {"tau": prev["tau"], "cal": prev["cal"], "aim": p["aim"], "feasible": True}
            prev = rplans[t]
        for t in targets:
            p = rplans[t]
            p["hold"] = certify(evaluate_router(log, hold, p["tau"], pred), len(hold), delta, binary, t)
        rp = rplans[round(target, 6)]
        report["router"] = {"kind": "router", "name": "Learned router", "tau": rp["tau"], "aim": rp["aim"],
                            "feasible": rp["feasible"], "cal": rp["cal"], "hold": rp["hold"], "bestCalQuality": float(rq.max())}
        if report["sensitivity"]:
            for row in report["sensitivity"]:
                row["router"] = rplans[row["target"]]["hold"]
        else:
            report["sensitivity"] = [{"target": t, "aim": rplans[t]["aim"], "router": rplans[t]["hold"]} for t in targets]

    # --- per-target comparison: cheapest certified single model and the oracle's price
    curve = orc_hold["curve"]
    for row in report["sensitivity"]:
        t = row["target"]
        cert = [s for s in singles if s["hold"]["qualityLB"] >= t - 1e-12]
        best = min(cert, key=lambda s: s["hold"]["cost"]) if cert else None
        row["bestSingle"] = {"name": best["name"], "cost": best["hold"]["cost"]} if best else None
        oc = next((p["cost"] for p in curve if p["quality"] >= t - 1e-12), None)
        row["oracleCost"] = oc
    report["breakEven"] = break_even(report)
    return report


def break_even(report: dict) -> Optional[dict]:
    """Where the cascade stops being the cheapest certified option."""
    rows = [r for r in report["sensitivity"] if r.get("cascade")]
    if not rows:
        return None
    wins = [r["target"] for r in rows if r["cascade"]["certified"] and
            (r["bestSingle"] is None or r["cascade"]["cost"] < r["bestSingle"]["cost"])]
    certified = [r["target"] for r in rows if r["cascade"]["certified"]]
    lo, hi = (min(wins), max(wins)) if wins else (None, None)
    in_range = sum(1 for r in rows if wins and lo - 1e-9 <= r["target"] <= hi + 1e-9)
    return {"winsFrom": lo, "winsTo": hi, "wins": len(wins), "inRange": in_range,
            "maxCertified": max(certified) if certified else None}


def run(path_or_log, **kw) -> dict:
    from .data import load_csv
    log = path_or_log if isinstance(path_or_log, EvalLog) else load_csv(path_or_log)
    return plan(log, **kw)