#!/usr/bin/env python3
"""Run one frozen oral-B-v2 development cell."""
import argparse
import hashlib
import json
import math
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 sdil.bci import generate_trajectories
from sdil.bci_v2 import BCIV2, BCIV2Config, run_day_v2
from sdil.bci_v2_metrics import annotate_events, signature_metrics_v2


ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PROTOCOL_PATH = os.path.join(ROOT, "ORAL_B_V2.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"
)
TASK_SEEDS = (20, 21, 22)
MODEL_SEEDS = (0,)
FORWARD_ETAS = (0.03, 0.1)
GAMMAS = (0.8, 0.95)
CRITIC_ETAS = (0.01, 0.03)
HORIZONS = (4, 8, 12, 16, 20, 24, 28)
CHALLENGE_EPISODES = 128
CONDITIONS = (
    "intact",
    "fixed_role",
    "plasticity_lesion",
    "critic_training_lesion",
    "outcome_training_lesion",
    "oracle_role",
)
ACUTE_MODES = (
    "intact",
    "acute_critic_lesion",
    "acute_outcome_lesion",
)


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 finite_tree(value):
    if isinstance(value, dict):
        return all(finite_tree(item) for item in value.values())
    if isinstance(value, (list, tuple)):
        return all(finite_tree(item) for item in value)
    if isinstance(value, (int, float)):
        return math.isfinite(value)
    return True


def source_paths():
    return {
        "runner": os.path.abspath(__file__),
        "development_analyzer": os.path.join(
            ROOT, "experiments", "analyze_bci_v2_development.py"
        ),
        "confirmation_runner": os.path.join(
            ROOT, "experiments", "bci_v2_confirmation.py"
        ),
        "confirmation_analyzer": os.path.join(
            ROOT, "experiments", "analyze_bci_v2_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"),
        "protocol": PROTOCOL_PATH,
        "d4_gate": D4_GATE_PATH,
        "old_r2_gate": OLD_R2_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)
    if not (
        d4.get("protocol")
        == "kp_dynamic_neutral_projection_confirmation_v1"
        and d4.get("status") == "passed"
        and d4.get("review_score_after") == 7
    ):
        raise ValueError("oral-B-v2 requires the complete D4 accept gate")
    if not (
        old_r2.get("protocol") == "oral_b_td_confirmation_v1"
        and old_r2.get("status") == "failed"
        and old_r2.get("review_score_after") == 7
    ):
        raise ValueError("oral-B-v2 preserves and requires the failed old R2")
    return d4, old_r2


def build_config(forward_eta, gamma, critic_eta):
    return BCIV2Config(
        n_plus=5,
        n_minus=5,
        n_background=30,
        context_dim=16,
        steps_per_episode=28,
        episodes_per_day=64,
        days=14,
        target=0.8,
        inertia=0.65,
        process_noise=0.12,
        context_ar=0.8,
        coupling_scale=1.0,
        predictor_eta=0.2,
        vectorizer_eta=0.03,
        forward_eta=forward_eta,
        perturb_sigma=0.03,
        perturb_every=4,
        kappa=0.0,
        feedback="performance_velocity",
        gamma=gamma,
        critic_eta=critic_eta,
        eligibility_decay=0.8,
        velocity_reward_scale=0.25,
        terminal_reward=1.0,
    )


def condition_settings(name):
    return {
        "intact": {},
        "fixed_role": {"learn_role": False},
        "plasticity_lesion": {"plasticity_gain": 0.0},
        "critic_training_lesion": {
            "learn_critic": False,
            "critic_enabled": False,
        },
        "outcome_training_lesion": {
            "terminal_outcome_enabled": False,
        },
        "oracle_role": {"learn_role": False},
    }[name]


