#!/usr/bin/env python3
from __future__ import annotations

import argparse
import re
import shlex
import shutil
import subprocess
import sys
from pathlib import Path


REPO = Path(__file__).resolve().parents[2]
SUITE = REPO / "validation/networks/pimcomp_models"
PIMCOMP_SOURCE = REPO / "third_party/PIMCOMP-NN"
COMPARE = REPO / "validation/tools/compare_raptor_pimcomp.py"
MODELS = {
    "vgg8": SUITE / "vgg8/vgg8-mnist-reconstructed.onnx",
    "resnet18": SUITE / "resnet18/resnet18-v1-7.onnx",
    "resnet34": SUITE / "resnet34/resnet34-v1-7.onnx",
    "googlenet": SUITE / "googlenet/googlenet-12-no-softmax.onnx",
}


def run(command: list[str], *, dry_run: bool, check: bool = True) -> int:
    print(f"$ {shlex.join(command)}", flush=True)
    if dry_run:
        return 0
    return subprocess.run(command, cwd=REPO, check=check).returncode


def prepare_pimcomp(work_dir: Path) -> None:
    shutil.copytree(
        PIMCOMP_SOURCE,
        work_dir,
        dirs_exist_ok=True,
        ignore=shutil.ignore_patterns(".git", "build", "output"),
    )
    header = work_dir / "backend/GeneticAlgorithm.h"
    source = header.read_text(encoding="utf-8")
    if "int population_num = 200;" not in source:
        raise RuntimeError("PIMCOMP GA population is not 200")
    source, replacements = re.subn(
        r"int max_iteration = \d+;",
        "int max_iteration = 1000;",
        source,
    )
    if replacements != 1:
        raise RuntimeError("Could not set PIMCOMP GA max_iteration")
    header.write_text(source, encoding="utf-8")
    shutil.copy2(SUITE / "configs/arch-a.json", work_dir / "config.json")


def comparison_command(model: Path, result_dir: Path, pimcomp_dir: Path, timeout: float) -> list[str]:
    return [
        sys.executable,
        str(COMPARE),
        "--model",
        str(model),
        "--out-dir",
        str(result_dir),
        "--pimcomp-dir",
        str(pimcomp_dir),
        "--core-count",
        "168",
        "--crossbar-count",
        "96",
        "--crossbar-size",
        "128",
        "--mesh-rows",
        "12",
        "--mesh-cols",
        "14",
        "--pimsim-mode",
        "latency",
        "--pimcomp-pipeline",
        "element",
        "--pimcomp-replication",
        "GA",
        "--timeout-seconds",
        str(timeout),
        "--fail-on-error",
    ]


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Reproduce the serial Arch-A latency comparison from the PIMCOMP paper."
    )
    parser.add_argument("--out-dir", required=True, type=Path)
    parser.add_argument("--models", nargs="+", choices=MODELS, default=list(MODELS))
    parser.add_argument("--timeout-seconds", type=float, default=3600.0)
    parser.add_argument(
        "--resume",
        action="store_true",
        help="Keep the existing work tree and skip models with a completed JSON report.",
    )
    parser.add_argument("--dry-run", action="store_true", help="Print commands without modifying files.")
    args = parser.parse_args()

    out_dir = args.out_dir.resolve()
    work_dir = out_dir / "pimcomp-ga1000"
    if not args.dry_run and out_dir.exists() and any(out_dir.iterdir()) and not args.resume:
        parser.error(f"{out_dir} is not empty; choose a fresh directory or pass --resume")

    missing = [str(MODELS[name]) for name in args.models if not MODELS[name].exists()]
    if missing:
        parser.error(f"missing model(s): {', '.join(missing)}")

    if args.dry_run:
        print(f"# prepare isolated PIMCOMP GA build in {work_dir}")
    else:
        out_dir.mkdir(parents=True, exist_ok=True)
        prepare_pimcomp(work_dir)

    run(["cmake", "--build", str(REPO / "build_release")], dry_run=args.dry_run)
    run(
        ["cmake", "-S", str(work_dir), "-B", str(work_dir / "build")],
        dry_run=args.dry_run,
    )
    run(
        ["cmake", "--build", str(work_dir / "build"), "--target", "PIMCOMP-NN"],
        dry_run=args.dry_run,
    )

    failed = []
    for name in args.models:
        result_dir = out_dir / name
        if args.resume and (result_dir / "comparison_report.json").exists():
            print(f"[{name}] completed report exists; skipping", flush=True)
            continue
        print(f"\n[{name}] Arch-A latency comparison", flush=True)
        returncode = run(
            comparison_command(MODELS[name], result_dir, work_dir, args.timeout_seconds),
            dry_run=args.dry_run,
            check=False,
        )
        if returncode:
            failed.append(name)

    if failed:
        print(f"\nCompleted with failed comparisons: {', '.join(failed)}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())