#!/usr/bin/env python3
"""Immutable driver for the complete 27-cell ResNet R2 crossover."""
import argparse
import hashlib
import json
import os
import subprocess
import sys
import time


ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from experiments.crossover_hardware import (
    DEFAULT_PROFILE,
    physical_gpu_report,
    policy_for_report,
    profile_choices,
)


PROTOCOL = os.path.join(ROOT, "RESNET_CROSSOVER.md")
RESULT_ROOT = os.path.join(ROOT, "results", "resnet_crossover")
DEFAULT_SELECTOR = os.path.join(RESULT_ROOT, "p1_selector.json")
METHODS = (
    "bp", "fa", "dfa", "pepita", "ff", "ep", "dualprop",
    "clean_kp", "sdil",
)
DEPTHS = (20, 32, 56)
ORDER = (
    ("ep", 56),
    ("dualprop", 56),
    ("ff", 56),
    ("ep", 32),
    ("dualprop", 32),
    ("ff", 32),
    ("ep", 20),
    ("dualprop", 20),
    ("ff", 20),
    ("pepita", 56),
    ("sdil", 56),
    ("clean_kp", 56),
    ("fa", 56),
    ("dfa", 56),
    ("bp", 56),
    ("pepita", 32),
    ("sdil", 32),
    ("clean_kp", 32),
    ("fa", 32),
    ("dfa", 32),
    ("bp", 32),
    ("pepita", 20),
    ("sdil", 20),
    ("clean_kp", 20),
    ("fa", 20),
    ("dfa", 20),
    ("bp", 20),
)


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


def read_json(path):
    with open(path, encoding="utf-8") as handle:
        return json.load(handle)


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


def selector_report(path):
    selector = read_json(path)
    if not (
        selector.get("gate") == "pass"
        and selector.get("stage") == "resnet_crossover_p1"
        and selector.get("num_expected_records") == 19
        and selector.get("num_audited_records") == 19
        and selector.get("missing_experiments") == []
        and selector.get("test_policy") == "none"
        and set(selector.get("selected", {})) == set(METHODS)
    ):
        raise RuntimeError("ResNet R2 requires the complete passed P1 selector")
    selected = {}
    for method in METHODS:
        row = selector["selected"][method]
        rate = float(row["rate"])
        if not rate > 0:
            raise RuntimeError(f"invalid selected rate for {method}")
        selected[method] = rate
    return {
        "path": os.path.relpath(os.path.abspath(path), ROOT),
        "sha256": sha256(path),
        "p1_source": selector["source"],
        "selected_rates": selected,
    }


def source_report(selector_path):
    dirty = git_output("status", "--porcelain", "--untracked-files=no")
    if dirty:
        raise RuntimeError("ResNet R2 requires clean tracked source")
    paths = [
        PROTOCOL,
        os.path.abspath(__file__),
        os.path.join(ROOT, "experiments", "analyze_resnet_crossover_r2.py"),
        os.path.join(ROOT, "experiments", "resnet_crossover_native.py"),
        os.path.join(ROOT, "experiments", "conv_run.py"),
        os.path.join(ROOT, "experiments", "crossover_hardware.py"),
        os.path.join(ROOT, "sdil", "conv_crossover.py"),
        os.path.join(ROOT, "sdil", "conv.py"),
        os.path.join(ROOT, "sdil", "data.py"),
        os.path.abspath(selector_path),
    ]
    for path in paths:
        relative = os.path.relpath(path, ROOT)
        tracked = subprocess.run(
            ["git", "ls-files", "--error-unmatch", relative],
            cwd=ROOT,
            capture_output=True,
        ).returncode == 0
        if not tracked:
            raise RuntimeError(f"untracked ResNet R2 source: {relative}")
    return {
        "git_commit": git_output("rev-parse", "HEAD"),
        "protocol_path": PROTOCOL,
        "protocol_sha256": sha256(PROTOCOL),
        "tracked_files": {
            os.path.relpath(path, ROOT): sha256(path) for path in paths
        },
        "python": sys.executable,
    }


def schedule(method):
    if method == "pepita":
        return {
            "epochs": 100,
            "lr_schedule": "pepita",
            "lr_milestones": "60,90",
        }
    if method == "ff":
        return {
            "epochs": 40,
            "lr_schedule": "constant",
            "lr_milestones": "100,150",
        }
    if method == "ep":
        return {
            "epochs": 100,
            "lr_schedule": "constant",
            "lr_milestones": "100,150",
        }
    return {
        "epochs": 200,
        "lr_schedule": "step",
        "lr_milestones": "100,150",
    }


