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

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

from colorama import Fore, Style


REPO = Path(__file__).resolve().parents[3]
SUITE = REPO / "validation/networks/pimcomp_models"
sys.path.insert(0, str(REPO / "validation"))

from raptor_validation.pimsim_nn import parse_pimsim_nn_metrics  # noqa: E402
from raptor_validation.validate_one import STAGE_COLORS  # noqa: E402

PIMCOMP_SOURCE = REPO / "third_party/PIMCOMP-NN"
PIMCOMP_CONFIGS = REPO / "validation/pimsim_configs/pimcomp"
COMPARE = Path(__file__).resolve().with_name("compare_raptor_pimcomp.py")
ARCHES = tuple(sorted(path.name for path in PIMCOMP_CONFIGS.iterdir() if path.is_dir()))
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-latency.onnx",
}
COMPARISONS = (
    ("latency", 1, "element"),
    ("throughput", 2, "batch"),
    ("throughput", 4, "batch"),
    ("throughput", 8, "batch"),
)


def result_dir(root: Path | None, name: str, mode: str, pipeline: int) -> Path:
    base = root / name if root is not None else MODELS[name].parent
    suffix = "latency" if mode == "latency" else f"throughput/pipeline{pipeline}"
    return base / suffix


def write_results_csv(root: Path | None, arch: str, models: list[str]) -> Path:
    output = (root or SUITE) / "results.csv"
    fields = (
        "model",
        "arch",
        "mode",
        "raptor_pipeline",
        "pimcomp_pipeline",
        "status",
        "raptor_throughput_samples_s",
        "pimcomp_throughput_samples_s",
        "raptor_latency_ms",
        "pimcomp_latency_ms",
        "raptor_power_mw",
        "pimcomp_power_mw",
        "raptor_energy_pj",
        "pimcomp_energy_pj",
        "better_compiler",
        "speedup",
    )
    rows = []
    for name in models:
        for mode, pipeline, pimcomp_pipeline in COMPARISONS:
            report_path = result_dir(root, name, mode, pipeline) / "pimcomp/comparison_report.json"
            if not report_path.exists():
                continue
            report = json.loads(report_path.read_text(encoding="utf-8"))
            raptor = report.get("raptor_performance") or {}
            pimcomp = report.get("pimcomp_performance") or {}
            raptor_values = performance_values(raptor)
            pimcomp_values = performance_values(pimcomp)
            raptor_metric = raptor_values["throughput"] if mode == "throughput" else raptor_values["latency"]
            pimcomp_metric = pimcomp_values["throughput"] if mode == "throughput" else pimcomp_values["latency"]
            status = "PASS" if comparison_passed(report) else "FAIL"
            if raptor_metric is None or pimcomp_metric is None:
                better = ""
                speedup = ""
            else:
                better = comparison_winner(mode, raptor_metric, pimcomp_metric)
                speedup = f"{max(raptor_metric, pimcomp_metric) / min(raptor_metric, pimcomp_metric):.2f}"
            rows.append({
                "model": name,
                "arch": arch,
                "mode": mode,
                "raptor_pipeline": pipeline,
                "pimcomp_pipeline": pimcomp_pipeline,
                "status": status,
                "raptor_throughput_samples_s": format_value(raptor_values["throughput"]),
                "pimcomp_throughput_samples_s": format_value(pimcomp_values["throughput"]),
                "raptor_latency_ms": format_value(raptor_values["latency"]),
                "pimcomp_latency_ms": format_value(pimcomp_values["latency"]),
                "raptor_power_mw": format_value(raptor_values["power"]),
                "pimcomp_power_mw": format_value(pimcomp_values["power"]),
                "raptor_energy_pj": format_value(raptor_values["energy"]),
                "pimcomp_energy_pj": format_value(pimcomp_values["energy"]),
                "better_compiler": better,
                "speedup": speedup,
            })
    with open(output, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fields, lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)
    return output


def performance_values(performance: dict) -> dict[str, float | None]:
    parsed = parse_pimsim_nn_metrics(performance.get("raw_output", ""))
    return {
        "throughput": performance.get("throughput") or parsed.get("throughput"),
        "latency": (
            performance.get("latency_ms")
            or performance.get("average_latency_ms")
            or parsed.get("latency_ms")
            or parsed.get("average_latency_ms")
        ),
        "power": performance.get("average_power_mw") or parsed.get("average_power_mw"),
        "energy": performance.get("average_energy_pj") or parsed.get("average_energy_pj"),
    }


def comparison_passed(report: dict) -> bool:
    if report.get("failures"):
        return False
    for key in ("raptor_validation", "pimcomp_validation"):
        result = report.get(key) or {}
        if result.get("status") != "done" or not result.get("passed"):
            return False
    for key in ("raptor_performance", "pimcomp_performance"):
        performance = report.get(key) or {}
        if performance.get("error") or performance.get("skipped"):
            return False
    return True


def comparison_winner(mode: str, raptor: float, pimcomp: float) -> str:
    if raptor == pimcomp:
        return "tie"
    if mode == "throughput":
        return "raptor" if raptor > pimcomp else "pimcomp"
    return "raptor" if raptor < pimcomp else "pimcomp"


def format_value(value: float | None) -> str:
    return "" if value is None else f"{value:.6f}"


def print_stage(title: str, color: str) -> None:
    print("\n" + Style.BRIGHT + color + f"[{title}]" + Style.RESET_ALL, flush=True)


