"""Example (c): an 8-step finance-QA agent; Monte Carlo Shapley with CIs and adaptive stopping.

intake -> router -> planner -> search -> extractor -> calculator -> drafter -> guardrail

Question: "What was the % change in travel reimbursment (sic) spend in FY2025?"
Ground truth: +12.5 %. What has to go right (each fix only shifts odds where
the step is an LLM call; noise is shared across counterfactuals):

* routing:   intake fix (spelling normalisation) OR router fix (LLM classifier)
* evidence:  routed to finance AND (planner fix adds FY2025 OR search fix filters by metadata)
* number:    evidence AND (extractor fix handles units OR calculator fix re-derives units)
* delivery:  the draft must contain the number (drafter LLM; fix = stronger model that also
             cites) AND pass the guardrail (as-run blocks uncited drafts; fix allows them)

With 8 players there are 256 coalitions, so exact Shapley is still available
as ground truth for the Monte Carlo estimators.

Run: ``python -m agentblame.demos.eight_step [--out DIR]``
"""

from __future__ import annotations

import argparse
import os
import time

import numpy as np

from agentblame import (BlameGame, Case, NoiseSource, Pipeline, blame_batch, exact_shapley,
                        format_batch, format_shapley, mc_shapley, step, write_html)

TABLES = {"FY2023": {"travel": (480, 520)}, "FY2025": {"travel": (800, 900)}}   # $k, prev/cur
TRUE_PCT = 12.5


@step("intake", cost=0.5)
def intake(question):
    return question


@intake.variant("spellfix", cost=1.0)
def intake_fixed(question):
    return question.replace("reimbursment", "reimbursement")


@step("router", cost=0.5)
def router(intake):
    return "finance" if "reimbursement" in intake else "general"


@router.variant("llm_classifier", cost=2.0)
def router_fixed(intake, noise):
    return noise.categorical({"finance": 2.5, "general": 0.0})


@step("planner", cost=2.0)
def planner(intake):
    return {"metric": "travel", "fiscal_year": None}


@planner.variant("explicit_year", cost=3.0)
def planner_fixed(intake):
    return {"metric": "travel", "fiscal_year": "FY2025"}


@step("search", cost=5.0)
def search(router, planner):
    if router != "finance":
        return None                                    # searched the wrong index
    return planner["fiscal_year"] or "FY2023"          # keyword match prefers the older table


@search.variant("metadata_filter", cost=6.0)
def search_fixed(router, planner):
    return "FY2025" if router == "finance" else None


@step("extractor", cost=2.0)
def extractor(search, noise):
    if search is None:
        return None
    prev, cur = TABLES[search]["travel"]
    units_ok = noise.categorical({"ok": 0.0, "unit_error": 1.0}) == "ok"
    return {"prev": prev, "cur": cur if units_ok else cur * 1000, "fy": search}


@extractor.variant("unit_aware", cost=3.0)
def extractor_fixed(search, noise):
    if search is None:
        return None
    prev, cur = TABLES[search]["travel"]
    units_ok = noise.categorical({"ok": 3.0, "unit_error": 0.0}) == "ok"
    return {"prev": prev, "cur": cur if units_ok else cur * 1000, "fy": search}


@step("calculator", cost=1.0)
def calculator(extractor):
    if extractor is None:
        return None
    return round(100 * (extractor["cur"] - extractor["prev"]) / extractor["prev"], 1)


@calculator.variant("rederive_units", cost=1.5)
def calculator_fixed(extractor):
    if extractor is None:
        return None
    prev, cur = TABLES[extractor["fy"]]["travel"]
    return round(100 * (cur - prev) / prev, 1)


@step("drafter", cost=3.0)
def drafter(calculator, noise):
    if calculator is None:
        return {"text": "I could not find that.", "cited": False}
    inc = noise.categorical({"include": 0.0, "omit": 0.4}) == "include"
    return {"text": f"Travel spend changed by {calculator}%." if inc else "Travel spend went up.",
            "cited": False}


@drafter.variant("strong_model", cost=8.0)
def drafter_fixed(calculator, noise):
    if calculator is None:
        return {"text": "I could not find that.", "cited": True}
    inc = noise.categorical({"include": 3.0, "omit": 0.4}) == "include"
    return {"text": f"Travel spend changed by {calculator}% [finance/{'FY'}]." if inc
            else "Travel spend went up [finance].", "cited": True}


@step("guardrail", cost=1.0)
def guardrail(drafter):
    return drafter["text"] if drafter["cited"] else "BLOCKED: uncited claim"


@guardrail.variant("calibrated", cost=1.0)
def guardrail_fixed(drafter):
    return drafter["text"]


def metric(out) -> float:
    return 1.0 if f"{TRUE_PCT}%" in out["guardrail"] else 0.0


def build_pipeline() -> Pipeline:
    return Pipeline([intake, router, planner, search, extractor, calculator, drafter, guardrail],
                    metric, name="eight_step_finance_qa")


Q = "What was the % change in travel reimbursment spend in FY2025?"


