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

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

from colorama import Fore, Style


REPO = Path(__file__).resolve().parents[2]
SUITE = REPO / "validation/networks/pimcomp_models"
PIMCOMP_SOURCE = REPO / "third_party/PIMCOMP-NN"
PIMCOMP_CONFIG = REPO / "validation/pimsim_configs/pimcomp/arch-a/latency_config.json"
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(Fore.CYAN + "$ " + Style.RESET_ALL + shlex.join(command), flush=True)
    if dry_run:
        return 0
    return subprocess.run(command, cwd=REPO, check=check).returncode


def validate_pimcomp_source() -> None:
    header = PIMCOMP_SOURCE / "backend/GeneticAlgorithm.h"
    source = header.read_text(encoding="utf-8")
    for setting in ("int population_num = 200;", "int max_iteration = 1000;"):
        if setting not in source:
            raise RuntimeError(f"PIMCOMP paper setting is missing: {setting}")


def comparison_command(model: Path, result_dir: Path, timeout: float) -> list[str]:
    return [
        sys.executable,
        str(COMPARE),
        "--model",
        str(model),
        "--out-dir",
        str(result_dir),
        "--pimcomp-dir",
        str(PIMCOMP_SOURCE),
        "--pimcomp-config",
        str(PIMCOMP_CONFIG),
        "--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",
        type=Path,
        help="Result root (default: artifacts beside each model under validation/).",
    )
    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="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() if args.out_dir is not None else None

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

    validate_pimcomp_source()
    if out_dir is not None and not args.dry_run:
        out_dir.mkdir(parents=True, exist_ok=True)

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

    failed = []
    for name in args.models:
        result_dir = out_dir / name if out_dir is not None else MODELS[name].parent
        if args.resume and (result_dir / "pimcomp/comparison_report.json").exists():
            print(
                Fore.YELLOW + f"[{name}] completed report exists; skipping" + Style.RESET_ALL,
                flush=True,
            )
            continue
        print(
            "\n" + Fore.CYAN + f"[{name}]" + Style.RESET_ALL
            + f" {Style.BRIGHT}Arch-A latency comparison{Style.RESET_ALL}",
            flush=True,
        )
        returncode = run(
            comparison_command(MODELS[name], result_dir, args.timeout_seconds),
            dry_run=args.dry_run,
            check=False,
        )
        if returncode:
            failed.append(name)

    if failed:
        print(
            "\n" + Style.BRIGHT + Fore.RED + "Result: FAIL" + Style.RESET_ALL,
            file=sys.stderr,
        )
        print(
            Fore.RED + f"Failed comparisons: {', '.join(failed)}" + Style.RESET_ALL,
            file=sys.stderr,
        )
        return 1
    if not args.dry_run:
        print("\n" + Style.BRIGHT + f"Result: {Fore.GREEN}PASS" + Style.RESET_ALL)
    return 0


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