#!/usr/bin/env python3
"""Frozen 17-cell B1 contrastive state-bias screen on author Dual Prop."""
import argparse
import hashlib
import json
import math
import os
from pathlib import Path
import subprocess
import time

import numpy as np


ROOT = Path(__file__).resolve().parents[1]
PROTOCOL = ROOT / "CONTRASTIVE_BIAS.md"
RESULT_ROOT = ROOT / "results" / "contrastive_bias" / "b1"
BIAS_PATCH = (
    ROOT / "external" / "dualprop_patches" /
    "0020-experiment-add-neuron-specific-bias-to-Dual-Prop.patch"
)
UPSTREAM = "7b2595b34421e1483a721dbfdeff8cdabda3a1ff"
RULES = ("raw", "innovation", "oracle")


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 git_output(repo, *args):
    return subprocess.run(
        ["git", *args], cwd=repo, check=True, capture_output=True, text=True
    ).stdout.strip()


def rate_tag(value):
    return f"{value:g}".replace(".", "p")


def bias_cells():
    cells = [{
        "cell_id": "clean", "kind": "none", "rule": "none", "ratio": 0.0,
    }, {
        "cell_id": "common-activity-r4-raw", "kind": "common", "rule": "raw",
        "ratio": 4.0,
    }]
    for kind, ratios in (("fixed", (1.0, 4.0)),
                         ("activity", (0.25, 1.0, 4.0))):
        for ratio in ratios:
            for rule in RULES:
                cells.append({
                    "cell_id": f"{kind}-r{rate_tag(ratio)}-{rule}",
                    "kind": kind, "rule": rule, "ratio": ratio,
                })
    if len(cells) != 17 or len({row["cell_id"] for row in cells}) != 17:
        raise AssertionError("B1 registry must contain 17 unique cells")
    return cells


def author_command(cell, author_python):
    name = "dp-bias-b1-" + cell["cell_id"]
    command = [
        author_python, "train.py",
        "--model", "miniCNN", "--dataset", "cifar10",
        "--num-epochs", "20", "--batch-size", "100",
        "--learning-rate", "0.025", "--learning-rate-final", "0.025",
        "--warmup-learning-rate", "0.025", "--warmup-epochs", "0",
        "--decay-epochs", "20", "--momentum", "0.9",
        "--weight-decay", "5e-4", "--dtype", "float32",
        "--param-dtype", "float32", "--percent-train", "90",
        "--percent-val", "10", "--seeds", "1988",
        "--feedback-seed", "1729", "--gradient-diagnostics", "none",
        "--spectral-diagnostics", "none", "--test-policy", "none",
        "--early-stop-policy", "none", "--learning-algorithm",
        "dualprop-lagr-ff", "--experiment-name", name,
        "--optimizer-schedule", "author", "--loss", "sce",
        "--alpha", "0.0", "--beta", "0.1",
        "--inference-sequence", "fwK", "--inference-passes-nudged", "16",
    ]
    if cell["rule"] != "none":
        command.extend([
            "--dp-bias-kind", cell["kind"], "--dp-bias-rule", cell["rule"],
            "--dp-bias-ratio", str(cell["ratio"]), "--dp-bias-seed", "6100",
            "--dp-bias-calibration-examples", "64",
        ])
    return name, command


def jobs(author_python):
    rows = []
    for cell in bias_cells():
        name, command = author_command(cell, author_python)
        rows.append({
            **cell, "experiment_name": name, "command": command,
            "output": str(RESULT_ROOT / (name + ".json")),
            "timeout_seconds": 2 * 60 * 60,
        })
    return rows


def registry_sha256(rows):
    payload = [
        {key: value for key, value in row.items() if key != "output"}
        for row in rows
    ]
    return hashlib.sha256(json.dumps(
        payload, sort_keys=True, separators=(",", ":")
    ).encode()).hexdigest()


def source_report(author_root):
    if git_output(ROOT, "status", "--porcelain", "--untracked-files=no"):
        raise RuntimeError("B1 requires clean tracked SDIL source")
    if git_output(author_root, "status", "--porcelain", "--untracked-files=no"):
        raise RuntimeError("B1 requires clean tracked author source")
    tracked = [
        PROTOCOL, Path(__file__).resolve(),
        ROOT / "experiments" / "analyze_contrastive_bias_b1.py",
        BIAS_PATCH,
    ]
    for path in tracked:
        relative = path.relative_to(ROOT)
        subprocess.run(
            ["git", "ls-files", "--error-unmatch", str(relative)], cwd=ROOT,
            check=True, capture_output=True,
        )
    return {
        "sdil_commit": git_output(ROOT, "rev-parse", "HEAD"),
        "author_commit": git_output(author_root, "rev-parse", "HEAD"),
        "author_upstream": UPSTREAM,
        "tracked_files": {
            str(path.relative_to(ROOT)): sha256(path) for path in tracked
        },
    }


def gpu_report(physical_index):
    output = subprocess.run([
        "nvidia-smi", f"--id={physical_index}",
        "--query-gpu=index,uuid,name,memory.total", "--format=csv,noheader,nounits",
    ], check=True, capture_output=True, text=True).stdout.strip()
    rows = [part.strip() for part in output.split(",")]
    if len(rows) != 4 or rows[0] != str(physical_index):
        raise RuntimeError(f"could not resolve physical GPU {physical_index}: {output}")
    visible = os.environ.get("CUDA_VISIBLE_DEVICES")
    if visible != str(physical_index):
        raise RuntimeError(
            f"CUDA_VISIBLE_DEVICES must equal physical GPU {physical_index}, got {visible}")
    return {
        "physical_index": int(rows[0]), "uuid": rows[1], "name": rows[2],
        "memory_total_mib": int(rows[3]), "cuda_visible_devices": visible,
    }


