"""Loading, the split, cascade evaluation, the oracle, the threshold search and the whole plan."""
import os

import numpy as np
import pytest

from router import bounds
from router.data import SAMPLE_LONG, SAMPLE_WIDE, load_csv, load_text
from router.plan import plan
from router.policies import (SKIP, enumerate_cascades, evaluate_cascade, evaluate_single, mulberry32, oracle,
                             pick_best, quantile_grid, split_indices)

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEMO = os.path.join(ROOT, "web", "demo.csv")


@pytest.fixture(scope="module")
def demo():
    return load_csv(DEMO)


@pytest.fixture(scope="module")
def demo_plan(demo):
    return plan(demo, target=0.9, delta=0.05)


# ------------------------------------------------------------------ loading --

def test_wide_and_long_samples_load_the_same():
    w, l = load_text(SAMPLE_WIDE), load_text(SAMPLE_LONG)
    assert w.models == ["tiny", "frontier"] and l.models == ["tiny", "frontier"]
    assert w.n == 3 and l.n == 2
    assert np.allclose(w.correct[:2], l.correct) and np.allclose(w.cost[:2], l.cost)
    assert w.confidence is not None and w.latency is not None and w.prompts[0].startswith("What is 17")


def test_models_are_ordered_by_mean_cost_and_aliases_work():
    text = ("id,question,big_score,big_price,small_score,small_price\n"
            "a,hello there,1,0.5,0,0.01\n"
            "b,second one,true,0.4,yes,0.02\n"
            "c,third,0,0.6,1,\n")                     # row c has a missing cost: dropped
    log = load_text(text)
    assert log.models == ["small", "big"]
    assert log.n == 2 and log.dropped == 1
    assert log.binary and log.confidence is None and log.latency is None
    assert log.correct[1].tolist() == [1.0, 1.0]


def test_long_drops_incomplete_prompts():
    text = "prompt_id,model,correct,cost\nq1,a,1,1\nq1,b,0,2\nq2,a,1,1\n"
    log = load_text(text)
    assert log.n == 1 and log.dropped == 1


def test_bad_files_raise():
    with pytest.raises(ValueError):
        load_text("x,y\n1,2\n")
    with pytest.raises(ValueError):
        load_text("prompt_id,a_correct,a_cost,b_correct,b_cost\nq,2,1,1,1\n")


# -------------------------------------------------------------------- split --

def test_mulberry32_matches_the_javascript_stream():
    # first three draws of mulberry32(1) in web/router.js
    r = mulberry32(1)
    assert [round(r(), 12) for _ in range(3)] == [0.627073940588, 0.002735721180, 0.527447039960]


def test_split_is_deterministic_and_disjoint():
    a, b = split_indices(101, 0.5, 7)
    a2, b2 = split_indices(101, 0.5, 7)
    assert a.tolist() == a2.tolist() and b.tolist() == b2.tolist()
    assert len(a) == 51 and len(b) == 50
    assert sorted(a.tolist() + b.tolist()) == list(range(101))
    a3, _ = split_indices(101, 0.5, 8)
    assert a3.tolist() != a.tolist()


# --------------------------------------------------------------- evaluation --

def test_cascade_matches_a_row_by_row_simulation(demo):
    rng = np.random.default_rng(3)
    idx = np.arange(demo.n)
    conf = demo.confidence
    for _ in range(20):
        thr = [SKIP if rng.random() < 0.25 else float(rng.random()) for _ in range(demo.m)]
        if all(t is SKIP for t in thr):
            thr[-1] = 0.0
        res = evaluate_cascade(demo, idx, thr, conf)
        active = [i for i, t in enumerate(thr) if t is not SKIP]
        q = c = 0.0
        for r in range(demo.n):
            for k, i in enumerate(active):
                c += demo.cost[r, i]
                if k == len(active) - 1 or conf[r, i] >= thr[i]:
                    q += demo.correct[r, i]
                    break
        assert abs(res["quality"] - q / demo.n) < 1e-12
        assert abs(res["cost"] - c / demo.n) < 1e-12
        assert abs(sum(res["shares"]) - 1) < 1e-9


def test_single_model_cascade_equals_the_single_model(demo):
    idx = np.arange(demo.n)
    for i in range(demo.m):
        thr = [SKIP] * demo.m
        thr[i] = 0.0
        a, b = evaluate_cascade(demo, idx, thr, demo.confidence), evaluate_single(demo, idx, i)
        assert abs(a["quality"] - b["quality"]) < 1e-12 and abs(a["cost"] - b["cost"]) < 1e-12


# ------------------------------------------------------------------- oracle --

