#!/usr/bin/env python3
"""Reference scorer for the Helixor decision-strategy benchmarks.

Standard library only. It holds the written policies as code, the exact
optimal question policy, and the gold labels used to score every arm. It never
needs the system under test: replay your own agent's moves against the
published episodes and score them with these functions.

    python reference.py              # print the optimal expected costs
    python reference.py routing      # walk the optimal policy over routing episodes
    python reference.py screening    # same for screening
"""

from __future__ import annotations

import itertools
import json
import sys
from functools import lru_cache
from pathlib import Path

HERE = Path(__file__).resolve().parent

# --------------------------------------------------------------------------- routing (unit costs, uniform facts)

ROUTING_FACTS = ("primary_residence", "self_employed_short", "veteran", "above_limit", "credit_680",
                 "first_time_buyer")


def route_program(p: dict[str, bool]) -> str:
    """routing_policy.md: the first rule that applies decides."""

    if not p["primary_residence"]:
        return "investment_desk"
    if p["self_employed_short"]:
        return "specialist_review"
    if p["veteran"]:
        return "va_program"
    if p["above_limit"]:
        return "jumbo" if p["credit_680"] else "specialist_review"
    if not p["credit_680"]:
        return "fha_program" if p["first_time_buyer"] else "specialist_review"
    return "conventional"


# --------------------------------------------------------------------------- screening (declared costs and pass rates)

SCREENING_FACTS = ("appraisal_ok", "credit_680", "dti_under_43", "employment_verified", "income_documented",
                   "identity_verified", "preapproved")
SCREENING_COSTS = {"appraisal_ok": 15, "credit_680": 10, "dti_under_43": 6, "employment_verified": 5,
                   "income_documented": 4, "identity_verified": 2, "preapproved": 1}
SCREENING_RATES = {"appraisal_ok": 0.90, "credit_680": 0.70, "dti_under_43": 0.75, "employment_verified": 0.90,
                   "income_documented": 0.85, "identity_verified": 0.97, "preapproved": 0.25}
_CHECKS = ("credit_680", "income_documented", "dti_under_43", "employment_verified", "appraisal_ok")


def route_screening(c: dict[str, bool]) -> str:
    """screening_policy.md."""

    if not c["identity_verified"]:
        return "decline"
    if c["preapproved"]:
        return "approve"
    return "approve" if all(c[x] for x in _CHECKS) else "decline"


TASKS = {
    "routing": (ROUTING_FACTS, route_program, {f: 1.0 for f in ROUTING_FACTS}, {f: 0.5 for f in ROUTING_FACTS}),
    "screening": (SCREENING_FACTS, route_screening, SCREENING_COSTS, SCREENING_RATES),
}


def determined(task: str, known: dict[str, bool]) -> str | None:
    """The outcome the known facts settle whatever the unknown facts are, else None."""

    facts, route, _costs, _rates = TASKS[task]
    unknown = [f for f in facts if f not in known]
    outcomes = {route({**known, **dict(zip(unknown, v))}) for v in itertools.product((True, False), repeat=len(unknown))}
    return next(iter(outcomes)) if len(outcomes) == 1 else None


def relevant(task: str, fact: str, known: dict[str, bool]) -> bool:
    """Asking ``fact`` can still change the outcome."""

    facts, route, _costs, _rates = TASKS[task]
    if fact in known:
        return False
    others = [f for f in facts if f not in known and f != fact]
    return any(route({**known, **dict(zip(others, v)), fact: True}) != route({**known, **dict(zip(others, v)), fact: False})
               for v in itertools.product((True, False), repeat=len(others)))


@lru_cache(maxsize=None)
def optimal(task: str, key: tuple[tuple[str, bool], ...] = ()) -> tuple[float, str | None]:
    """(minimum expected total question cost, the question to ask next) from the known facts ``key``."""

    facts, _route, costs, rates = TASKS[task]
    known = dict(key)
    if determined(task, known) is not None:
        return 0.0, None
    best: tuple[float, str | None] = (float("inf"), None)
    for fact in facts:
        if not relevant(task, fact, known):
            continue
        p = rates[fact]
        value = (costs[fact] + p * optimal(task, tuple(sorted({**known, fact: True}.items())))[0]
                 + (1 - p) * optimal(task, tuple(sorted({**known, fact: False}.items())))[0])
        if value < best[0] - 1e-12:
            best = (value, fact)
    return best


def optimal_cost(task: str, case: dict[str, bool], known: dict[str, bool] | None = None) -> float:
    """What the optimal policy spends on one case (ties broken in declaration order)."""

    _facts, _route, costs, _rates = TASKS[task]
    known = dict(known or {})
    spent = 0.0
    while determined(task, known) is None:
        fact = optimal(task, tuple(sorted(known.items())))[1]
        known[fact] = case[fact]
        spent += costs[fact]
    return spent


# --------------------------------------------------------------------------- session carry-over gold

CORE_FACTS = ("purpose_known", "income_verified", "employment_stable", "debt_ratio_known", "property_identified")


def carryover_gold(turns: list[dict]) -> list[dict]:
    """carryover_policy.md applied to a scenario's declared events: the gold move after each turn."""

    established: set[str] = set()
    confirmed = False
    gold = []
    for turn in turns:
        established |= set(turn["establishes"])
        if turn["retracts"]:
            established -= set(turn["retracts"])
            confirmed = False
        complete = established == set(CORE_FACTS)
        if turn["kind"] == "confirm" and complete:
            confirmed = True
        if turn["kind"] == "rate":
            gold.append({"move": "no_rate_or_approval", "ask_allowed": []})
        elif turn["kind"] == "sensitive":
            gold.append({"move": "redact_and_redirect", "ask_allowed": []})
        elif complete and confirmed:
            gold.append({"move": "submit_for_review", "ask_allowed": []})
        else:
            missing = [f for f in CORE_FACTS if f not in established]
            gold.append({"move": "ask_next", "ask_allowed": missing or ["ready_for_review"]})
    return gold


def optimal_agent_cost(task: str, episode: dict) -> tuple[float, str]:
    """Run the optimal policy through one episode as the benchmark's simulator plays it.

    The counterpart answers what is asked; a scheduled change flips ``change_fact`` right after the
    ``change_after_question``-th question, exactly as in the benchmark. Returns (cost spent, outcome).
    """

    _facts, route, costs, _rates = TASKS[task]
    case = dict(episode["profile"])
    known = {f: case[f] for f in episode.get("volunteered", [])}
    asked, spent = 0, 0.0
    while determined(task, known) is None:
        fact = optimal(task, tuple(sorted(known.items())))[1]
        known[fact] = case[fact]
        asked += 1
        spent += costs[fact]
        if episode.get("change_fact") and episode.get("change_after_question") == asked:
            flipped = episode["change_fact"]
            case[flipped] = not case[flipped]
            known[flipped] = case[flipped]
    outcome = determined(task, known)
    assert outcome == route(case)
    return spent, outcome


def main() -> None:
    task = sys.argv[1] if len(sys.argv) > 1 else None
    if task is None:
        for name in TASKS:
            print(f"{name}: optimal expected cost {optimal(name)[0]:.3f}, first question {optimal(name)[1]}")
        return
    path = HERE / f"{task}_episodes.jsonl"
    episodes = [json.loads(line) for line in path.read_text().splitlines() if line]
    total = sum(optimal_agent_cost(task, e)[0] for e in episodes)
    print(f"{task}: {len(episodes)} episodes, the optimal policy spends {total / len(episodes):.3f} per episode")


if __name__ == "__main__":
    main()