def neutral_warmup(initial, task_seed, model_seed, condition, batches=100):
    """Fit neutral traffic and expose every condition to paired role probes."""
    model = initial.clone()
    generator = torch.Generator(device="cpu").manual_seed(
        600_000 + 100 * task_seed + model_seed
    )
    active = torch.ones(model.cfg.episodes_per_day, dtype=torch.bool)
    update_role = condition not in ("fixed_role", "oracle_role")
    scalar_observations = 0
    for _ in range(batches):
        soma = torch.randn(
            model.cfg.episodes_per_day,
            model.cfg.n_neurons,
            generator=generator,
        ).tanh()
        model.update_predictor(
            soma, model.coupling * soma, active
        )
        perturbations = torch.empty_like(soma).bernoulli_(
            0.5, generator=generator
        ).mul_(2).sub_(1)
        sigma = model.cfg.perturb_sigma
        cursor_plus = (soma + sigma * perturbations) @ model.role
        cursor_minus = (soma - sigma * perturbations) @ model.role
        scalar_observations += 2 * soma.shape[0]
        target = model.causal_role_targets(
            cursor_plus, cursor_minus, perturbations
        )
        if update_role:
            model.update_role(target, active)
    if condition == "oracle_role":
        model.A.copy_(model.role)
    return model, {
        "batches": batches,
        "examples": batches * model.cfg.episodes_per_day,
        "instruction_present": False,
        "role_cursor_scalar_observations": scalar_observations,
        "role_update_enabled": update_role,
        "predictor_max_abs_error": float(
            (model.P - model.coupling).abs().max()
        ),
        "role_cosine_after_warmup": float(
            torch.nn.functional.cosine_similarity(
                model.A, model.role, dim=0
            ).item()
        ),
        "rng_seed": 600_000 + 100 * task_seed + model_seed,
    }


def train(model, trajectories, condition, collect):
    settings = condition_settings(condition)
    daily_success = []
    events = []
    global_step = 0
    active_transitions = 0
    role_probe_examples = 0
    started = time.perf_counter()
    for day in range(model.cfg.days):
        report = run_day_v2(
            model,
            trajectories,
            day,
            global_step=global_step,
            collect=collect,
            learn_predictor=True,
            probe_role=True,
            **settings,
        )
        global_step = report["global_step"]
        active_transitions += report["active_transitions"]
        role_probe_examples += report["role_probe_examples"]
        daily_success.append(report["success_rate"])
        if collect:
            annotate_events(
                report["events"],
                day,
                report["success"],
                day * model.cfg.episodes_per_day,
                model.cfg.steps_per_episode,
            )
            events.extend(report["events"])
    return events, {
        "daily_success": daily_success,
        "early_success": sum(daily_success[:3]) / 3,
        "late_success": sum(daily_success[-3:]) / 3,
        "learning_gain": (
            sum(daily_success[-3:])
            - sum(daily_success[:3])
        ) / 3,
        "training_wall_s": time.perf_counter() - started,
        "role_cosine_after_training": float(
            torch.nn.functional.cosine_similarity(
                model.A, model.role, dim=0
            ).item()
        ),
        "critic_l2_after_training": float(model.critic.norm().item()),
        "cost": {
            "maximum_state_episode_steps": (
                model.cfg.days
                * model.cfg.episodes_per_day
                * model.cfg.steps_per_episode
            ),
            "active_state_episode_steps": active_transitions,
            "role_probe_examples": role_probe_examples,
            "cursor_scalar_observations": 2 * role_probe_examples,
            "terminal_outcome_observations": (
                model.cfg.days * model.cfg.episodes_per_day
            ),
            "task_loss_queries": 0,
            "reverse_mode_calls": 0,
        },
    }


def evaluate_performance(model, trajectories):
    return run_day_v2(
        model.clone(),
        trajectories,
        0,
        horizon=model.cfg.steps_per_episode,
        plasticity_gain=0.0,
        learn_role=False,
        probe_role=False,
        learn_predictor=False,
        learn_critic=False,
        collect=False,
    )