def test_oracle_is_a_lower_bound_on_cost(demo):
    idx = np.arange(demo.n)
    orc = oracle(demo, idx)
    curve = orc["curve"]
    assert all(curve[k]["cost"] <= curve[k + 1]["cost"] + 1e-15 for k in range(len(curve) - 1))
    assert abs(curve[-1]["quality"] - orc["quality"]) < 1e-12
    rng = np.random.default_rng(5)
    for _ in range(30):
        thr = [float(rng.random()) for _ in range(demo.m)]
        res = evaluate_cascade(demo, idx, thr, demo.confidence)
        # cheapest oracle point with at least this quality costs no more than the cascade
        oc = next(p["cost"] for p in curve if p["quality"] >= res["quality"] - 1e-12)
        assert oc <= res["cost"] + 1e-12
        assert res["quality"] <= orc["quality"] + 1e-12


# ------------------------------------------------------------------ search --

def test_enumeration_is_exhaustive_and_pick_best_is_optimal(demo):
    cal, _ = split_indices(demo.n, 0.3, 2)
    conf = demo.confidence
    grids = [quantile_grid(conf[cal, i], 8) for i in range(demo.m)]
    thrs, costs, quals = enumerate_cascades(demo, cal, conf, grids)
    scale = float(demo.cost.mean(axis=0).max())
    # every enumerated plan is evaluated correctly
    for j in range(0, len(thrs), 97):
        res = evaluate_cascade(demo, cal, thrs[j], conf)
        assert abs(res["cost"] - costs[j]) < 1e-12 and abs(res["quality"] - quals[j]) < 1e-12
    # brute force: random grid cascades never beat the pick
    rng = np.random.default_rng(9)
    for target in (0.7, 0.8, 0.88):
        j, feasible = pick_best(costs, quals, target, scale)
        assert feasible and quals[j] >= target
        for _ in range(300):
            thr = [SKIP if rng.random() < 0.3 else float(rng.choice(grids[i])) for i in range(demo.m)]
            if all(t is SKIP for t in thr):
                continue
            res = evaluate_cascade(demo, cal, thr, conf)
            if res["quality"] >= target:
                assert res["cost"] >= costs[j] - 1e-12


def test_cost_is_monotone_in_the_target(demo_plan):
    """Paying for more quality never gets cheaper: calibration cost is nondecreasing along the sweep,
    and so is the certified (held-out) cost whenever consecutive targets are both certified."""
    rows = demo_plan["sensitivity"]
    log = load_csv(DEMO)
    cal, _ = split_indices(log.n, 0.5, 1)
    cal_costs = [evaluate_cascade(log, cal, r["thresholds"], log.confidence)["cost"] for r in rows]
    assert all(cal_costs[k] <= cal_costs[k + 1] + 1e-12 for k in range(len(rows) - 1)), cal_costs
    aims = [r["aim"] for r in rows if r["aim"] is not None]
    assert aims == sorted(aims)


# -------------------------------------------------------------------- plan --

def test_demo_plan_headline(demo_plan):
    rep = demo_plan
    cas, frontier = rep["cascade"]["hold"], rep["singles"][-1]["hold"]
    assert rep["signal"] == "confidence" and rep["nCal"] == 750 and rep["nHold"] == 750
    assert cas["certified"] and cas["qualityLB"] >= 0.9
    assert cas["cost"] < 0.5 * frontier["cost"]
    assert rep["oracle"]["hold"]["cost"] < cas["cost"]
    assert rep["router"] is not None
    # the guarantee is exactly Clopper-Pearson on the held-out successes
    k = round(cas["qualitySum"])
    assert abs(cas["qualityLB"] - bounds.clopper_pearson_lower(k, 750, 0.05)) < 1e-12
    # traffic shares add up and every model that answers was called
    assert abs(sum(cas["shares"]) - 1) < 1e-9
    assert all(c >= s - 1e-12 for c, s in zip(cas["calls"], cas["shares"]))


def test_plan_without_confidence_learns_a_router(demo):
    log = load_csv(DEMO)
    log.confidence = None
    rep = plan(log, target=0.85, delta=0.1)
    assert rep["signal"] == "learned" and rep["cascade"] is not None and rep["router"] is not None


def test_plan_with_nothing_to_route_on(demo):
    log = load_csv(DEMO)
    log.confidence = None
    log.prompts = None
    rep = plan(log, target=0.85)
    assert rep["signal"] == "none" and rep["cascade"] is None and rep["router"] is None
    assert len(rep["singles"]) == 4 and rep["oracle"]["hold"]["quality"] > 0.95


def test_unreachable_target_is_reported_plainly(demo):
    rep = plan(demo, target=0.998, delta=0.05)
    assert rep["aim"] is None                      # 750 rows cannot certify 99.8%
    assert not rep["cascade"]["feasible"] and not rep["cascade"]["hold"]["certified"]