def make_cases(n: int = 60) -> list[Case]:
    return [Case(f"c{i:03d}", {"question": Q}, seed=i) for i in range(n)]


def fresh_game(case: Case) -> tuple[Pipeline, BlameGame]:
    """New pipeline (empty cache, zeroed meter) so each estimator's cost is measured alone."""
    p = build_pipeline()
    return p, BlameGame(p, case.inputs, world=NoiseSource(case.seed))


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--case-seed", type=int, default=0)
    ap.add_argument("--cases", type=int, default=400)
    ap.add_argument("--stream", type=int, default=3000,
                    help="max cases for the adaptive-stopping batch")
    ap.add_argument("--out", default="agentblame_out")
    args = ap.parse_args(argv)
    t0 = time.time()
    case = Case(f"c{args.case_seed:03d}", {"question": Q}, seed=args.case_seed)

    _, g = fresh_game(case)
    ex = exact_shapley(g)
    naive_cost = sum(s.variants[g.fixes.get(s.name, "as_run")][1] for s in g.pipeline.steps) * 256
    print(f"Case {case.id}: exact Shapley over 2^8 = 256 coalitions "
          f"(step cost with caching {ex.cost:g}; without caching ~{naive_cost:g})")
    print(format_shapley(ex), "\n")
    truth = np.array(list(ex.values.values()))

    lines = [f"{'estimator':<34}{'perms':>6}{'coalitions':>11}{'cost':>8}{'max|err|':>10}"
             f"{'max CI radius':>15}  stop"]
    rows = []
    for k in (16, 64, 256, 1024):
        _, gk = fresh_game(case)
        r = mc_shapley(gk, k, bound="bernstein", seed=1)
        rows.append((f"MC, fixed {k} perms (Bernstein)", r))
    for bound in ("hoeffding", "bernstein"):
        _, gk = fresh_game(case)
        rows.append((f"MC adaptive top1 ({bound})",
                     mc_shapley(gk, 5000, bound=bound, stop="top1", seed=1)))
    _, gk = fresh_game(case)
    rows.append(("MC adaptive top1, budget=100", mc_shapley(gk, 5000, stop="top1", budget=100, seed=1)))
    for name, r in rows:
        err = np.max(np.abs(np.array(list(r.values.values())) - truth))
        rad = max((hi - lo) / 2 for lo, hi in r.ci.values())
        lines.append(f"{name:<34}{r.n_permutations:>6}{r.coalitions_evaluated:>11}{r.cost:>8g}"
                     f"{err:>10.3f}{rad:>15.3f}  {r.stop_reason}")
    table = "\n".join(lines)
    print("Monte Carlo permutation Shapley vs exact (delta = 0.05, simultaneous over steps)")
    print(table, "\n")
    adaptive = rows[5][1]
    print("Adaptive (Bernstein, top1) estimate with CIs:")
    print(format_shapley(adaptive), "\n")

    pipe = build_pipeline()
    cases = make_cases(args.cases)
    full = blame_batch(pipe, cases, method="exact")
    print(f"Batch over {args.cases} cases, exact Shapley per failure:")
    print(format_batch(full), "\n")
    stream = make_cases(args.stream)
    early = {}
    for bound in ("bernstein", "normal"):
        p2 = build_pipeline()
        early[bound] = blame_batch(p2, stream, method="exact", stop_when_separated=True,
                                   min_failures=8, stop_bound=bound)
    print(f"Batch, budget-aware: analyse failures from a stream of up to {args.stream} cases and stop "
          f"once the top step's interval separates from all others (delta=0.05):")
    for bound, b in early.items():
        top = b.ranking()[0]
        print(f"  {bound:<10} {len(b.cases):>5} failures analysed, cost {b.cost:>8g}, top = {top.name} (mean, 95% normal CI) "
              f"{top.shapley:+.3f} [{top.shapley_ci[0]:+.3f}, {top.shapley_ci[1]:+.3f}]  ({b.stop_reason})")
    print("  (bernstein: finite-sample, valid under optional stopping; "
          "normal: asymptotic heuristic, no such guarantee)")
    early = early["bernstein"]

    os.makedirs(args.out, exist_ok=True)
    path = write_html(
        os.path.join(args.out, "eight_step.html"), "8-step finance agent: Monte Carlo blame",
        [(f"Single case {case.id}: exact Shapley", format_shapley(ex)),
         ("Monte Carlo vs exact", table),
         ("Adaptive MC estimate (empirical Bernstein, stop when top step separates)",
          format_shapley(adaptive)),
         (f"Batch of {args.cases} cases (exact)", full),
         ("Batch with adaptive stopping (empirical Bernstein)", early)],
        intro="intake → router → planner → search → extractor → calculator → drafter → guardrail. "
              "LLM-like steps sample via Gumbel-max with noise shared across counterfactuals.")
    print(f"\nHTML report: {path}\nruntime {time.time() - t0:.1f}s")
    return ex, rows, full, early


if __name__ == "__main__":
    main()