def evaluate_challenge_ladder(model, task_seed, seed_offset):
    """Paired acute readouts over a fixed, non-selected horizon ladder."""
    mode_events = {name: [] for name in ACUTE_MODES}
    costs = {name: 0 for name in ACUTE_MODES}
    seeds = {}
    for horizon_index, horizon in enumerate(HORIZONS):
        trajectory_seed = (
            seed_offset + 1_000 * task_seed + horizon_index
        )
        seeds[str(horizon)] = trajectory_seed
        trajectories = generate_trajectories(
            model.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]
            report = run_day_v2(
                model.clone(),
                trajectories,
                0,
                horizon=horizon,
                plasticity_gain=0.0,
                learn_role=False,
                probe_role=False,
                learn_predictor=False,
                learn_critic=False,
                collect=True,
                **settings,
            )
            episode_offset = (
                1_000_000 + horizon_index * CHALLENGE_EPISODES
            )
            annotate_events(
                report["events"],
                0,
                report["success"],
                episode_offset,
                horizon,
            )
            mode_events[mode].extend(report["events"])
            costs[mode] += report["active_transitions"]
    return mode_events, {
        "horizons": list(HORIZONS),
        "episodes_per_horizon": CHALLENGE_EPISODES,
        "trajectory_seeds": seeds,
        "active_state_episode_steps_by_mode": costs,
        "selection_over_horizons": False,
    }


def run_cell(
    task_seed,
    model_seed,
    forward_eta,
    gamma,
    critic_eta,
    *,
    split,
    performance_seed_offset,
    challenge_seed_offset,
):
    cfg = build_config(forward_eta, gamma, critic_eta)
    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_challenge_ladder(
        trained["intact"], task_seed, challenge_seed_offset
    )
    signatures = signature_metrics_v2(
        training_events,
        challenge_events,
        cfg,
        trained["intact"].role,
    )
    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(
        "--forward-eta", type=float, choices=FORWARD_ETAS, required=True
    )
    parser.add_argument(
        "--gamma", type=float, choices=GAMMAS, required=True
    )
    parser.add_argument(
        "--critic-eta", type=float, choices=CRITIC_ETAS, required=True
    )
    parser.add_argument("--outdir", default="results/bci_v2_dev")
    args = parser.parse_args()
    require_parent_gates()
    source = provenance()
    if (
        source["git_tracked_dirty"]
        or not all(source["tracked_inputs"].values())
    ):
        raise RuntimeError(
            "oral-B-v2 R1 requires clean, tracked, frozen inputs"
        )
    cell = run_cell(
        args.task_seed,
        args.model_seed,
        args.forward_eta,
        args.gamma,
        args.critic_eta,
        split="development",
        performance_seed_offset=400_000,
        challenge_seed_offset=410_000,
    )
    result = {
        "schema_version": 2,
        "protocol": {
            "name": "oral_b_v2_development_v1",
            "selection_split": "development",
            "confirmation_seeds_touched": False,
            "grid_size": (
                len(TASK_SEEDS)
                * len(FORWARD_ETAS)
                * len(GAMMAS)
                * len(CRITIC_ETAS)
            ),
            "protocol_sha256": source["input_sha256"]["protocol"],
            "d4_gate_sha256": source["input_sha256"]["d4_gate"],
            "old_r2_gate_sha256": source["input_sha256"]["old_r2_gate"],
        },
        "args": vars(args),
        "provenance": source,
        **cell,
    }
    result["finite"] = finite_tree(result)
    if not result["finite"]:
        raise RuntimeError("non-finite oral-B-v2 development record")
    os.makedirs(args.outdir, exist_ok=True)
    eta = str(args.forward_eta).replace(".", "p")
    gamma = str(args.gamma).replace(".", "p")
    critic = str(args.critic_eta).replace(".", "p")
    path = os.path.join(
        args.outdir,
        (
            f"bci_v2_e{eta}_g{gamma}_c{critic}"
            f"_t{args.task_seed}_m{args.model_seed}.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"],
        "terminal_outcome_accuracy": result["signatures"][
            "terminal_residual_outcome_balanced_acc"
        ],
        "finite": result["finite"],
    }, indent=2))


if __name__ == "__main__":
    main()