#!/usr/bin/env python3
"""Run the synchronization and Spatial-planning ablation matrix on Pimcomp models."""

from __future__ import annotations

import argparse
import csv
import math
import os
import shlex
import signal
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path


REPO = Path(__file__).resolve().parents[4]
RUNNER = REPO / "validation/tools/pim/pimcomp/compare/run_pimcomp_models.py"
DEFAULT_OUT_DIR = REPO / "validation/networks/pimcomp_models"
sys.path.insert(0, str(REPO / "validation"))

from raptor_validation.pimcomp_models import (
    ABLATION_DEFAULT_MODELS,
    add_models_argument,
)
from raptor_validation.artifacts import remove_lock_files


@dataclass(frozen=True)
class Variant:
    name: str
    raptor_args: tuple[str, ...]


VARIANTS = (
    Variant("baseline", ()),
    Variant("no-synchronization", ("--pim-disable-synchronization",)),
    Variant("no-spatial-planning", ("--pim-disable-spatial-planning",)),
    Variant(
        "no-synchronization-no-spatial-planning",
        ("--pim-disable-synchronization", "--pim-disable-spatial-planning"),
    ),
)
VARIANT_BY_NAME = {variant.name: variant for variant in VARIANTS}
REFERENCE_VARIANT = "no-synchronization-no-spatial-planning"
COMPARISON_RESULTS_FILENAME = "results_comparison.csv"
ABLATION_RESULTS_FILENAME = "results_ablation.csv"
CASE_FIELDS = ("arch", "model", "mode", "raptor_pipeline")
PERCENTAGE_FIELDS = (
    ("raptor_latency_ms", "latency_percent"),
    ("raptor_throughput_samples_s", "throughput_percent"),
    ("raptor_power_mw", "power_percent"),
    ("raptor_energy_pj", "energy_percent"),
)
RESULT_FIELDS = (*CASE_FIELDS, *(target for _, target in PERCENTAGE_FIELDS))
DEFAULT_CASE_ARGUMENTS = (
    ("latency", ("--mode", "latency")),
    ("throughput/pipeline4", ("--mode", "throughput", "--pipeline", "4")),
)
DEFAULT_CASE_KEYS = frozenset({("latency", "1"), ("throughput", "4")})


def parse_args(argv: list[str] | None = None) -> tuple[argparse.Namespace, list[str]]:
    parser = argparse.ArgumentParser(
        description="Run the complete compiler ablation matrix on the Pimcomp model suite.",
        allow_abbrev=False,
    )
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=DEFAULT_OUT_DIR,
        help="Artifact root (default: validation/networks/pimcomp_models).",
    )
    add_models_argument(parser, ABLATION_DEFAULT_MODELS)
    parser.add_argument(
        "--variant",
        choices=tuple(VARIANT_BY_NAME),
        action="append",
        dest="variants",
        help="Run only this variant; baseline is added when needed. Repeat as needed.",
    )
    parser.add_argument("--dry-run", action="store_true", help="Print runner commands without modifying files.")
    args, forwarded = parser.parse_known_args(argv)
    if any(option == "--only" or option.startswith("--only=") for option in forwarded):
        parser.error("--only is not supported; the ablation wrapper always runs Raptor only")
    return args, forwarded


def selected_variants(names: list[str] | None) -> list[Variant]:
    requested = set(names or VARIANT_BY_NAME)
    requested.add("baseline")
    requested.add(REFERENCE_VARIANT)
    return [variant for variant in VARIANTS if variant.name in requested]


def has_option(arguments: list[str], option: str) -> bool:
    return any(argument == option or argument.startswith(f"{option}=") for argument in arguments)


def runner_argument_sets(forwarded: list[str]) -> tuple[tuple[str, list[str]], ...]:
    if has_option(forwarded, "--mode") or has_option(forwarded, "--pipeline"):
        return (("requested", forwarded),)
    return tuple(
        (label, [*forwarded, *case_arguments])
        for label, case_arguments in DEFAULT_CASE_ARGUMENTS
    )


def variant_output_dir(out_dir: Path, variant: Variant) -> Path:
    return out_dir if variant.name == "baseline" else out_dir / "ablation" / variant.name


def remove_variant_summaries(out_dir: Path) -> None:
    summary_root = out_dir / "ablation"
    if not summary_root.is_dir():
        return
    for summary in summary_root.glob("*/results_comparison.csv"):
        summary.unlink()
    for variant_dir in summary_root.iterdir():
        if variant_dir.is_dir() and not any(variant_dir.iterdir()):
            variant_dir.rmdir()
    if not any(summary_root.iterdir()):
        summary_root.rmdir()


def runner_command(
    out_dir: Path,
    variant: Variant,
    models: list[str],
    forwarded: list[str],
    common_root: Path,
    dry_run: bool,
) -> list[str]:
    command = [
        sys.executable,
        str(RUNNER),
        "--out-dir",
        str(out_dir),
        "--raptor-only",
        "--models",
        *models,
    ]
    if variant.name != "baseline":
        command.extend(("--ablation-variant", variant.name))
    if not has_option(forwarded, "--common-dir"):
        command.extend(("--common-dir", str(common_root)))
    command.extend(forwarded)
    command.extend(f"--raptor-extra-arg={arg}" for arg in variant.raptor_args)
    if dry_run:
        command.append("--dry-run")
    return command


