#!/usr/bin/env python3
"""Run one frozen oral-B-v2 cold-start recovery development cell."""
import argparse
from dataclasses import replace
import hashlib
import json
import os
import platform
import resource
import subprocess
import sys
import time

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from experiments.bci_v2_run import (
    ACUTE_MODES,
    CONDITIONS,
    build_config,
    evaluate_performance,
    finite_tree,
    neutral_warmup,
    train,
)
from sdil.bci import generate_trajectories
from sdil.bci_v2 import BCIV2, run_day_v2
from sdil.bci_v2_recovery_metrics import (
    annotate_target_events,
    recovery_signature_metrics,
)


ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PROTOCOL_PATH = os.path.join(ROOT, "ORAL_B_V2_RECOVERY.md")
D4_GATE_PATH = os.path.join(
    ROOT, "results", "kp_dynamic_projection_confirmation_gate.json"
)
OLD_R2_GATE_PATH = os.path.join(
    ROOT, "results", "bci_td_confirmation_gate.json"
)
FAILED_V2_GATE_PATH = os.path.join(
    ROOT, "results", "bci_v2_dev_gate.json"
)
TASK_SEEDS = (23, 24, 25)
MODEL_SEEDS = (0,)
TARGETS = (1.55, 1.60, 1.65, 1.70, 1.75, 1.80)
CHALLENGE_EPISODES = 128
FIXED_CONFIG = {
    "forward_eta": 0.1,
    "gamma": 0.8,
    "critic_eta": 0.03,
    "velocity_reward_scale": 1.0,
}


def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def source_paths():
    return {
        "runner": os.path.abspath(__file__),
        "development_analyzer": os.path.join(
            ROOT,
            "experiments",
            "analyze_bci_v2_recovery_development.py",
        ),
        "confirmation_runner": os.path.join(
            ROOT, "experiments", "bci_v2_recovery_confirmation.py"
        ),
        "confirmation_analyzer": os.path.join(
            ROOT,
            "experiments",
            "analyze_bci_v2_recovery_confirmation.py",
        ),
        "base_dynamics": os.path.join(ROOT, "sdil", "bci.py"),
        "v2_dynamics": os.path.join(ROOT, "sdil", "bci_v2.py"),
        "v2_metrics": os.path.join(ROOT, "sdil", "bci_v2_metrics.py"),
        "recovery_metrics": os.path.join(
            ROOT, "sdil", "bci_v2_recovery_metrics.py"
        ),
        "protocol": PROTOCOL_PATH,
        "d4_gate": D4_GATE_PATH,
        "old_r2_gate": OLD_R2_GATE_PATH,
        "failed_v2_gate": FAILED_V2_GATE_PATH,
    }


def provenance(extra_paths=None):
    def run(command):
        return subprocess.run(
            command,
            cwd=ROOT,
            check=True,
            capture_output=True,
            text=True,
        ).stdout.strip()

    paths = source_paths()
    if extra_paths:
        paths.update(extra_paths)
    relative = {
        name: os.path.relpath(path, ROOT)
        for name, path in paths.items()
    }
    tracked = {
        name: subprocess.run(
            ["git", "ls-files", "--error-unmatch", path],
            cwd=ROOT,
            capture_output=True,
        ).returncode == 0
        for name, path in relative.items()
    }
    return {
        "git_commit": run(["git", "rev-parse", "HEAD"]),
        "git_tracked_dirty": bool(run([
            "git", "status", "--porcelain", "--untracked-files=no"
        ])),
        "tracked_inputs": tracked,
        "input_sha256": {
            name: sha256(path) for name, path in paths.items()
        },
    }


def require_parent_gates():
    with open(D4_GATE_PATH) as handle:
        d4 = json.load(handle)
    with open(OLD_R2_GATE_PATH) as handle:
        old_r2 = json.load(handle)
    with open(FAILED_V2_GATE_PATH) as handle:
        failed_v2 = json.load(handle)
    if not (
        d4.get("status") == "passed"
        and d4.get("review_score_after") == 7
        and old_r2.get("status") == "failed"
        and old_r2.get("review_score_after") == 7
        and failed_v2.get("protocol") == "oral_b_v2_development_v1"
        and failed_v2.get("status") == "failed"
        and failed_v2.get("complete_grid") is True
        and failed_v2.get("oral_b_v2_confirmation_opened") is False
        and failed_v2.get("review_score_after") == 7
    ):
        raise ValueError(
            "cold-start recovery requires D4 and both preserved failures"
        )


def build_recovery_config():
    return replace(
        build_config(
            FIXED_CONFIG["forward_eta"],
            FIXED_CONFIG["gamma"],
            FIXED_CONFIG["critic_eta"],
        ),
        velocity_reward_scale=FIXED_CONFIG[
            "velocity_reward_scale"
        ],
    )


