#!/usr/bin/env python3

from __future__ import annotations

import argparse
import csv
import datetime as dt
import json
import os
import pathlib
import platform
import random
import shutil
import statistics
import subprocess
import sys
import time
from dataclasses import dataclass
from typing import Any


JSON_PREFIX = "MFEM_BENCHMARK_JSON "


@dataclass(frozen=True)
class Case:
    name: str
    device: str
    ranks: int
    threads: int


def positive_int(value: str) -> int:
    parsed = int(value)
    if parsed < 1:
        raise argparse.ArgumentTypeError("must be at least one")
    return parsed


def integer_list(value: str) -> list[int]:
    try:
        parsed = [int(item.strip()) for item in value.split(",") if item.strip()]
    except ValueError as error:
        raise argparse.ArgumentTypeError("expected comma-separated integers") from error
    if not parsed or any(item < 1 for item in parsed):
        raise argparse.ArgumentTypeError("all entries must be positive")
    return parsed


def command_output(command: list[str]) -> str | None:
    try:
        result = subprocess.run(
            command,
            check=False,
            capture_output=True,
            text=True,
            timeout=10,
        )
    except (OSError, subprocess.TimeoutExpired):
        return None
    if result.returncode != 0:
        return None
    output = result.stdout.strip()
    return output if output else None


def resolve_program(value: str) -> str | None:
    candidate = pathlib.Path(value)
    if candidate.is_file():
        return str(candidate.resolve())
    discovered = shutil.which(value)
    return str(pathlib.Path(discovered).resolve()) if discovered else None


def compact_cpu_list(cpus: list[int]) -> str:
    if not cpus:
        return ""
    groups: list[str] = []
    start = previous = cpus[0]
    for current in cpus[1:]:
        if current == previous + 1:
            previous = current
            continue
        groups.append(str(start) if start == previous else f"{start}-{previous}")
        start = previous = current
    groups.append(str(start) if start == previous else f"{start}-{previous}")
    return ",".join(groups)


def affinity_prefix(case: Case, setting: str) -> tuple[list[str], str | None]:
    if setting == "none" or shutil.which("taskset") is None:
        return [], None
    if setting != "auto":
        return ["taskset", "-c", setting], setting
    if not hasattr(os, "sched_getaffinity"):
        return [], None
    available = sorted(os.sched_getaffinity(0))
    count = min(len(available), max(case.ranks, case.threads))
    chosen = compact_cpu_list(available[:count])
    return (["taskset", "-c", chosen] if chosen else []), chosen or None


def expand_cases(args: argparse.Namespace) -> list[Case]:
    requested = [item.strip() for item in args.cases.split(",") if item.strip()]
    valid = {"cpu", "ceed-cpu", "omp", "mpi", "cuda", "ceed-cuda"}
    unknown = sorted(set(requested) - valid)
    if unknown:
        raise ValueError(f"unknown benchmark case(s): {', '.join(unknown)}")
    cases: list[Case] = []
    for requested_case in requested:
        if requested_case == "cpu":
            cases.append(Case("cpu-1r-1t", "cpu", 1, 1))
        elif requested_case == "ceed-cpu":
            cases.append(Case("ceed-cpu-1r-1t", args.ceed_cpu_device, 1, 1))
        elif requested_case == "omp":
            cases.extend(
                Case(f"omp-1r-{threads}t", "omp", 1, threads)
                for threads in args.omp_threads
            )
        elif requested_case == "mpi":
            cases.extend(
                Case(f"mpi-{ranks}r-1t", "cpu", ranks, 1)
                for ranks in args.mpi_ranks
            )
        elif requested_case == "cuda":
            cases.append(Case("cuda-1r-1t", "cuda", 1, 1))
        elif requested_case == "ceed-cuda":
            cases.append(Case("ceed-cuda-1r-1t", args.ceed_cuda_device, 1, 1))
    deduplicated: list[Case] = []
    for case in cases:
        if case not in deduplicated:
            deduplicated.append(case)
    return deduplicated