def terminate_process_group(process: subprocess.Popen[bytes]) -> None:
    if process.poll() is not None:
        return
    try:
        os.killpg(process.pid, signal.SIGTERM)
    except ProcessLookupError:
        return
    try:
        process.wait(timeout=5)
    except subprocess.TimeoutExpired:
        try:
            os.killpg(process.pid, signal.SIGKILL)
        except ProcessLookupError:
            return
        process.wait()


def percentage_difference(value: str | None, reference: str | None) -> str:
    try:
        current = float(value) if value is not None else math.nan
        baseline = float(reference) if reference is not None else math.nan
    except ValueError:
        return "NA"
    if not math.isfinite(current) or not math.isfinite(baseline) or baseline == 0:
        return "NA"
    return f"{(current - baseline) / baseline * 100:+.2f}%"


def aggregate_results(
    out_dir: Path,
    variants: list[Variant],
    selected_cases: frozenset[tuple[str, str]] | None = None,
    selected_models: frozenset[str] | None = None,
) -> tuple[Path | None, list[str]]:
    fields: list[str] | None = None
    rows_by_variant: dict[str, list[dict[str, str]]] = {}
    failures = []
    for variant in variants:
        results_path = variant_output_dir(out_dir, variant) / COMPARISON_RESULTS_FILENAME
        if not results_path.is_file():
            failures.append(f"{variant.name}: missing {results_path}")
            continue
        with results_path.open(newline="", encoding="utf-8") as stream:
            reader = csv.DictReader(stream)
            if reader.fieldnames is None:
                failures.append(f"{variant.name}: empty {results_path}")
                continue
            if fields is None:
                fields = reader.fieldnames
            elif reader.fieldnames != fields:
                failures.append(f"{variant.name}: inconsistent columns in {results_path}")
                continue
            selected_rows = rows_by_variant.setdefault(variant.name, [])
            for row in reader:
                if selected_models is not None and row.get("model") not in selected_models:
                    continue
                if selected_cases is not None and (
                    row.get("mode"), row.get("raptor_pipeline")
                ) not in selected_cases:
                    continue
                selected_rows.append(row)
    if fields is None:
        return None, failures
    reference_rows = {
        tuple(row.get(field, "") for field in CASE_FIELDS): row
        for row in rows_by_variant.get(REFERENCE_VARIANT, [])
    }
    if REFERENCE_VARIANT not in rows_by_variant:
        failures.append(f"{REFERENCE_VARIANT}: missing percentage reference results")
    output = out_dir / ABLATION_RESULTS_FILENAME
    with output.open("w", newline="", encoding="utf-8") as stream:
        writer = csv.DictWriter(stream, fieldnames=("variant", *RESULT_FIELDS), lineterminator="\n")
        writer.writeheader()
        for variant in variants:
            for row in rows_by_variant.get(variant.name, []):
                reference = reference_rows.get(tuple(row.get(field, "") for field in CASE_FIELDS))
                writer.writerow(
                    {
                        "variant": variant.name,
                        **{field: row.get(field, "") for field in CASE_FIELDS},
                        **{
                            target: percentage_difference(
                                row.get(source), reference.get(source) if reference else None
                            )
                            for source, target in PERCENTAGE_FIELDS
                        },
                    }
                )
    return output, failures


def main(argv: list[str] | None = None) -> int:
    args, forwarded = parse_args(argv)
    if not RUNNER.is_file():
        print(f"Missing Pimcomp runner: {RUNNER}", file=sys.stderr)
        return 1
    variants = selected_variants(args.variants)
    out_dir = args.out_dir.resolve()
    case_sets = runner_argument_sets(forwarded)
    if args.dry_run:
        for variant in variants:
            for _, case_forwarded in case_sets:
                print(
                    shlex.join(
                        runner_command(
                            out_dir,
                            variant,
                            args.models,
                            case_forwarded,
                            out_dir,
                            True,
                        )
                    )
                )
        return 0

    out_dir.mkdir(parents=True, exist_ok=True)
    failed = []
    current_process = None
    try:
        for variant in variants:
            for case_label, case_forwarded in case_sets:
                command = runner_command(
                    out_dir, variant, args.models, case_forwarded, out_dir, False
                )
                print(f"[{variant.name}/{case_label}] {shlex.join(command)}")
                current_process = subprocess.Popen(
                    command,
                    cwd=REPO,
                    start_new_session=True,
                )
                returncode = current_process.wait()
                current_process = None
                if returncode:
                    failed.append(f"{variant.name}/{case_label}: runner exited with {returncode}")
    except KeyboardInterrupt:
        if current_process is not None:
            terminate_process_group(current_process)
        remove_lock_files(out_dir)
        remove_variant_summaries(out_dir)
        print("Interrupted; terminated the active ablation job.", file=sys.stderr)
        return 130

    remove_lock_files(out_dir)
    selected_cases = DEFAULT_CASE_KEYS if len(case_sets) > 1 else None
    output, aggregation_failures = aggregate_results(
        out_dir,
        variants,
        selected_cases,
        frozenset(args.models),
    )
    remove_variant_summaries(out_dir)
    failed.extend(aggregation_failures)
    if output is not None:
        print(f"Ablation results: {output}")
    if failed:
        print("Failed: " + "; ".join(failed), file=sys.stderr)
        return 1
    return 0


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