def evaluate_target_ladder(model, task_seed, seed_offset):
    mode_events = {name: [] for name in ACUTE_MODES}
    costs = {name: 0 for name in ACUTE_MODES}
    seeds = {}
    for target_index, target in enumerate(TARGETS, start=1):
        challenge_cfg = replace(model.cfg, target=target)
        trajectory_seed = (
            seed_offset + 1_000 * task_seed + target_index - 1
        )
        seeds[str(target)] = trajectory_seed
        trajectories = generate_trajectories(
            challenge_cfg,
            trajectory_seed,
            days=1,
            episodes=CHALLENGE_EPISODES,
        )
        for mode in ACUTE_MODES:
            settings = {
                "intact": {},
                "acute_critic_lesion": {"critic_enabled": False},
                "acute_outcome_lesion": {
                    "terminal_outcome_enabled": False
                },
            }[mode]
            assay_model = model.clone()
            assay_model.cfg = challenge_cfg
            report = run_day_v2(
                assay_model,
                trajectories,
                0,
                horizon=challenge_cfg.steps_per_episode,
                plasticity_gain=0.0,
                learn_role=False,
                probe_role=False,
                learn_predictor=False,
                learn_critic=False,
                collect=True,
                **settings,
            )
            episode_offset = (
                2_000_000
                + (target_index - 1) * CHALLENGE_EPISODES
            )
            annotate_target_events(
                report["events"],
                0,
                report["success"],
                episode_offset,
                target_index,
            )
            mode_events[mode].extend(report["events"])
            costs[mode] += report["active_transitions"]
    return mode_events, {
        "targets": list(TARGETS),
        "episodes_per_target": CHALLENGE_EPISODES,
        "trajectory_seeds": seeds,
        "active_state_episode_steps_by_mode": costs,
        "selection_over_targets": False,
        "maximum_steps_per_episode": model.cfg.steps_per_episode,
    }


def run_cell(
    task_seed,
    model_seed,
    *,
    split,
    performance_seed_offset,
    challenge_seed_offset,
):
    cfg = build_recovery_config()
    trajectories = generate_trajectories(cfg, task_seed)
    performance_seed = performance_seed_offset + task_seed
    performance_trajectories = generate_trajectories(
        cfg, performance_seed, days=1, episodes=256
    )
    initial = BCIV2(cfg, model_seed)
    conditions = {}
    trained = {}
    warmups = {}
    training_events = None
    started = time.perf_counter()
    for name in CONDITIONS:
        model, warmup = neutral_warmup(
            initial, task_seed, model_seed, name
        )
        events, report = train(
            model, trajectories, name, collect=name == "intact"
        )
        warmups[name] = warmup
        conditions[name] = report
        trained[name] = model
        if name == "intact":
            training_events = events
    for name in CONDITIONS:
        report = evaluate_performance(
            trained[name], performance_trajectories
        )
        conditions[name]["final_success"] = report["success_rate"]
        conditions[name]["evaluation_active_state_episode_steps"] = (
            report["active_transitions"]
        )
    challenge_events, challenge_cost = evaluate_target_ladder(
        trained["intact"], task_seed, challenge_seed_offset
    )
    signatures = recovery_signature_metrics(
        training_events,
        challenge_events,
        cfg,
        trained["intact"].role,
        TARGETS,
    )
    return {
        "config": vars(cfg),
        "warmup": warmups,
        "conditions": conditions,
        "signatures": signatures,
        "assays": {
            "performance_evaluation_seed": performance_seed,
            "performance_evaluation_episodes": 256,
            "challenge": challenge_cost,
        },
        "wall_s": time.perf_counter() - started,
        "peak_rss_mib": (
            resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
        ),
        "hardware": {
            "device": "cpu",
            "platform": platform.platform(),
            "torch_version": torch.__version__,
            "threads": torch.get_num_threads(),
        },
        "split": split,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--task-seed", type=int, choices=TASK_SEEDS, required=True
    )
    parser.add_argument(
        "--model-seed", type=int, choices=MODEL_SEEDS, default=0
    )
    parser.add_argument(
        "--outdir", default="results/bci_v2_recovery_dev"
    )
    args = parser.parse_args()
    require_parent_gates()
    source = provenance()
    if (
        source["git_tracked_dirty"]
        or not all(source["tracked_inputs"].values())
    ):
        raise RuntimeError(
            "v2 recovery R1 requires clean, tracked, frozen inputs"
        )
    cell = run_cell(
        args.task_seed,
        args.model_seed,
        split="development",
        performance_seed_offset=460_000,
        challenge_seed_offset=470_000,
    )
    result = {
        "schema_version": 3,
        "protocol": {
            "name": "oral_b_v2_cold_start_recovery_development_v1",
            "selection_split": "development_validation",
            "hyperparameter_selection": False,
            "confirmation_seeds_touched": False,
            "fixed_config": FIXED_CONFIG,
            "protocol_sha256": source["input_sha256"]["protocol"],
            "d4_gate_sha256": source["input_sha256"]["d4_gate"],
            "old_r2_gate_sha256": source["input_sha256"]["old_r2_gate"],
            "failed_v2_gate_sha256": source["input_sha256"][
                "failed_v2_gate"
            ],
        },
        "args": vars(args),
        "provenance": source,
        **cell,
    }
    result["finite"] = finite_tree(result)
    if not result["finite"]:
        raise RuntimeError("non-finite v2 recovery development record")
    os.makedirs(args.outdir, exist_ok=True)
    path = os.path.join(
        args.outdir,
        f"bci_v2_recovery_t{args.task_seed}_m0.json",
    )
    if os.path.exists(path):
        raise FileExistsError(f"refusing to overwrite {path}")
    with open(path, "w") as handle:
        json.dump(result, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(json.dumps({
        "path": path,
        "intact_final": result["conditions"]["intact"]["final_success"],
        "challenge_success_fraction": result["signatures"][
            "challenge_success_fraction"
        ],
        "terminal_outcome_accuracy": result["signatures"][
            "terminal_residual_outcome_balanced_acc"
        ],
        "finite": result["finite"],
    }, indent=2))


if __name__ == "__main__":
    main()