def run_case(
    args: argparse.Namespace, case: Case, trial: int, warmup_trial: bool
) -> dict[str, Any]:
    prefix, affinity = affinity_prefix(case, args.affinity)
    command = prefix + [args.launcher, "-n", str(case.ranks)]
    command.extend(args.launcher_arg)
    command.extend(
        [
            args.executable,
            "--device",
            case.device,
            "--mesh-n",
            str(args.mesh_n),
            "--order",
            str(args.order),
            "--applications",
            str(args.applications),
            "--max-applications",
            str(args.max_applications),
            "--minimum-apply-seconds",
            str(args.minimum_apply_seconds),
            "--warmup-applications",
            str(args.warmup_applications),
            "--relative-tolerance",
            str(args.relative_tolerance),
            "--max-iterations",
            str(args.max_iterations),
            "--trial",
            str(trial),
        ]
    )
    if args.solve:
        command.append("--solve")
    environment = os.environ.copy()
    environment.update(
        {
            "OMP_NUM_THREADS": str(case.threads),
            "OMP_DYNAMIC": "FALSE",
            "OMP_PROC_BIND": "close",
            "OMP_PLACES": "cores",
            "FI_PROVIDER": environment.get("FI_PROVIDER", "sockets"),
        }
    )
    environment.setdefault("OMPI_ALLOW_RUN_AS_ROOT", "1")
    environment.setdefault("OMPI_ALLOW_RUN_AS_ROOT_CONFIRM", "1")
    selected_interface = args.network_interface
    if selected_interface == "local":
        selected_interface = "lo0" if sys.platform == "darwin" else "lo"
    if selected_interface == "auto":
        environment.pop("FI_SOCKETS_IFACE", None)
        environment.pop("HYDRA_IFACE", None)
        environment.pop("MPICH_INTERFACE_HOSTNAME", None)
    else:
        environment["FI_SOCKETS_IFACE"] = selected_interface
        environment["HYDRA_IFACE"] = selected_interface
        environment.pop("MPICH_INTERFACE_HOSTNAME", None)
    if sys.platform == "darwin" and args.network_interface == "local":
        environment["MPICH_INTERFACE_HOSTNAME"] = "127.0.0.1"
    if case.device.startswith("cuda") or case.device.startswith("ceed-cuda"):
        if args.gpu_aware_mpi:
            environment["MFEM_GPU_AWARE_MPI"] = "1"
        else:
            environment.pop("MFEM_GPU_AWARE_MPI", None)
        if args.cuda_visible_devices is not None:
            environment["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices

    started = time.perf_counter()
    completed = subprocess.run(
        command,
        check=False,
        capture_output=True,
        text=True,
        env=environment,
        timeout=args.timeout,
    )
    wall_seconds = time.perf_counter() - started
    combined_output = "\n".join((completed.stdout, completed.stderr))
    records = [
        json.loads(line[len(JSON_PREFIX) :])
        for line in combined_output.splitlines()
        if line.startswith(JSON_PREFIX)
    ]
    if completed.returncode != 0 or len(records) != 1:
        tail = "\n".join(combined_output.splitlines()[-30:])
        raise RuntimeError(
            f"{case.name} trial {trial} failed with exit code "
            f"{completed.returncode}; JSON records={len(records)}\n{tail}"
        )
    record = records[0]
    record.update(
        {
            "case": case.name,
            "requested_device": case.device,
            "runner_trial": trial,
            "discarded_warmup_trial": warmup_trial,
            "wall_seconds": wall_seconds,
            "cpu_affinity": affinity,
            "network_interface": environment.get("HYDRA_IFACE"),
            "cuda_visible_devices": environment.get("CUDA_VISIBLE_DEVICES"),
            "gpu_aware_mpi": environment.get("MFEM_GPU_AWARE_MPI") == "1",
            "command": command,
        }
    )
    if not record.get("valid", False):
        raise RuntimeError(f"{case.name} trial {trial} failed numerical validation")
    return record


def median_absolute_deviation(values: list[float]) -> float:
    center = statistics.median(values)
    return statistics.median(abs(value - center) for value in values)


def validate_cross_backend(
    records: list[dict[str, Any]], relative_tolerance: float
) -> dict[str, float]:
    measured = [row for row in records if not row["discarded_warmup_trial"]]
    baseline = [row for row in measured if row["case"] == "cpu-1r-1t"]
    if not baseline:
        return {}
    validation_fields = ["probe_norm"]
    if measured and all(row.get("solve_ran", False) for row in measured):
        validation_fields.append("solution_norm")
    references = {
        field: statistics.median(float(row[field]) for row in baseline)
        for field in validation_fields
    }
    maximum_errors = {field: 0.0 for field in references}
    for row in measured:
        for field, reference in references.items():
            relative_error = abs(float(row[field]) - reference) / max(
                abs(reference), 1.0e-300
            )
            maximum_errors[field] = max(maximum_errors[field], relative_error)
            if relative_error > relative_tolerance:
                raise RuntimeError(
                    f"{row['case']} disagrees with CPU for {field}: relative "
                    f"error {relative_error:.3e} > {relative_tolerance:.3e}"
                )
    return maximum_errors


def checkpoint_results(
    output_dir: pathlib.Path,
    metadata: dict[str, Any],
    records: list[dict[str, Any]],
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    destination = output_dir / "partial-results.json"
    temporary = output_dir / "partial-results.json.tmp"
    temporary.write_text(
        json.dumps({"schema_version": 1, "metadata": metadata, "raw": records}, indent=2)
        + "\n",
        encoding="utf-8",
    )
    temporary.replace(destination)


def summarize(cases: list[Case], records: list[dict[str, Any]]) -> list[dict[str, Any]]:
    summaries: list[dict[str, Any]] = []
    for case in cases:
        rows = [
            row
            for row in records
            if row["case"] == case.name and not row["discarded_warmup_trial"]
        ]
        if not rows:
            continue
        throughput = [float(row["apply_mdof_per_second"]) for row in rows]
        apply_seconds = [float(row["apply_seconds"]) for row in rows]
        solve_seconds = [float(row["solve_seconds"]) for row in rows]
        total_seconds = [float(row["total_seconds"]) for row in rows]
        timed_applications = [int(row["applications"]) for row in rows]
        summaries.append(
            {
                "case": case.name,
                "device": case.device,
                "ranks": case.ranks,
                "threads": case.threads,
                "trials": len(rows),
                "global_true_dofs": rows[0]["global_true_dofs"],
                "requested_applications": rows[0]["requested_applications"],
                "applications_min": min(timed_applications),
                "applications_max": max(timed_applications),
                "apply_mdof_per_second_median": statistics.median(throughput),
                "apply_mdof_per_second_mad": median_absolute_deviation(throughput),
                "apply_mdof_per_second_min": min(throughput),
                "apply_mdof_per_second_max": max(throughput),
                "apply_seconds_median": statistics.median(apply_seconds),
                "solve_seconds_median": statistics.median(solve_seconds),
                "total_seconds_median": statistics.median(total_seconds),
                "cg_iterations": sorted({row["cg_iterations"] for row in rows}),
                "solve_ran": all(bool(row["solve_ran"]) for row in rows),
                "verified_relative_residual_max": max(
                    float(row["verified_relative_residual"]) for row in rows
                ),
            }
        )
    baseline = next((row for row in summaries if row["case"] == "cpu-1r-1t"), None)
    baseline_rate = (
        float(baseline["apply_mdof_per_second_median"]) if baseline else None
    )
    for row in summaries:
        row["speedup_vs_cpu"] = (
            float(row["apply_mdof_per_second_median"]) / baseline_rate
            if baseline_rate
            else None
        )
    return summaries


def markdown_table(summaries: list[dict[str, Any]]) -> str:
    lines = [
        "| Case | Device | Ranks | Threads | Timed apps | Apply window (s) | Apply MDoF/s, median [range] | MAD | Speedup | Solve (s) | CG iters |",
        "|---|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|",
    ]
    for row in summaries:
        speedup = row["speedup_vs_cpu"]
        speedup_text = f"{speedup:.2f}x" if speedup is not None else "n/a"
        solve_text = f"{row['solve_seconds_median']:.4f}" if row["solve_ran"] else "—"
        iteration_text = (
            ",".join(str(value) for value in row["cg_iterations"])
            if row["solve_ran"]
            else "—"
        )
        applications_text = (
            str(row["applications_min"])
            if row["applications_min"] == row["applications_max"]
            else f"{row['applications_min']}-{row['applications_max']}"
        )
        lines.append(
            f"| {row['case']} | `{row['device']}` | {row['ranks']} | "
            f"{row['threads']} | {applications_text} | "
            f"{row['apply_seconds_median']:.3f} | "
            f"{row['apply_mdof_per_second_median']:.2f} "
            f"[{row['apply_mdof_per_second_min']:.2f}, "
            f"{row['apply_mdof_per_second_max']:.2f}] | "
            f"{row['apply_mdof_per_second_mad']:.2f} | "
            f"{speedup_text} | {solve_text} | {iteration_text} |"
        )
    return "\n".join(lines) + "\n"


def collect_metadata(args: argparse.Namespace) -> dict[str, Any]:
    executable = pathlib.Path(args.executable).resolve()
    source_root = pathlib.Path(__file__).resolve().parents[1]
    return {
        "timestamp_utc": dt.datetime.now(dt.timezone.utc).isoformat(),
        "hostname": platform.node(),
        "platform": platform.platform(),
        "python": sys.version,
        "cpu_count": os.cpu_count(),
        "lscpu": command_output(["lscpu"]) if shutil.which("lscpu") else None,
        "executable": str(executable),
        "launcher": str(pathlib.Path(args.launcher).resolve()),
        "launcher_version": command_output([args.launcher, "--version"]),
        "compiler": command_output(["c++", "--version"]),
        "nvidia_smi": command_output(
            [
                "nvidia-smi",
                "--query-gpu=name,driver_version,temperature.gpu,pstate,clocks.sm,power.draw",
                "--format=csv,noheader",
            ]
        )
        if shutil.which("nvidia-smi")
        else None,
        "git_commit": command_output(
            ["git", "-C", str(source_root), "rev-parse", "HEAD"]
        ),
        "git_dirty": bool(
            command_output(["git", "-C", str(source_root), "status", "--porcelain"])
        ),
        "parameters": {
            "mesh_n": args.mesh_n,
            "order": args.order,
            "applications": args.applications,
            "max_applications": args.max_applications,
            "minimum_apply_seconds": args.minimum_apply_seconds,
            "warmup_applications": args.warmup_applications,
            "trials": args.trials,
            "discarded_warmup_trials": args.warmup_trials,
            "relative_tolerance": args.relative_tolerance,
            "max_iterations": args.max_iterations,
            "solve": args.solve,
            "affinity": args.affinity,
            "network_interface": args.network_interface,
            "cuda_visible_devices": args.cuda_visible_devices,
            "gpu_aware_mpi": args.gpu_aware_mpi,
            "cross_backend_relative_tolerance": args.cross_backend_relative_tolerance,
        },
    }


def write_results(
    output_dir: pathlib.Path,
    metadata: dict[str, Any],
    records: list[dict[str, Any]],
    summaries: list[dict[str, Any]],
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "results.json").write_text(
        json.dumps(
            {"schema_version": 1, "metadata": metadata, "raw": records, "summary": summaries},
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )
    fieldnames = [
        "case",
        "device",
        "ranks",
        "threads",
        "trials",
        "global_true_dofs",
        "requested_applications",
        "applications_min",
        "applications_max",
        "apply_mdof_per_second_median",
        "apply_mdof_per_second_mad",
        "apply_mdof_per_second_min",
        "apply_mdof_per_second_max",
        "speedup_vs_cpu",
        "apply_seconds_median",
        "solve_seconds_median",
        "solve_ran",
        "total_seconds_median",
        "cg_iterations",
        "verified_relative_residual_max",
    ]
    with (output_dir / "results.csv").open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for summary in summaries:
            writer.writerow(summary)
    first = records[0]
    if first["minimum_apply_seconds"] > 0.0:
        timing_description = (
            "Application counts were selected with a "
            f"{first['minimum_apply_seconds']:.3g}-second calibration window, then "
            "measured in a fresh timing window."
        )
    else:
        timing_description = (
            "The requested application count was measured directly in one timing "
            "window."
        )
    report = "# MFEM backend benchmark\n\n"
    report += (
        f"Host: `{metadata['hostname']}`  \n"
        f"UTC timestamp: `{metadata['timestamp_utc']}`  \n"
        f"Workload: {first['global_elements']:,} 3D hex elements, H1 order "
        f"{first['order']}, {first['global_true_dofs']:,} global true DoFs, "
        f"at least {first['requested_applications']} timed operator applications "
        f"per trial. {timing_description}\n\n"
    )
    report += markdown_table(summaries)
    report += (
        "\nThe table reports medians from fresh processes. Operator throughput uses "
        "the same global partial-assembly diffusion operator and true-DOF field in "
        "every case. MPI timings use the slowest rank.\n"
    )
    (output_dir / "results.md").write_text(report, encoding="utf-8")


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--executable", required=True)
    parser.add_argument("--launcher", required=True)
    parser.add_argument(
        "--cases",
        default="cpu,ceed-cpu,omp,mpi,cuda,ceed-cuda",
        help="Comma-separated subset of cpu, ceed-cpu, omp, mpi, cuda, ceed-cuda.",
    )
    parser.add_argument("--omp-threads", type=integer_list, default=[8])
    parser.add_argument("--mpi-ranks", type=integer_list, default=[4])
    parser.add_argument("--trials", type=positive_int, default=5)
    parser.add_argument("--warmup-trials", type=int, default=0)
    parser.add_argument("--mesh-n", type=positive_int, default=48)
    parser.add_argument("--order", type=positive_int, default=3)
    parser.add_argument("--applications", type=positive_int, default=50)
    parser.add_argument("--max-applications", type=positive_int, default=1_000_000)
    parser.add_argument(
        "--minimum-apply-seconds",
        type=float,
        default=1.0,
        help=(
            "Double the requested applications until a calibration window is this "
            "long, then measure a fresh window."
        ),
    )
    parser.add_argument("--warmup-applications", type=int, default=5)
    parser.add_argument("--relative-tolerance", type=float, default=1.0e-6)
    parser.add_argument("--max-iterations", type=positive_int, default=1000)
    parser.add_argument(
        "--solve",
        action="store_true",
        help="Also run and time the Jacobi-preconditioned CG validation solve.",
    )
    parser.add_argument("--timeout", type=positive_int, default=900)
    parser.add_argument(
        "--affinity",
        default="auto",
        help="CPU list for taskset, 'auto' for a compact set, or 'none'.",
    )
    parser.add_argument(
        "--network-interface",
        default="local",
        help="'local' for platform loopback, 'auto' for routing, or an interface name.",
    )
    parser.add_argument(
        "--cuda-visible-devices",
        default=None,
        help="Optional CUDA_VISIBLE_DEVICES value for GPU cases.",
    )
    parser.add_argument(
        "--gpu-aware-mpi",
        action="store_true",
        help="Opt in only when the selected MPI implementation is CUDA-aware.",
    )
    parser.add_argument(
        "--cross-backend-relative-tolerance",
        type=float,
        default=1.0e-5,
        help="Maximum relative difference from CPU for probe and solution norms.",
    )
    parser.add_argument(
        "--ceed-cpu-device",
        default="ceed-cpu",
        help="Full libCEED CPU resource or portable ceed-cpu alias.",
    )
    parser.add_argument(
        "--ceed-cuda-device",
        default="ceed-cuda",
        help="Full libCEED CUDA resource or portable ceed-cuda alias.",
    )
    parser.add_argument(
        "--launcher-arg",
        action="append",
        default=[],
        help="Additional MPI launcher argument; repeat as needed.",
    )
    parser.add_argument(
        "--output-dir",
        type=pathlib.Path,
        default=pathlib.Path("benchmark-results"),
    )
    args = parser.parse_args()
    if args.warmup_trials < 0 or args.warmup_applications < 0:
        parser.error("warm-up counts may not be negative")
    if args.relative_tolerance <= 0.0 or args.cross_backend_relative_tolerance <= 0.0:
        parser.error("tolerances must be positive")
    if args.minimum_apply_seconds < 0.0:
        parser.error("minimum apply seconds may not be negative")
    if args.max_applications < args.applications:
        parser.error("max applications must be at least the requested applications")
    for attribute in ("executable", "launcher"):
        value = getattr(args, attribute)
        resolved = resolve_program(value)
        if resolved is None:
            parser.error(f"program not found: {value}")
        setattr(args, attribute, resolved)
    return args


def main() -> int:
    args = parse_arguments()
    try:
        cases = expand_cases(args)
    except ValueError as error:
        print(f"error: {error}", file=sys.stderr)
        return 2
    metadata = collect_metadata(args)
    records: list[dict[str, Any]] = []

    for warmup in range(args.warmup_trials):
        for case in cases:
            print(f"warm-up {warmup + 1}/{args.warmup_trials}: {case.name}", flush=True)
            records.append(run_case(args, case, -(warmup + 1), True))
            checkpoint_results(args.output_dir, metadata, records)

    schedule = [(trial, case) for trial in range(1, args.trials + 1) for case in cases]
    random.Random(20260905).shuffle(schedule)
    for index, (trial, case) in enumerate(schedule, start=1):
        print(
            f"run {index}/{len(schedule)}: {case.name}, trial {trial}/{args.trials}",
            flush=True,
        )
        records.append(run_case(args, case, trial, False))
        checkpoint_results(args.output_dir, metadata, records)

    summaries = summarize(cases, records)
    workload_signatures = {
        (
            row["global_elements"],
            row["global_true_dofs"],
            row["mesh_n"],
            row["order"],
            row["requested_applications"],
            row["minimum_apply_seconds"],
        )
        for row in records
        if not row["discarded_warmup_trial"]
    }
    if len(workload_signatures) != 1:
        raise RuntimeError(
            f"cases did not use one fixed global workload: {workload_signatures}"
        )
    metadata["cross_backend_maximum_relative_errors"] = validate_cross_backend(
        records, args.cross_backend_relative_tolerance
    )
    if shutil.which("nvidia-smi"):
        metadata["nvidia_smi_after"] = command_output(
            [
                "nvidia-smi",
                "--query-gpu=name,driver_version,temperature.gpu,pstate,clocks.sm,power.draw",
                "--format=csv,noheader",
            ]
        )
    write_results(args.output_dir, metadata, records, summaries)
    table = markdown_table(summaries)
    print("\n" + table)
    print(f"JSON, CSV, and Markdown written to {args.output_dir.resolve()}")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, subprocess.TimeoutExpired, RuntimeError) as error:
        print(f"benchmark failed: {error}", file=sys.stderr)
        raise SystemExit(1)