def r2_command(method, depth, rate):
    name = (
        f"resnet-r2-{method}-d{depth}-lr{rate_tag(rate)}"
    )
    output = os.path.join(RESULT_ROOT, "r2", name + ".json")
    method_schedule = schedule(method)
    output_rate = 0.1 if method in ("fa", "dfa") else rate
    common = [
        "--device", "cuda",
        "--depth", str(depth),
        "--width", "16",
        "--seed", "0",
        "--loader_seed", "0",
        "--split_seed", "2027",
        "--batch_size", "128",
        "--epochs", str(method_schedule["epochs"]),
        "--train_limit", "0",
        "--val_examples", "5000",
        "--eval_split", "validation",
        "--eval_every", "1",
        "--augment_train", "1",
        "--lr", str(rate),
        "--output_lr", str(output_rate),
        "--lr_schedule", method_schedule["lr_schedule"],
        "--lr_milestones", method_schedule["lr_milestones"],
        "--lr_gamma", "0.1",
        "--momentum", "0.9",
        "--weight_decay", "1e-4",
    ]
    if method in ("pepita", "ff", "ep", "dualprop"):
        command = [
            sys.executable,
            "experiments/resnet_crossover_native.py",
            "--method", method,
            "--out", output,
            "--feedback_seed", "1729",
            *common,
        ]
        if method == "pepita":
            command.extend(["--pepita_projection_scale", "0.05"])
        elif method == "ff":
            command.extend([
                "--ff_threshold", "2.0",
                "--ff_score_from_layer", "1",
            ])
        elif method == "ep":
            command.extend([
                "--ep_beta", "0.5",
                "--ep_dt", "0.5",
                "--ep_free_steps", "20",
                "--ep_nudge_steps", "4",
            ])
        else:
            command.extend([
                "--dp_alpha", "0.0",
                "--dp_beta", "0.1",
                "--dp_inference_passes", "16",
            ])
    else:
        modes = {
            "bp": "bp",
            "fa": "hfa",
            "dfa": "dfa",
            "clean_kp": "kp",
            "sdil": "kp_traffic",
        }
        command = [
            sys.executable,
            "experiments/conv_run.py",
            "--mode", modes[method],
            "--out", output,
            "--normalization", "batchnorm",
            "--a_scale", "1",
            "--alignment_probe", "0" if method == "bp" else "32",
            *common,
        ]
        if method == "dfa":
            command.extend(["--vectorizer_mode", "spatial_template"])
        elif method == "sdil":
            command.extend([
                "--traffic_rule", "innovation",
                "--predictor_mode", "closed_form",
                "--neutral_projection", "1",
                "--traffic_seed", "5000",
                "--traffic_ratio", "4",
                "--traffic_calibration_examples", "64",
                "--learn_P", "1",
                "--eta_P", "0.1",
                "--predictor_warmup_steps", "1",
                "--predictor_every", "0",
            ])
    return {
        "stage": "r2",
        "method": method,
        "architecture": f"resnet{depth}",
        "depth": depth,
        "rate": rate,
        "epochs": method_schedule["epochs"],
        "lr_schedule": method_schedule["lr_schedule"],
        "experiment_name": name,
        "output": output,
        "timeout_seconds": 48 * 60 * 60,
        "command": command,
    }


def r2_jobs(selected_rates):
    if set(selected_rates) != set(METHODS):
        raise ValueError("R2 requires one selected rate for every method")
    jobs = [
        r2_command(method, depth, float(selected_rates[method]))
        for method, depth in ORDER
    ]
    identifiers = [job["experiment_name"] for job in jobs]
    cells = {(job["method"], job["depth"]) for job in jobs}
    if (
        len(jobs) != 27
        or len(set(identifiers)) != 27
        or cells != {(method, depth) for method in METHODS for depth in DEPTHS}
    ):
        raise AssertionError("ResNet R2 registry must contain 27 unique cells")
    return jobs