def ensure_launch(source, rows):
    path = RESULT_ROOT / "launch.json"
    expected = {
        "stage": "contrastive_bias_b1", "source": source,
        "registry_sha256": registry_sha256(rows), "num_jobs": len(rows),
        "allowed_physical_gpus": [5, 7],
    }
    if path.is_file():
        with open(path, encoding="utf-8") as handle:
            if json.load(handle) != expected:
                raise RuntimeError("B1 launch lock drift")
    else:
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w", encoding="utf-8") as handle:
            json.dump(expected, handle, indent=2, sort_keys=True)
            handle.write("\n")
    return path


def to_float_list(value, count):
    array = np.asarray(value)[:count]
    return [float(item) for item in array]


def summarize_hist(path):
    hist = np.load(path, allow_pickle=True).item()
    completed = int(hist["epochs_completed"])
    keys = (
        "val_loss", "val_accuracy", "train_loss", "train_accuracy",
        "train_time", "val_time", "raw_bias_clean_difference_rms_ratio",
        "post_bias_raw_bias_rms_ratio", "used_clean_difference_rms_ratio",
        "maximum_used_clean_difference_relative_error", "neutral_observations",
        "instruction_observations_for_predictor",
    )
    curves = {key: to_float_list(hist[key], completed) for key in keys}
    finite = completed == 20 and all(
        math.isfinite(value)
        for key in ("val_loss", "val_accuracy", "train_loss")
        for value in curves[key]
    )
    return {
        "epochs_completed": completed, "finite": finite,
        "final_validation_accuracy": curves["val_accuracy"][-1],
        "best_validation_accuracy": float(hist["best_validation_accuracy"]),
        "best_epoch": int(hist["best_epoch"]),
        "test_accuracy": float(hist["test_accuracy"]),
        "dp_bias_initialization": hist.get("dp_bias_initialization"),
        "curves": curves,
    }


def find_hist(author_root, experiment_name):
    paths = list((author_root / "runs" / experiment_name).glob("*/hist.npy"))
    if len(paths) != 1:
        raise RuntimeError(
            f"expected one history for {experiment_name}, found {len(paths)}")
    return paths[0]


def run_job(job, author_root, source, registry_hash, gpu, dry_run):
    output = Path(job["output"])
    if output.exists():
        print(f"preserving {job['cell_id']}", flush=True)
        return
    print("RUN", " ".join(job["command"]), flush=True)
    if dry_run:
        return
    if git_output(ROOT, "rev-parse", "HEAD") != source["sdil_commit"]:
        raise RuntimeError("SDIL commit changed after B1 launch")
    if git_output(author_root, "rev-parse", "HEAD") != source["author_commit"]:
        raise RuntimeError("author commit changed after B1 launch")
    started = time.time()
    try:
        result = subprocess.run(
            job["command"], cwd=author_root, timeout=job["timeout_seconds"])
        return_code = result.returncode
        status = "completed" if return_code == 0 else "nonzero_exit"
    except subprocess.TimeoutExpired:
        return_code, status = None, "timeout"
    history = None
    history_path = None
    if status == "completed":
        try:
            resolved = find_hist(author_root, job["experiment_name"])
            history_path = str(resolved)
            history = summarize_hist(resolved)
        except Exception as error:
            status = "missing_or_invalid_history"
            history = {"error": repr(error)}
    record = {
        **job, "stage": "contrastive_bias_b1", "source": source,
        "registry_sha256": registry_hash, "hardware": gpu, "status": status,
        "return_code": return_code, "driver_wall_seconds": time.time() - started,
        "completed_unix_time": time.time(), "author_history": history_path,
        "history": history,
    }
    output.parent.mkdir(parents=True, exist_ok=True)
    with open(output, "w", encoding="utf-8") as handle:
        json.dump(record, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(f"DONE status={status} {job['cell_id']}", flush=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--author-root", type=Path, required=True)
    parser.add_argument("--author-python", required=True)
    parser.add_argument("--physical-gpu", type=int, choices=(5, 7), required=True)
    parser.add_argument("--shard-index", type=int, default=0)
    parser.add_argument("--num-shards", type=int, default=1)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    args.author_root = args.author_root.resolve()
    if not 0 <= args.shard_index < args.num_shards:
        raise ValueError("invalid B1 shard")
    rows = jobs(args.author_python)
    selected = [
        row for index, row in enumerate(rows)
        if index % args.num_shards == args.shard_index
    ]
    if args.dry_run:
        for row in selected:
            print(row["cell_id"], " ".join(row["command"]))
        return
    source = source_report(args.author_root)
    gpu = gpu_report(args.physical_gpu)
    launch = ensure_launch(source, rows)
    print(f"B1 launch lock: {launch}", flush=True)
    registry_hash = registry_sha256(rows)
    for row in selected:
        run_job(row, args.author_root, source, registry_hash, gpu, False)


if __name__ == "__main__":
    main()