"""Repair the exposed-fact admission by checking the established mismatch control.

No new learning law: this reuses acquisition_controls.py's complete pre-query
drive and wrong-row policy. The failed attention recipes remain frozen. All
cases are taught; confirmation is over founders, not unseen fact combinations.
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import sys
import time
from dataclasses import asdict
from pathlib import Path
from typing import Any

import numpy as np

import cadence
from cadence import Brain, NeuralGraph
from cadence.receipts import Receipt, canonical_sha256, source_manifest

HERE = Path(__file__).resolve().parent
PROTOCOL = HERE / "repair-protocol.json"
SPEC = importlib.util.spec_from_file_location(
    "attention_repair_controls", HERE.parent / "relevance" / "acquisition_controls.py")
assert SPEC is not None and SPEC.loader is not None
controls = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = controls
SPEC.loader.exec_module(controls)
sr = controls.sr


def sources() -> list[tuple[str, Path]]:
    package = Path(cadence.__file__).resolve().parent
    return [
        ("benchmarks/attention/repair_admission.py", Path(__file__).resolve()),
        ("benchmarks/attention/repair-protocol.json", PROTOCOL),
        ("benchmarks/relevance/acquisition_controls.py", Path(controls.__file__)),
        ("benchmarks/relevance/selective_recall.py", Path(sr.__file__)),
    ] + [("src/cadence/" + str(p.relative_to(package)), p)
         for p in sorted(package.rglob("*.py"))]


def train(brain: Brain, order: np.ndarray, policy: str, *, deadline: float,
          neutral_events: int = 1) -> dict[str, Any]:
    """Act then correct from the captured drive, retaining row identity and context.

    Both policies use the same order. A new stimulus on a smaller wrong-row
    batch would lose the prefix trace; a new stimulus after acting would use
    the query-advanced trace. Neither is the event whose answer was judged.
    """
    if policy not in ("every", "wrong"):
        raise ValueError("policy must be every or wrong")
    work = sr.Work()
    rounds, stop = 0, None
    wrong_counts = []
    for permutation in order:
        if time.monotonic() >= deadline:
            stop = "wall_budget"
            break
        rows = sr.CASES[permutation]
        prefix, query, target = sr.inputs(
            rows[:, :3], rows[:, 3], neutral_events, exposed=True)
        brain.reset()
        for observation in prefix:
            if work.act(brain, observation) is None:
                stop = "refused_prefix"
                break
        if stop is not None:
            break
        drive = np.array(brain.stimulus(query), copy=True)
        answer = work.act(brain, query)
        if answer is None:
            stop = "refused_query"
            break
        wrong = np.asarray(answer).reshape(-1) != target
        wrong_counts.append(int(wrong.sum()))
        selected = np.ones(len(target), dtype=bool) if policy == "every" else wrong
        if selected.any():
            controls.teach_drive(work, brain, drive[selected], target[selected])
        rounds += 1
    return {"completed": rounds == len(order), "rounds": rounds, "stop": stop,
            "wrong_counts": wrong_counts, "work": asdict(work)}


def gate(training: dict[str, Any], evaluation: dict[str, Any], floor: float) -> bool:
    # An optimistic completion flag must never erase a refused or unrun branch.
    try:
        branches = evaluation["branches"]
        if set(branches) != {"intact", "erased"}:
            return False
        scores = {name: sr.scores(branch["answers"]) for name, branch in branches.items()}
        complete = all(score["unanswered"] == 0 for score in scores.values())
        intact = scores["intact"]
        return bool(
            training["completed"] and evaluation["completed"] and complete
            and evaluation["original_preserved"]
            and min(intact["accuracy"], intact["unequal_fact_accuracy"],
                    *intact["per_query_accuracy"]) >= floor
        )
    except (KeyError, TypeError, ValueError):
        return False


def make_brain(recipe: dict[str, Any], seed: int) -> Brain:
    brain = Brain.compose(sr.INPUTS, 2, seed=seed, **recipe["brain"])
    original = brain.brain
    brain.learner.brain = NeuralGraph(
        original.connectome, original.neuron_model.replace(dt=recipe["dt"]),
        efficacy=original.efficacy, bias=original.bias, log_gain=original.log_gain)
    return brain


def trial(recipe: dict[str, Any], seed: int, policy: str, protocol: dict[str, Any],
          deadline: float, checkpoints: Path | None) -> dict[str, Any]:
    brain = make_brain(recipe, seed)
    initial_parameters = sr.parameter_digest(brain)
    initial = sr.evaluate(brain, protocol["neutral_events"], exposed=True)
    rng = np.random.default_rng(np.random.SeedSequence([31014, seed]))
    order = np.stack([rng.permutation(len(sr.CASES)) for _ in range(protocol["rounds"])])
    arm_deadline = min(deadline, time.monotonic() + protocol["training_seconds_per_arm"])
    training = train(brain, order, policy, deadline=arm_deadline,
                     neutral_events=protocol["neutral_events"])
    evaluation = sr.evaluate(brain, protocol["neutral_events"], exposed=True)
    checkpoint = None
    if checkpoints is not None:
        path = checkpoints / f"{recipe['name']}-{seed}-{policy}.npz"
        if path.exists():
            raise FileExistsError(path)
        checkpoint = sr.checkpoint_digest(brain.save(path))
    random = np.random.default_rng(np.random.SeedSequence([31015, seed]))
    random_answers = random.integers(0, 2, size=(8, 2))
    return {
        "recipe": recipe["name"], "seed": seed, "policy": policy,
        "resolved": brain.describe(), "neuron_model": brain.brain.neuron_model.to_dict(),
        "initial_parameters": initial_parameters, "final_parameters": sr.parameter_digest(brain),
        "schedule": {"cases": sr.CASES.tolist(), "order_sha256": canonical_sha256(order)},
        "initial": initial, "training": training, "evaluation": evaluation,
        "passed": gate(training, evaluation, protocol["acquisition_floor"]),
        "checkpoint_content_sha256": checkpoint,
        "uniform_random": {"answers": random_answers.tolist(), "scores": sr.scores(random_answers),
                           "expected_accuracy": 0.5},
    }


def selected_recipe(runs: list[dict[str, Any]], protocol: dict[str, Any]) -> str | None:
    """First recipe whose wrong-policy lives ALL pass; not a best-checkpoint search."""
    for recipe in protocol["recipes"]:
        group = [r for r in runs if r["recipe"] == recipe["name"] and r["policy"] == "wrong"]
        if ([r["seed"] for r in group] == protocol["development_seeds"]
                and all(r["passed"] for r in group)):
            return str(recipe["name"])
    return None


def run(stage: str, development: Path | None = None,
        checkpoints: Path | None = None) -> Receipt:
    protocol = json.loads(PROTOCOL.read_text())
    manifest = source_manifest(sources())
    began = time.monotonic()
    spent, reference = 0.0, None
    selected = None
    if stage == "confirmation":
        if development is None:
            raise ValueError("confirmation requires the frozen development receipt")
        valid, message = Receipt.verify(development, sources=sources(), check=check)
        if not valid:
            raise ValueError(message)
        prior = Receipt.read(development)
        if prior.body["stage"] != "development" or not prior.body["passed"]:
            raise ValueError("development must pass before opening confirmation seeds")
        selected = prior.body["selected_recipe"]
        spent = prior.body["elapsed_host_seconds"]
        reference = {"digest": prior.digest, "selected_recipe": selected}
    elif stage != "development":
        raise ValueError("unknown stage")
    deadline = began + max(0.0, protocol["total_seconds"] - spent)
    if checkpoints is not None:
        checkpoints.mkdir(parents=True, exist_ok=True)
    runs = []
    for recipe in protocol["recipes"]:
        if stage == "confirmation" and recipe["name"] != selected:
            continue
        seeds = protocol[f"{stage}_seeds"]
        for seed in seeds:
            for policy in ("every", "wrong") if stage == "development" else ("wrong",):
                result = trial(recipe, seed, policy, protocol, deadline, checkpoints)
                runs.append(result)
                scores = result["evaluation"]["branches"]["intact"]["scores"]
                print(json.dumps({"recipe": recipe["name"], "seed": seed, "policy": policy,
                                  "passed": result["passed"],
                                  "rounds": result["training"]["rounds"],
                                  "accuracy": scores["accuracy"]}),
                      flush=True)
        if stage == "development":
            selected = selected_recipe(runs, protocol)
            if selected is not None:
                break
    if manifest != source_manifest(sources()):
        raise RuntimeError("source changed during the experiment; retain diagnostics and rerun")
    passed = selected is not None
    if stage == "confirmation":
        passed = (len(runs) == protocol["confirmation_required"]
                  and all(r["passed"] for r in runs))
    body = {
        "stage": stage, "protocol": protocol, "runs": runs, "selected_recipe": selected,
        "passed": passed, "development_reference": reference,
        "elapsed_host_seconds": time.monotonic() - began, "previous_host_seconds": spent,
        "numpy_version": np.__version__, "cadence_version": cadence.__version__,
        "source_revision": sr.revision(),
        "source_identity": (
            "The manifest binds imported code and protocol, including uncommitted bytes."),
        "interpretation": (
            "Finite exposed-fact acquisition is admitted; learned retrieval, semantic relevance "
            "and lifecycle integration require their own gates." if passed else
            "Acquisition admission failed; retain every failed/refused/unrun founder and stop."),
    }
    problem = check(body)
    if problem:
        raise AssertionError(problem)
    return Receipt.build("cadence-attention-acquisition-repair/1", body, sources())


def check(body: dict[str, Any]) -> str | None:
    """Recompute decisions from actual answers, never trust a stored pass flag."""
    try:
        protocol, runs, stage = body["protocol"], body["runs"], body["stage"]
        if protocol != json.loads(PROTOCOL.read_text()):
            return "protocol differs from the frozen declaration"
        if stage not in ("development", "confirmation"):
            return "unknown stage"
        recipes = protocol["recipes"]
        names = [recipe["name"] for recipe in recipes]
        for result in runs:
            if result["recipe"] not in names:
                return "undeclared recipe"
            for evaluation in (result["initial"], result["evaluation"]):
                branches = evaluation["branches"]
                if set(branches) != {"intact", "erased"}:
                    return "evaluation branches differ"
                complete = True
                for branch in branches.values():
                    scores = sr.scores(branch["answers"])
                    if scores != branch["scores"]:
                        return "scores differ from recorded answers"
                    complete = complete and scores["unanswered"] == 0
                if evaluation["completed"] != complete:
                    return "evaluation completion differs from recorded answers"
                if not evaluation["original_preserved"]:
                    return "evaluation changed the acquired brain"
            random = result["uniform_random"]
            if sr.scores(random["answers"]) != random["scores"]:
                return "random scores differ"
            training = result["training"]
            rounds = training["rounds"]
            if (type(rounds) is not int or not 0 <= rounds <= protocol["rounds"]
                    or training["completed"] != (rounds == protocol["rounds"])):
                return "training completion differs"
            if result["passed"] != gate(
                training, result["evaluation"], protocol["acquisition_floor"]
            ):
                return "founder gate differs"
            rng = np.random.default_rng(np.random.SeedSequence([31014, result["seed"]]))
            order = np.stack([rng.permutation(len(sr.CASES)) for _ in range(protocol["rounds"])])
            if result["schedule"] != {
                "cases": sr.CASES.tolist(), "order_sha256": canonical_sha256(order)
            }:
                return "schedule differs from the declared founder order"
            recipe = recipes[names.index(result["recipe"])]
            if result["neuron_model"]["dt"] != recipe["dt"]:
                return "neuron model differs from the declared recipe"
        triples = [(r["recipe"], r["seed"], r["policy"]) for r in runs]
        if stage == "development":
            expected = selected_recipe(runs, protocol)
            used = recipes if expected is None else recipes[:names.index(expected) + 1]
            expected_triples = [
                (recipe["name"], seed, policy)
                for recipe in used for seed in protocol["development_seeds"]
                for policy in ("every", "wrong")
            ]
            if triples != expected_triples:
                return "development founders, arms or recipe order differ"
            for every, wrong in zip(runs[::2], runs[1::2], strict=True):
                for key in ("initial_parameters", "schedule", "resolved", "neuron_model"):
                    if every[key] != wrong[key]:
                        return "development policies have unmatched " + key
            if expected != body["selected_recipe"] or body["passed"] != (expected is not None):
                return "development selection differs"
        else:
            selected = body["selected_recipe"]
            if (selected not in names or triples != [
                (selected, seed, "wrong") for seed in protocol["confirmation_seeds"]
            ]):
                return "confirmation founders or recipe differ"
            passed = (len(runs) == protocol["confirmation_required"]
                      and all(r["passed"] for r in runs))
            if body["passed"] != passed:
                return "confirmation gate differs"
    except (KeyError, TypeError, ValueError, IndexError):
        return "malformed acquisition repair readings"
    return None


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--stage", choices=("development", "confirmation"))
    parser.add_argument("--development", type=Path)
    parser.add_argument("--checkpoints", type=Path)
    parser.add_argument("--out", type=Path)
    parser.add_argument("--verify", type=Path)
    args = parser.parse_args()
    if args.verify:
        ok, message = Receipt.verify(args.verify, sources=sources(), check=check)
        print(message)
        raise SystemExit(0 if ok else 1)
    if not args.stage or not args.out:
        parser.error("provide --stage and --out")
    if args.out.exists():
        parser.error("output already exists; preserve the earlier receipt")
    run(args.stage, args.development, args.checkpoints).write(args.out)