def registry_sha256(jobs):
    encoded = json.dumps(
        jobs, sort_keys=True, separators=(",", ":")
    ).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def ensure_launch(source, selector, jobs, hardware_policy):
    path = os.path.join(RESULT_ROOT, "r2_launch.json")
    expected = {
        "stage": "r2",
        "source": source,
        "selector": selector,
        "registry_sha256": registry_sha256(jobs),
        "num_jobs": len(jobs),
        "hardware_policy": hardware_policy,
        "allowed_physical_gpus":
            hardware_policy["allowed_physical_gpu_indices"],
    }
    if os.path.isfile(path):
        if read_json(path) != expected:
            raise RuntimeError("ResNet R2 launch lock drift")
        return path
    os.makedirs(os.path.dirname(path), 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 assert_source_unchanged(source):
    if git_output("status", "--porcelain", "--untracked-files=no"):
        raise RuntimeError("tracked source changed after ResNet R2 launch")
    if git_output("rev-parse", "HEAD") != source["git_commit"]:
        raise RuntimeError("source commit changed after ResNet R2 launch")
    for relative, expected_hash in source["tracked_files"].items():
        if sha256(os.path.join(ROOT, relative)) != expected_hash:
            raise RuntimeError(f"ResNet R2 source drift: {relative}")


def run_job(job, source, selector, gpu, dry_run):
    if not dry_run:
        assert_source_unchanged(source)
        if sha256(os.path.join(ROOT, selector["path"])) != selector["sha256"]:
            raise RuntimeError("ResNet R2 selector changed after launch")
    manifest_path = job["output"] + ".manifest.json"
    if os.path.exists(manifest_path):
        print(f"preserving {job['experiment_name']}", flush=True)
        return
    if os.path.exists(job["output"]):
        raise RuntimeError(f"orphaned R2 output requires audit: {job['output']}")
    print("RUN", " ".join(job["command"]), flush=True)
    if dry_run:
        return
    os.makedirs(os.path.dirname(job["output"]), exist_ok=True)
    started = time.time()
    try:
        result = subprocess.run(
            job["command"], cwd=ROOT, timeout=job["timeout_seconds"]
        )
        return_code = result.returncode
        status = "completed" if return_code == 0 else "nonzero_exit"
    except subprocess.TimeoutExpired:
        return_code = None
        status = "timeout"
    output_exists = os.path.isfile(job["output"])
    if status == "completed" and not output_exists:
        status = "missing_output"
    manifest = {
        **job,
        "source": source,
        "selector": selector,
        "hardware_lock": gpu,
        "status": status,
        "return_code": return_code,
        "output_exists": output_exists,
        "output_sha256": sha256(job["output"]) if output_exists else None,
        "driver_wall_seconds": time.time() - started,
        "completed_unix_time": time.time(),
        "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
    }
    with open(manifest_path, "w", encoding="utf-8") as handle:
        json.dump(manifest, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(f"DONE status={status} {job['experiment_name']}", flush=True)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--selector", default=DEFAULT_SELECTOR)
    parser.add_argument("--shard-index", type=int, default=0)
    parser.add_argument("--num-shards", type=int, default=1)
    parser.add_argument("--method", choices=METHODS)
    parser.add_argument("--depth", type=int, choices=DEPTHS)
    parser.add_argument(
        "--hardware-profile",
        choices=profile_choices(),
        default=os.environ.get("SDIL_HARDWARE_PROFILE", DEFAULT_PROFILE),
    )
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    if not 0 <= args.shard_index < args.num_shards:
        raise ValueError("invalid ResNet R2 shard")
    selector = selector_report(args.selector)
    source = source_report(args.selector)
    gpu = physical_gpu_report(args.hardware_profile, args.dry_run)
    complete_jobs = r2_jobs(selector["selected_rates"])
    if not args.dry_run:
        hardware_policy = policy_for_report(args.hardware_profile, gpu)
        launch = ensure_launch(
            source, selector, complete_jobs, hardware_policy
        )
        print(f"R2 launch lock: {launch}", flush=True)
    jobs = complete_jobs
    if args.method:
        jobs = [job for job in jobs if job["method"] == args.method]
    if args.depth is not None:
        jobs = [job for job in jobs if job["depth"] == args.depth]
    jobs = [
        job
        for index, job in enumerate(jobs)
        if index % args.num_shards == args.shard_index
    ]
    if not jobs:
        raise ValueError("no ResNet R2 jobs selected")
    for job in jobs:
        run_job(job, source, selector, gpu, args.dry_run)


if __name__ == "__main__":
    main()