def run(command: list[str], *, dry_run: bool, check: bool = True) -> int:
    print(f"  cwd: {REPO}", flush=True)
    print(f"  $ {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,
    config: Path,
    mode: str,
    pipeline: int,
    pimcomp_pipeline: str,
    pimsim_time_ms: int,
    timeout: float,
) -> list[str]:
    time_args = ["--pimsim-time-ms", str(pimsim_time_ms)] if mode == "throughput" else []
    return [
        sys.executable,
        str(COMPARE),
        "--model",
        str(model),
        "--out-dir",
        str(result_dir),
        "--pimcomp-dir",
        str(PIMCOMP_SOURCE),
        "--pimcomp-config",
        str(config),
        "--pimsim-mode",
        mode,
        *time_args,
        "--pimcomp-pipeline",
        pimcomp_pipeline,
        "--pimcomp-replication",
        "GA",
        f"--raptor-extra-arg=--pipeline={pipeline}",
        "--timeout-seconds",
        str(timeout),
        "--fail-on-error",
    ]


def config_path(arch: str, mode: str) -> Path:
    path = PIMCOMP_CONFIGS / arch / f"{mode}_config.json"
    if not path.exists():
        raise ValueError(f"{arch} has no {mode} config: {path}")
    return path


def core_count(config: Path) -> int:
    with open(config, encoding="utf-8") as f:
        return int(json.load(f)["chip_config"]["core_cnt"])


def completed_report(path: Path, mode: str, pipeline: int, config: Path, pimsim_time_ms: int) -> bool:
    if not path.exists():
        return False
    report = json.loads(path.read_text(encoding="utf-8"))
    return (
        report.get("pimsim_mode") == mode
        and report.get("pimcomp_pipeline") == ("element" if mode == "latency" else "batch")
        and report.get("pimsim_time_ms") == pimsim_time_ms
        and report.get("pimcomp_config") == str(config.resolve())
        and f"--pipeline={pipeline}" in report.get("raptor_extra_args", [])
    )


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Compare supported PIMCOMP models with Raptor latency and throughput schedules."
    )
    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(
        "--arch", choices=ARCHES, default="arch-a", help="PIM architecture (default: arch-a)."
    )
    parser.add_argument(
        "--pimsim-time-ms",
        type=int,
        default=100,
        help="throughput pimsim-nn horizon in ms (default: 100).",
    )
    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()

    if args.pimsim_time_ms <= 0:
        parser.error("--pimsim-time-ms must be positive")
    configs = {mode: config_path(args.arch, mode) for mode, _, _ in COMPARISONS}
    unsupported = [pipeline for mode, pipeline, _ in COMPARISONS if core_count(configs[mode]) % pipeline]
    if unsupported:
        parser.error(
            f"{args.arch} has {core_count(configs['throughput'])} cores; "
            f"throughput pipelines must divide that count (invalid: {unsupported})"
        )

    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)

    print(Style.BRIGHT + f"Found {len(args.models)} PIMCOMP model(s) to compare." + Style.RESET_ALL)
    print(f"Architecture: {args.arch}")
    print(f"Throughput pimsim time: {args.pimsim_time_ms} ms")
    print(f"Results root: {out_dir or SUITE}")
    print("=" * 72)

    print_stage("Build Raptor", STAGE_COLORS["Build Runner"])
    run(["cmake", "--build", str(REPO / "build_release")], dry_run=args.dry_run)

    print_stage("Build PIMCOMP", STAGE_COLORS["Build Runner"])
    run(
        ["cmake", "--build", str(PIMCOMP_SOURCE / "build"), "--target", "PIMCOMP-NN"],
        dry_run=args.dry_run,
    )

    failed = []
    for index, name in enumerate(args.models, start=1):
        for mode, pipeline, pimcomp_pipeline in COMPARISONS:
            model_result_dir = result_dir(out_dir, name, mode, pipeline)
            print(
                "\n" + Fore.CYAN + f"[{index}/{len(args.models)}]" + Style.RESET_ALL
                + f" {Style.BRIGHT}Comparing {name} ({mode}, pipeline={pipeline}){Style.RESET_ALL}",
                flush=True,
            )
            if args.resume and completed_report(
                model_result_dir / "pimcomp/comparison_report.json",
                mode,
                pipeline,
                configs[mode],
                args.pimsim_time_ms,
            ):
                print(
                    Fore.YELLOW + "  Completed report exists; skipping" + Style.RESET_ALL,
                    flush=True,
                )
                continue
            returncode = run(
                comparison_command(
                    MODELS[name],
                    model_result_dir,
                    configs[mode],
                    mode,
                    pipeline,
                    pimcomp_pipeline,
                    args.pimsim_time_ms,
                    args.timeout_seconds,
                ),
                dry_run=args.dry_run,
                check=False,
            )
            if returncode:
                failed.append(f"{name}/{mode}/pipeline{pipeline}")

    if args.dry_run:
        return 1 if failed else 0

    results_path = write_results_csv(out_dir, args.arch, args.models)
    print_stage("Results", STAGE_COLORS["Compare Outputs"])
    print(results_path.read_text(encoding="utf-8"), end="")
    print("\n" + Style.BRIGHT + Fore.CYAN + "Summary" + Style.RESET_ALL)
    total_jobs = len(args.models) * len(COMPARISONS)
    print(Style.BRIGHT + f"Passed: {total_jobs - len(failed)}" + Style.RESET_ALL)
    print(Style.BRIGHT + f"Failed: {len(failed)}" + Style.RESET_ALL)
    print(Style.BRIGHT + f"Results: {results_path}" + Style.RESET_ALL)
    if failed:
        print(
            Fore.RED + f"Failed comparisons: {', '.join(failed)}" + Style.RESET_ALL,
            file=sys.stderr,
        )
        return 1
    return 0


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