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

import argparse
import gzip
import importlib.util
import json
import mmap
import os
import re
import shlex
import shutil
import subprocess
import sys
import time
import types
from collections import Counter
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import numpy as np
import onnx
from colorama import Fore, Style


REPO = Path(__file__).resolve().parents[3]
VALIDATION_DIR = REPO / "validation"
PIMSIM_CONFIG_DIR = VALIDATION_DIR / "pimsim_configs/pimcomp"
PIMCOMP_OUTPUT_FILES = ("SimulationInfo.gz", "VerificationInfo.json", "MappingResult.txt")
sys.path.insert(0, str(VALIDATION_DIR))

from raptor_validation.gen_network_runner import gen_network_runner  # noqa: E402
from raptor_validation.onnx_utils import (  # noqa: E402
    _ONNX_TO_NP,
    gen_random_inputs,
    onnx_io,
    save_inputs_to_files,
    write_inputs_to_memory_bin,
)
from raptor_validation.raptor import compile_with_raptor  # noqa: E402
from raptor_validation.pimsim_nn import (  # noqa: E402
    export_raptor_latency_artifact,
    parse_pimsim_nn_metrics,
)
from raptor_validation.validate_one import (  # noqa: E402
    STAGE_COLORS,
    build_dump_ranges,
    parse_pim_simulator_outputs,
)


@dataclass
class StepRecord:
    name: str
    duration_sec: float
    command: str
    status: str = "passed"
    returncode: int | None = None
    error: str | None = None
    output_tail: str | None = None


@dataclass
class CompareResult:
    passed: bool
    max_diffs: dict[str, float]
    status: str = "done"
    error: str | None = None


def load_pimcomp_exporter():
    path = REPO / "third_party/PIMCOMP-NN/verification/export_to_pim_simulator.py"
    spec = importlib.util.spec_from_file_location("pimcomp_exporter", path)
    module = importlib.util.module_from_spec(spec)
    assert spec is not None and spec.loader is not None
    sys.modules.setdefault("cv2", types.ModuleType("cv2"))
    write_bytecode = sys.dont_write_bytecode
    sys.dont_write_bytecode = True
    try:
        spec.loader.exec_module(module)
    finally:
        sys.dont_write_bytecode = write_bytecode
    return module


def shell_join(cmd: list[str]) -> str:
    return shlex.join(str(arg) for arg in cmd)


def print_step(
    name: str,
    cmd: list[str] | None = None,
    cwd: Path | None = None,
    stage: str | None = None,
):
    color = STAGE_COLORS.get(stage or name, Fore.WHITE)
    print("\n" + Style.BRIGHT + color + f"[{name}]" + Style.RESET_ALL)
    if cmd is not None:
        print(f"  cwd: {cwd or REPO}")
        print(f"  $ {shell_join(cmd)}")


def output_tail(output: str | bytes | None, limit: int = 4000) -> str:
    if output is None:
        return ""
    if isinstance(output, bytes):
        output = output.decode(errors="replace")
    return output[-limit:]


def exception_message(exc: BaseException) -> str:
    if isinstance(exc, subprocess.CalledProcessError):
        command = shell_join([str(arg) for arg in exc.cmd]) if isinstance(exc.cmd, list) else str(exc.cmd)
        tail = output_tail(exc.output)
        message = f"command failed with exit code {exc.returncode}: {command}"
        if tail:
            message += f"\n--- output tail ---\n{tail}"
        return message
    if isinstance(exc, subprocess.TimeoutExpired):
        command = shell_join([str(arg) for arg in exc.cmd]) if isinstance(exc.cmd, list) else str(exc.cmd)
        tail = output_tail(exc.output)
        message = f"command timed out after {exc.timeout} seconds: {command}"
        if tail:
            message += f"\n--- output tail ---\n{tail}"
        return message
    return f"{type(exc).__name__}: {exc}"


def print_failure(name: str, exc: BaseException | str) -> None:
    message = exc if isinstance(exc, str) else exception_message(exc)
    print(
        "\n" + Style.BRIGHT + Fore.RED + f"[{name} FAILED]" + Style.RESET_ALL,
        file=sys.stderr,
    )
    for line in message.splitlines()[:20]:
        print(Fore.RED + f"  {line}" + Style.RESET_ALL, file=sys.stderr)


def run_logged(
    name: str,
    cmd: list[str],
    *,
    cwd: Path,
    timeout_sec: float,
    steps: list[StepRecord],
    stage: str | None = None,
) -> str:
    print_step(name, cmd, cwd, stage)
    start = time.perf_counter()
    command = shell_join(cmd)
    try:
        proc = subprocess.run(
            cmd,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            timeout=timeout_sec,
        )
    except subprocess.TimeoutExpired as exc:
        duration = time.perf_counter() - start
        tail = output_tail(exc.output)
        steps.append(
            StepRecord(
                name=name,
                duration_sec=duration,
                command=command,
                status="timeout",
                error=f"Timed out after {timeout_sec} seconds",
                output_tail=tail or None,
            )
        )
        raise

    duration = time.perf_counter() - start
    if proc.returncode != 0:
        tail = output_tail(proc.stdout)
        steps.append(
            StepRecord(
                name=name,
                duration_sec=duration,
                command=command,
                status="failed",
                returncode=proc.returncode,
                error=f"Exited with status {proc.returncode}",
                output_tail=tail or None,
            )
        )
        raise subprocess.CalledProcessError(proc.returncode, cmd, output=tail)

    steps.append(StepRecord(name=name, duration_sec=duration, command=command))
    return proc.stdout


def load_model_inputs(model_path: Path, seed: int):
    inputs_desc, outputs_desc = onnx_io(model_path)
    arrays_in_order, _ = gen_random_inputs(inputs_desc, seed=seed)
    return inputs_desc, outputs_desc, arrays_in_order


def load_saved_inputs(
    inputs_desc: list[tuple[int, str, int, list[int]]],
    inputs_dir: Path,
) -> list[np.ndarray]:
    arrays = []
    for idx, name, elem_type, shape in inputs_desc:
        array = np.loadtxt(inputs_dir / f"in{idx}.csv", delimiter=",", dtype=_ONNX_TO_NP[elem_type]).reshape(shape)
        arrays.append(array)
    return arrays


def prepare_pimcomp_model(model_path: Path, out_dir: Path) -> Path:
    model = onnx.load(model_path)
    if not any(node.op_type == "BatchNormalization" for node in model.graph.node):
        return model_path

    out_dir.mkdir(parents=True, exist_ok=True)
    output_path = out_dir / f"{model_path.stem}_pimcomp.onnx"
    from onnxsim import simplify

    model, equivalent = simplify(model, check_n=1)
    if not equivalent:
        raise RuntimeError("Conv+BatchNormalization folding changed the model output")
    if any(node.op_type == "BatchNormalization" for node in model.graph.node):
        import onnxruntime as ort

        options = ort.SessionOptions()
        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
        options.optimized_model_filepath = str(output_path)
        ort.InferenceSession(str(model_path), options, providers=["CPUExecutionProvider"])
        model = onnx.load(output_path)
        if any(node.op_type == "BatchNormalization" for node in model.graph.node):
            raise RuntimeError("PIMCOMP model preparation did not eliminate BatchNormalization")
    else:
        onnx.save(model, output_path)
    return output_path


def compare_simulator_outputs(
    output_bin: Path,
    outputs_desc: list[tuple[int, str, int, list[int]]],
    reference_dir: Path,
    *,
    threshold: float,
    rtol: float,
    channel_last: bool = False,
) -> CompareResult:
    sim_arrays = parse_pim_simulator_outputs(output_bin, outputs_desc)
    max_diffs: dict[str, float] = {}
    passed = True
    for sim_array, (idx, name, _, shape) in zip(sim_arrays, outputs_desc):
        if channel_last and len(shape) == 4:
            n, c, h, w = shape
            sim_array = sim_array.reshape(n, h, w, c).transpose(0, 3, 1, 2)
        csv_name = reference_dir / f"output{idx}_{sanitize_output_name(name)}.csv"
        ref = np.loadtxt(csv_name, delimiter=",", dtype=np.float32).reshape(shape)
        diff = np.abs(sim_array.astype(np.float64) - ref.astype(np.float64))
        allowed = threshold + rtol * np.abs(ref.astype(np.float64))
        max_diffs[name] = float(np.max(diff))
        if not np.all(diff <= allowed):
            passed = False
    return CompareResult(passed=passed, max_diffs=max_diffs)


def sanitize_output_name(name: str) -> str:
    return "".join(ch if ch.isalnum() or ch in "_.-" else "_" for ch in name[:255])


def load_effective_hardware(args: argparse.Namespace) -> dict[str, int]:
    with open(args.pimcomp_config, "r", encoding="utf-8") as f:
        config = json.load(f)
    rows, cols = config["chip_config"]["network_config"]["layout"]
    xbar_h, xbar_w = config["chip_config"]["core_config"]["matrix_config"]["xbar_size"]
    hardware = {
        "mesh_rows": args.mesh_rows or rows,
        "mesh_cols": args.mesh_cols or cols,
        "crossbar_count": args.crossbar_count or config["chip_config"]["core_config"]["matrix_config"]["xbar_array_count"],
        "crossbar_size": args.crossbar_size or xbar_h,
    }
    if xbar_h != xbar_w:
        raise ValueError(f"Only square crossbars are supported, got {xbar_h}x{xbar_w}")
    hardware["core_count"] = args.core_count or hardware["mesh_rows"] * hardware["mesh_cols"]
    return hardware


def select_pimsim_config(args: argparse.Namespace, hardware: dict[str, int]) -> Path:
    for path in sorted(PIMSIM_CONFIG_DIR.glob(f"*/{args.pimsim_mode}_config.json")):
        with open(path, encoding="utf-8") as f:
            config = json.load(f)
        chip = config["chip_config"]
        matrix = chip["core_config"]["matrix_config"]
        network = chip["network_config"]
        if (
            chip["core_cnt"] == hardware["core_count"]
            and matrix["xbar_array_count"] == hardware["crossbar_count"]
            and matrix["xbar_size"] == [hardware["crossbar_size"]] * 2
            and network["layout"] == [hardware["mesh_rows"], hardware["mesh_cols"]]
            and config["sim_config"]["sim_mode"] == (1 if args.pimsim_mode == "latency" else 0)
            and config["sim_config"]["sim_time"] == args.pimsim_time_ms
        ):
            return path
    raise ValueError(
        f"No pre-generated {args.pimsim_mode} pimsim-nn config matches "
        f"{hardware} with sim_time={args.pimsim_time_ms}"
    )


def compile_reference(
    args: argparse.Namespace,
    model_path: Path,
    work_dir: Path,
    steps: list[StepRecord],
) -> Path:
    raptor_dir = work_dir / "raptor"
    runner_dir = work_dir / "runner"
    build_dir = runner_dir / "build"
    raptor_dir.mkdir(parents=True, exist_ok=True)
    build_dir.mkdir(parents=True, exist_ok=True)
    stem = model_path.stem
    onnx_ir_base = raptor_dir / stem
    runner_base = runner_dir / stem

    run_logged(
        "Compile Reference ONNX IR",
        [str(args.raptor_path), str(model_path), "-o", str(onnx_ir_base), "--EmitONNXIR",
         "--mlir-elide-elementsattrs-if-larger=16", "--enable-conv-opt-pass=false"],
        cwd=REPO,
        timeout_sec=args.timeout_seconds,
        steps=steps,
        stage="Compile ONNX",
    )
    run_logged(
        "Compile Reference Native",
        [str(args.raptor_path), "-O3", str(model_path), "-o", str(runner_base)],
        cwd=REPO,
        timeout_sec=args.timeout_seconds,
        steps=steps,
        stage="Compile ONNX",
    )
    network_so = runner_base.with_suffix(".so")

    print_step("Generate Runner Source", stage="Build Runner")
    gen_network_runner(
        model_path,
        network_so,
        args.onnx_include_dir,
        entry="run_main_graph",
        out=runner_dir / "runner.c",
        verbose=False,
    )

    run_logged(
        "Configure Runner",
        ["cmake", str(runner_dir), "-DCMAKE_BUILD_TYPE=Release", "-DCMAKE_C_FLAGS_RELEASE=-O3"],
        cwd=build_dir,
        timeout_sec=args.timeout_seconds,
        steps=steps,
        stage="Build Runner",
    )
    run_logged(
        "Build Runner",
        ["cmake", "--build", ".", "-j"],
        cwd=build_dir,
        timeout_sec=args.timeout_seconds,
        steps=steps,
    )
    return build_dir / "runner"


def generate_reference_outputs(
    runner_path: Path,
    runner_build_dir: Path,
    model_path: Path,
    arrays_in_order: list[np.ndarray],
    steps: list[StepRecord],
    args: argparse.Namespace,
    out_dir: Path,
) -> Path:
    inputs_dir = out_dir / "inputs"
    reference_dir = out_dir / "outputs"
    inputs_dir.mkdir(parents=True, exist_ok=True)
    reference_dir.mkdir(parents=True, exist_ok=True)
    flags, _ = save_inputs_to_files(model_path, arrays_in_order, inputs_dir)
    run_logged(
        "Run Reference",
        [str(runner_path), *flags, "--save-csv-dir", str(reference_dir)],
        cwd=runner_build_dir,
        timeout_sec=args.timeout_seconds,
        steps=steps,
    )
    return reference_dir


def compile_raptor_target(
    model_path: Path,
    out_dir: Path,
    hardware: dict[str, int],
    args: argparse.Namespace,
    steps: list[StepRecord],
) -> tuple[Path, dict[str, float]]:
    out_dir.mkdir(parents=True, exist_ok=True)
    cmd = [
        str(args.raptor_path),
        str(model_path),
        "-o",
        str(out_dir / "model"),
        "--maccel=PIM",
        "--EmitPimCodegen",
        f"--crossbar-size={hardware['crossbar_size']}",
        f"--crossbar-count={hardware['crossbar_count']}",
        f"--core-count={hardware['core_count']}",
        f"--pim-target-config={args.pimcomp_config}",
        "--pim-emit-json",
        *args.raptor_extra_arg,
    ]
    print_step("Compile Raptor PIM", cmd, REPO, "Compile PIM")
    start = time.perf_counter()
    command = shell_join(cmd)
    raptor_extra_args = [
        f"--pim-target-config={args.pimcomp_config}",
        "--pim-emit-json",
        *args.raptor_extra_arg,
    ]
    try:
        timings = compile_with_raptor(
            model_path,
            args.raptor_path,
            out_dir / "model",
            hardware["crossbar_size"],
            hardware["crossbar_count"],
            core_count=hardware["core_count"],
            raptor_extra_args=raptor_extra_args,
            cwd=out_dir,
            verbose=args.verbose_raptor_compile,
            reporter=None,
            timeout_sec=args.timeout_seconds,
        )
    except Exception as exc:
        steps.append(
            StepRecord(
                name="Compile Raptor PIM",
                duration_sec=time.perf_counter() - start,
                command=command,
                status="failed",
                error=exception_message(exc),
            )
        )
        raise

    steps.append(
        StepRecord(
            name="Compile Raptor PIM",
            duration_sec=time.perf_counter() - start,
            command=command,
        )
    )
    return out_dir / "pim", timings


def run_functional_validation(
    label: str,
    pim_dir: Path,
    config_path: Path,
    output_bin: Path,
    outputs_desc: list[tuple[int, str, int, list[int]]],
    reference_dir: Path,
    steps: list[StepRecord],
    args: argparse.Namespace,
    *,
    channel_last: bool = False,
) -> CompareResult:
    dump_ranges = build_dump_ranges(config_path, outputs_desc)
    cmd = [
        "cargo",
        "run",
        "--no-default-features",
        "--release",
        "--package",
        "pim-simulator",
        "--bin",
        "pim-simulator",
        "--",
        "-f",
        str(pim_dir),
        "-o",
        str(output_bin),
        "-d",
        dump_ranges,
    ]
    output_bin.parent.mkdir(parents=True, exist_ok=True)
    run_logged(
        label,
        cmd,
        cwd=args.pim_simulator_dir,
        timeout_sec=args.timeout_seconds,
        steps=steps,
        stage="Run Functional Simulation",
    )
    return compare_simulator_outputs(
        output_bin,
        outputs_desc,
        reference_dir,
        threshold=args.threshold,
        rtol=args.rtol,
        channel_last=channel_last,
    )


def copy_pimcomp_outputs(source_dir: Path, out_dir: Path):
    out_dir.mkdir(parents=True, exist_ok=True)
    for name in PIMCOMP_OUTPUT_FILES:
        shutil.copy2(source_dir / name, out_dir / name)


def compile_pimcomp(
    args: argparse.Namespace,
    model_path: Path,
    out_dir: Path,
    steps: list[StepRecord],
) -> tuple[Path, Path]:
    frontend_json_dir = out_dir / "models/JSON"
    frontend_json_dir.mkdir(parents=True, exist_ok=True)
    runtime_config = out_dir / "config.json"
    if args.pimcomp_config != runtime_config:
        shutil.copy2(args.pimcomp_config, runtime_config)
    pimcomp_output_dir = out_dir / "output"
    pimcomp_output_dir.mkdir(parents=True, exist_ok=True)
    for name in PIMCOMP_OUTPUT_FILES:
        (pimcomp_output_dir / name).unlink(missing_ok=True)
    model_name = args.pimcomp_model_name or f"compare_{model_path.stem}"
    frontend_json = frontend_json_dir / f"{model_name}.json"
    frontend_cmd = [
        sys.executable,
        str(args.pimcomp_dir / "frontend/frontend.py"),
        "--model_path",
        str(model_path),
        "--save_path",
        str(frontend_json),
    ]
    run_logged(
        "Compile PIMCOMP Frontend",
        frontend_cmd,
        cwd=args.pimcomp_dir / "frontend",
        timeout_sec=args.timeout_seconds,
        steps=steps,
        stage="Compile PIM",
    )
    backend_cmd = [
        str(args.pimcomp_dir / "build" / "PIMCOMP-NN"),
        f"-m={model_name}",
        f"-r={args.pimcomp_replication}",
        f"-p={args.pimcomp_pipeline}",
        "-v=YES",
        "-s=YES",
    ]
    run_logged(
        "Compile PIMCOMP Backend",
        backend_cmd,
        cwd=frontend_json_dir.parent,
        timeout_sec=args.timeout_seconds,
        steps=steps,
        stage="Compile PIM",
    )
    shutil.rmtree(frontend_json_dir.parent)
    return pimcomp_output_dir / "VerificationInfo.json", pimcomp_output_dir / "SimulationInfo.gz"


def export_pimcomp_for_pimsim_nn(simulation_info: Path, output_dir: Path) -> Path:
    if output_dir.exists():
        shutil.rmtree(output_dir)
    with gzip.open(simulation_info, "rt", encoding="utf-8") as f:
        sim_info = json.load(f)

    output_dir.mkdir(parents=True, exist_ok=True)
    sim_config = sim_info["config"]
    core_count = int(sim_config["core_cnt"])
    if core_count <= 0:
        raise ValueError("PIMCOMP SimulationInfo.gz must configure at least one core")
    core_indices = range(core_count)

    config = {
        "core_cnt": core_count,
        "xbar_size": sim_config["xbar_size"],
        "xbar_array_count": sim_config["xbar_array_count"],
        "cell_precision": sim_config["cell_precision"],
        "adc_count": sim_config["adc_count"],
        "array_group_map": {},
    }
    for core_idx in core_indices:
        core_name = f"core{core_idx}"
        config["array_group_map"][core_name] = sim_config["array_group_map"].get(core_name, [])

    with open(output_dir / "config.json", "w", encoding="utf-8") as f:
        json.dump(config, f, separators=(",", ":"))
        f.write("\n")

    for core_idx in core_indices:
        core_key = f"core{core_idx}"
        instructions = sim_info.get(core_key, [])
        with open(output_dir / f"core_{core_idx}.json", "w", encoding="utf-8") as f:
            json.dump(instructions, f, separators=(",", ":"))
            f.write("\n")
    return output_dir


def flatten_pimcomp_input(array: np.ndarray) -> np.ndarray:
    tensor = array.astype(np.float32, copy=False)
    if tensor.ndim == 4:
        tensor = tensor.transpose((0, 2, 3, 1))
    return tensor.reshape(-1)


def export_pimcomp_for_rust(
    model_path: Path,
    verification_info: Path,
    simulation_info: Path,
    runtime_inputs: list[np.ndarray],
    output_dir: Path,
) -> Path:
    if len(runtime_inputs) != 1:
        raise ValueError("PIMCOMP export currently requires exactly one runtime input tensor")
    if output_dir.exists():
        shutil.rmtree(output_dir)
    exporter = load_pimcomp_exporter()
    with open(verification_info, "r", encoding="utf-8") as f:
        final_info = json.load(f)
    with gzip.open(simulation_info, "rt", encoding="utf-8") as f:
        sim_info = json.load(f)

    onnx_model, weights, gemm_weights, output_to_weight, output_to_bias = exporter.load_model_info(
        model_path, final_info
    )
    input_tensor = flatten_pimcomp_input(runtime_inputs[0])
    node_list = final_info["node_list"]
    max_output = exporter.max_output_element_num(node_list)
    local_group_map = exporter.map_local_groups(final_info, sim_info)

    output_dir.mkdir(parents=True, exist_ok=True)
    weights_dir = output_dir / "weights"
    weights_dir.mkdir(parents=True, exist_ok=True)

    input_addr = 0
    cursor = exporter.byte_offset(len(input_tensor))
    bias_addrs: dict[str, int] = {}
    for node_name, bias_name in output_to_bias.items():
        bias = weights[bias_name].astype(np.float32).flatten()
        bias_addrs[node_name] = cursor
        cursor += exporter.byte_offset(len(bias))

    lldi_addrs: dict[tuple[bytes, int], int] = {}
    for core_idx in range(sim_info["config"]["core_cnt"]):
        for inst in sim_info.get(f"core{core_idx}", []) or []:
            if inst["op"] != "lldi":
                continue
            key = (exporter.float32_bytes(inst["imm"]), inst["len"])
            if key not in lldi_addrs:
                lldi_addrs[key] = cursor
                cursor += exporter.byte_offset(inst["len"])

    output_base = (cursor + 255) & ~255
    memory_size = output_base + exporter.byte_offset(max_output * len(node_list))
    memory = bytearray(memory_size)
    memory[input_addr : input_addr + input_tensor.nbytes] = input_tensor.tobytes()
    for node_name, bias_name in output_to_bias.items():
        bias = weights[bias_name].astype(np.float32).flatten()
        start = bias_addrs[node_name]
        memory[start : start + bias.nbytes] = bias.tobytes()
    for (value_bytes, element_num), start in lldi_addrs.items():
        value = np.frombuffer(value_bytes, dtype=np.float32)[0]
        blob = np.full(element_num, value, dtype=np.float32)
        memory[start : start + blob.nbytes] = blob.tobytes()

    config = {
        "core_cnt": sim_info["config"]["core_cnt"],
        "xbar_size": sim_info["config"]["xbar_size"],
        "xbar_array_count": sim_info["config"]["xbar_array_count"],
        "cell_precision": sim_info["config"]["cell_precision"],
        "adc_count": sim_info["config"]["adc_count"],
        "array_group_map": {},
        "inputs_addresses": [input_addr],
        "outputs_addresses": [],
    }
    output_name_to_node = {node["name"]: node for node in node_list}
    for graph_output in onnx_model.graph.output:
        node = output_name_to_node[graph_output.name]
        config["outputs_addresses"].append(output_base + exporter.byte_offset(node["new_node_index"] * max_output))

    ag_info = final_info["AG_info"]
    weight_counter = 0
    xbar_size = int(sim_info["config"]["xbar_size"][0])
    for core_idx in range(config["core_cnt"]):
        core_name = f"core{core_idx}"
        core_dir = output_dir / f"core_{core_idx}"
        core_dir.mkdir(parents=True, exist_ok=True)
        local_to_global = local_group_map.get(core_idx, {})
        ag_counts = sim_info["config"]["array_group_map"].get(core_name, [])
        local_group_to_physical = {}
        total_crossbars = 0
        for local_group in sorted(local_to_global):
            width = ag_counts[local_group]
            local_group_to_physical[local_group] = total_crossbars
            total_crossbars += width
        config["array_group_map"][core_name] = list(range(total_crossbars))

        for local_group, global_ag in sorted(local_to_global.items()):
            info = ag_info[global_ag]
            weight_name = output_to_weight[info["node_name"]]
            matrix = gemm_weights[weight_name]
            row_slice = slice(info["height_start"], info["height_end"] + 1)
            first_physical = local_group_to_physical[local_group]
            for crossbar_idx, crossbar in enumerate(info["crossbar"]):
                col_slice = slice(crossbar["width_start"], crossbar["width_end"] + 1)
                tile = np.zeros((xbar_size, col_slice.stop - col_slice.start), dtype=np.float32)
                tile_rows = matrix[row_slice, col_slice].astype(np.float32)
                tile[: tile_rows.shape[0], :] = tile_rows
                weight_path = weights_dir / f"crossbar_{weight_counter}.bin"
                weight_path.write_bytes(tile.tobytes(order="C"))
                os.symlink(weight_path.resolve(), core_dir / f"crossbar_{first_physical + crossbar_idx}.bin")
                weight_counter += 1

        instructions = []
        last_sldi_by_rd: dict[int, int] = {}
        ver_ops = exporter.filtered_verification_ops(final_info, core_idx)
        ver_index = 0
        for sim_inst in sim_info.get(core_name, []) or []:
            op = sim_inst["op"]
            if op == "setbw":
                instructions.append({"op": "setbw", "ibiw": 32, "obiw": 32})
                continue
            if op == "sldi":
                translated = {"op": "sldi", "rd": sim_inst["rd"], "imm": exporter.byte_offset(sim_inst["imm"])}
                instructions.append(translated)
                last_sldi_by_rd[sim_inst["rd"]] = len(instructions) - 1
                continue
            if ver_index >= len(ver_ops):
                raise RuntimeError(f"core{core_idx}: simulation op {op} has no matching verification op")
            ver_inst = ver_ops[ver_index]
            ver_index += 1
            ver_op = ver_inst["operation"].lower()
            if ver_op != op:
                raise RuntimeError(
                    f"core{core_idx}: simulation/verification op mismatch {op} vs {ver_op} at {ver_index - 1}"
                )
            if op == "ld":
                if ver_inst["stage"] == "INPUT":
                    if ver_inst["node_index"] == 1:
                        src = input_addr + exporter.byte_offset(ver_inst["source_offset"])
                    else:
                        provider_index = -ver_inst["source_address"]
                        src = output_base + exporter.byte_offset(
                            provider_index * max_output + ver_inst["source_offset"]
                        )
                elif ver_inst["stage"] == "BIAS":
                    src = bias_addrs[node_list[ver_inst["node_index"]]["name"]] + exporter.byte_offset(ver_inst["source_offset"])
                else:
                    raise RuntimeError(f"Unsupported LD stage {ver_inst['stage']}")
                instructions[last_sldi_by_rd[sim_inst["rs1"]]]["imm"] = src
                translated = dict(sim_inst)
                translated["size"] = exporter.byte_offset(sim_inst["size"])
                instructions.append(translated)
            elif op == "st":
                dst = output_base + exporter.byte_offset(
                    ver_inst["node_index"] * max_output + ver_inst["destination_offset"]
                )
                instructions[last_sldi_by_rd[sim_inst["rd"]]]["imm"] = dst
                translated = dict(sim_inst)
                translated["size"] = exporter.byte_offset(sim_inst["size"])
                instructions.append(translated)
            elif op == "lldi":
                key = (exporter.float32_bytes(sim_inst["imm"]), sim_inst["len"])
                src = lldi_addrs[key]
                temp_rd = 1 if sim_inst["rd"] == 0 else 0
                instructions.append({"op": "sldi", "rd": temp_rd, "imm": src})
                instructions.append(
                    {
                        "op": "ld",
                        "rd": sim_inst["rd"],
                        "rs1": temp_rd,
                        "size": exporter.byte_offset(sim_inst["len"]),
                        "offset": sim_inst["offset"],
                    }
                )
            elif op == "lmv":
                translated = dict(sim_inst)
                translated["len"] = exporter.byte_offset(sim_inst["len"])
                instructions.append(translated)
            elif op in ("vvadd", "vvmul", "vvmax", "vrelu"):
                instructions.append(sim_inst)
            elif op in ("send", "recv"):
                translated = dict(sim_inst)
                translated["size"] = exporter.byte_offset(sim_inst["size"])
                instructions.append(translated)
            elif op == "mvmul":
                local_group = sim_inst["group"]
                global_ag = local_to_global[local_group]
                first_physical = local_group_to_physical[local_group]
                widths = [
                    crossbar["width_end"] - crossbar["width_start"] + 1
                    for crossbar in ag_info[global_ag]["crossbar"]
                ]
                dst = instructions[last_sldi_by_rd[sim_inst["rd"]]]["imm"]
                src = instructions[last_sldi_by_rd[sim_inst["rs1"]]]["imm"]
                out_offset = 0
                for idx, width in enumerate(widths):
                    instructions.append({"op": "sldi", "rd": sim_inst["rd"], "imm": dst + exporter.byte_offset(out_offset)})
                    instructions.append({"op": "sldi", "rd": sim_inst["rs1"], "imm": src})
                    translated = dict(sim_inst)
                    translated["group"] = first_physical + idx
                    instructions.append(translated)
                    out_offset += width
            else:
                raise RuntimeError(f"Unsupported PIMCOMP op {op}")

        with open(output_dir / f"core_{core_idx}.json", "w", encoding="utf-8") as f:
            json.dump(instructions, f, separators=(",", ":"))
            f.write("\n")

    with open(output_dir / "config.json", "w", encoding="utf-8") as f:
        json.dump(config, f, separators=(",", ":"))
        f.write("\n")
    (output_dir / "memory.bin").write_bytes(memory)
    return output_dir


def run_pimsim_nn(
    label: str,
    inst_path: Path,
    config_path: Path,
    steps: list[StepRecord],
    args: argparse.Namespace,
) -> dict[str, Any]:
    cmd = [
        str(args.pimsim_nn_build_dir / "ChipTest"),
        str(inst_path),
        str(config_path),
        "--gui=false",
    ]
    output = run_logged(
        label,
        cmd,
        cwd=args.pimsim_nn_build_dir,
        timeout_sec=args.timeout_seconds * 10.0,
        steps=steps,
        stage="Run Non-functional Simulation",
    )
    return parse_pimsim_nn_metrics(output)


def parse_raptor_instructions(pim_dir: Path) -> dict[str, Any]:
    op_re = re.compile(br'"op":"([^"]+)"')
    counts = Counter()
    per_core = []
    for path in sorted(pim_dir.glob("core_*.json"), key=lambda p: int(p.stem.split("_")[1])):
        with path.open("rb") as f:
            mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
            core_counts = Counter(m.group(1).decode() for m in op_re.finditer(mm))
            mm.close()
        total = sum(core_counts.values())
        counts.update(core_counts)
        per_core.append(
            {
                "core": path.stem,
                "total": total,
                "send": core_counts.get("send", 0),
                "recv": core_counts.get("recv", 0),
                "mvmul": core_counts.get("mvmul", 0),
            }
        )
    return {
        "active_cores": sum(1 for entry in per_core if entry["total"]),
        "total_instructions": int(sum(counts.values())),
        "op_counts": dict(counts),
        "top_cores_by_total": sorted(per_core, key=lambda entry: entry["total"], reverse=True)[:10],
        "top_cores_by_send": sorted(per_core, key=lambda entry: entry["send"], reverse=True)[:10],
        "top_cores_by_recv": sorted(per_core, key=lambda entry: entry["recv"], reverse=True)[:10],
    }


def parse_pimcomp_instructions(simulation_info: Path) -> dict[str, Any]:
    with gzip.open(simulation_info, "rt", encoding="utf-8") as f:
        data = json.load(f)
    per_core = []
    counts = Counter()
    for key in sorted((name for name in data if name.startswith("core")), key=lambda name: int(name[4:])):
        insts = data[key]
        core_counts = Counter((inst.get("operation") or inst.get("op") or "unknown").lower() for inst in insts)
        counts.update(core_counts)
        per_core.append(
            {
                "core": key,
                "total": int(sum(core_counts.values())),
                "send": core_counts.get("send", 0),
                "recv": core_counts.get("recv", 0),
                "mvmul": core_counts.get("mvmul", 0),
            }
        )
    return {
        "active_cores": sum(1 for entry in per_core if entry["total"]),
        "total_instructions": int(sum(counts.values())),
        "op_counts": dict(counts),
        "top_cores_by_total": sorted(per_core, key=lambda entry: entry["total"], reverse=True)[:10],
        "top_cores_by_send": sorted(per_core, key=lambda entry: entry["send"], reverse=True)[:10],
        "top_cores_by_recv": sorted(per_core, key=lambda entry: entry["recv"], reverse=True)[:10],
    }


def format_op_table(counts: dict[str, int], total: int) -> list[str]:
    if total <= 0:
        return ["| n/a | 0 | n/a |"]
    rows = []
    for op, count in sorted(counts.items(), key=lambda item: item[1], reverse=True):
        rows.append(f"| `{op}` | {count} | {100.0 * count / total:.2f}% |")
    return rows


def validation_status(result: CompareResult) -> str:
    if result.status == "done":
        return "PASS" if result.passed else "FAIL"
    return result.status.upper()


def skipped_validation(reason: str) -> CompareResult:
    return CompareResult(passed=False, max_diffs={}, status="skipped", error=reason)


def failed_validation(error: BaseException | str) -> CompareResult:
    message = error if isinstance(error, str) else exception_message(error)
    return CompareResult(passed=False, max_diffs={}, status="failed", error=message)


def skipped_perf(reason: str) -> dict[str, Any]:
    return {"skipped": True, "reason": reason}


def failed_perf(error: BaseException | str) -> dict[str, Any]:
    message = error if isinstance(error, str) else exception_message(error)
    return {"error": message}


def perf_status(perf: dict[str, Any]) -> str:
    if perf.get("skipped"):
        return "SKIPPED"
    if perf.get("error"):
        return "FAILED"
    return "DONE"


def perf_value(perf: dict[str, Any], key: str) -> Any:
    return perf[key] if key in perf else "n/a"


def empty_instruction_summary(reason: str | None = None, error: str | None = None) -> dict[str, Any]:
    result: dict[str, Any] = {
        "active_cores": 0,
        "total_instructions": 0,
        "op_counts": {},
        "top_cores_by_total": [],
        "top_cores_by_send": [],
        "top_cores_by_recv": [],
    }
    if reason is not None:
        result["skipped"] = True
        result["reason"] = reason
    if error is not None:
        result["error"] = error
    return result


def optional_path(path: Path | None) -> str | None:
    return str(path) if path is not None else None


def record_failure(failures: list[dict[str, str]], stage: str, exc: BaseException | str) -> None:
    message = exc if isinstance(exc, str) else exception_message(exc)
    failures.append({"stage": stage, "error": message})
    print_failure(stage, message)


def try_stage(
    failures: list[dict[str, str]],
    stage: str,
    func,
    *args,
    **kwargs,
):
    try:
        return func(*args, **kwargs)
    except Exception as exc:
        record_failure(failures, stage, exc)
        return None


def try_stage_success(
    failures: list[dict[str, str]],
    stage: str,
    func,
    *args,
    **kwargs,
) -> bool:
    try:
        func(*args, **kwargs)
        return True
    except Exception as exc:
        record_failure(failures, stage, exc)
        return False


def write_report(
    report_path: Path,
    *,
    model_path: Path,
    hardware: dict[str, int],
    steps: list[StepRecord],
    failures: list[dict[str, str]],
    raptor_validation: CompareResult,
    pimcomp_validation: CompareResult,
    raptor_perf: dict[str, Any],
    pimcomp_perf: dict[str, Any],
    raptor_instr: dict[str, Any],
    pimcomp_instr: dict[str, Any],
    raptor_pass_timings: dict[str, float],
    pimsim_mode: str,
    pimcomp_pipeline: str,
    pimcomp_replication: str,
):
    lines = [
        "# Raptor vs PIMCOMP Comparison Report",
        "",
        f"- Model: `{model_path}`",
        f"- Hardware: `{hardware.get('core_count', 'n/a')} cores`, `{hardware.get('crossbar_count', 'n/a')} xbars/core`, `{hardware.get('crossbar_size', 'n/a')}x{hardware.get('crossbar_size', 'n/a')}` crossbars, mesh `{hardware.get('mesh_rows', 'n/a')}x{hardware.get('mesh_cols', 'n/a')}`",
        f"- PIMCOMP pipeline: `{pimcomp_pipeline}`",
        f"- PIMCOMP replication: `{pimcomp_replication}`",
        "",
    ]

    if failures or any(step.status != "passed" for step in steps):
        lines.extend(
            [
                "## Failures / Skipped Work",
                "",
                "The script did not abort. The failed stage was recorded and any dependent stage was skipped when its inputs were not available.",
                "",
            ]
        )
        if failures:
            lines.extend(["| Stage | Error |", "|---|---|"])
            for failure in failures:
                error = failure["error"].replace("\n", "<br>")
                lines.append(f"| {failure['stage']} | {error} |")
            lines.append("")

    lines.extend(
        [
            "## Functional Validation",
            "",
            f"- Raptor via `pim-simulator`: `{validation_status(raptor_validation)}`",
            f"- PIMCOMP via exported `pim-simulator`: `{validation_status(pimcomp_validation)}`",
        ]
    )
    if raptor_validation.error:
        lines.append(f"- Raptor validation note: `{raptor_validation.error.splitlines()[0]}`")
    if pimcomp_validation.error:
        lines.append(f"- PIMCOMP validation note: `{pimcomp_validation.error.splitlines()[0]}`")

    lines.extend(["", "### Max Output Differences", ""])
    diff_names = sorted(set(raptor_validation.max_diffs) | set(pimcomp_validation.max_diffs))
    if diff_names:
        lines.extend(["| Output | Raptor max diff | PIMCOMP max diff |", "|---|---:|---:|"])
        for name in diff_names:
            lines.append(
                f"| `{name}` | {raptor_validation.max_diffs.get(name, float('nan')):.6e} | "
                f"{pimcomp_validation.max_diffs.get(name, float('nan')):.6e} |"
            )
    else:
        lines.append("No output differences are available because validation did not run or failed before comparison.")

    lines.extend(
        [
            "",
            "## pimsim-nn Performance",
            "",
            f"- Mode: `{pimsim_mode}`",
            "",
        ]
    )
    if pimsim_mode == "throughput":
        lines.extend(
            [
                "| Compiler | Status | Throughput (samples/s) | Avg latency (ms) | Avg power (mW) | Avg energy (pJ/it) | Output count |",
                "|---|---|---:|---:|---:|---:|---:|",
                f"| Raptor | {perf_status(raptor_perf)} | {perf_value(raptor_perf, 'throughput')} | {perf_value(raptor_perf, 'average_latency_ms')} | "
                f"{perf_value(raptor_perf, 'average_power_mw')} | {perf_value(raptor_perf, 'average_energy_pj')} | {perf_value(raptor_perf, 'output_count')} |",
                f"| PIMCOMP | {perf_status(pimcomp_perf)} | {perf_value(pimcomp_perf, 'throughput')} | {perf_value(pimcomp_perf, 'average_latency_ms')} | "
                f"{perf_value(pimcomp_perf, 'average_power_mw')} | {perf_value(pimcomp_perf, 'average_energy_pj')} | {perf_value(pimcomp_perf, 'output_count')} |",
                "",
            ]
        )
    else:
        lines.extend(
            [
                "| Compiler | Status | Latency (ms) | Avg power (mW) | Avg energy (pJ) |",
                "|---|---|---:|---:|---:|",
                f"| Raptor | {perf_status(raptor_perf)} | {perf_value(raptor_perf, 'latency_ms')} | "
                f"{perf_value(raptor_perf, 'average_power_mw')} | {perf_value(raptor_perf, 'average_energy_pj')} |",
                f"| PIMCOMP | {perf_status(pimcomp_perf)} | {perf_value(pimcomp_perf, 'latency_ms')} | "
                f"{perf_value(pimcomp_perf, 'average_power_mw')} | {perf_value(pimcomp_perf, 'average_energy_pj')} |",
                "",
            ]
        )
    if raptor_perf.get("reason") or raptor_perf.get("error"):
        lines.append(f"- Raptor pimsim-nn note: `{(raptor_perf.get('reason') or raptor_perf.get('error')).splitlines()[0]}`")
    if pimcomp_perf.get("reason") or pimcomp_perf.get("error"):
        lines.append(f"- PIMCOMP pimsim-nn note: `{(pimcomp_perf.get('reason') or pimcomp_perf.get('error')).splitlines()[0]}`")
    if lines[-1] != "":
        lines.append("")

    lines.extend(
        [
            "## Instruction Summary",
            "",
            "| Compiler | Status | Active cores | Total instructions | Sends | Receives | MVMUL |",
            "|---|---|---:|---:|---:|---:|---:|",
            f"| Raptor | {'FAILED' if raptor_instr.get('error') else 'SKIPPED' if raptor_instr.get('skipped') else 'DONE'} | {raptor_instr.get('active_cores', 0)} | {raptor_instr.get('total_instructions', 0)} | {raptor_instr.get('op_counts', {}).get('send', 0)} | {raptor_instr.get('op_counts', {}).get('recv', 0)} | {raptor_instr.get('op_counts', {}).get('mvmul', 0)} |",
            f"| PIMCOMP | {'FAILED' if pimcomp_instr.get('error') else 'SKIPPED' if pimcomp_instr.get('skipped') else 'DONE'} | {pimcomp_instr.get('active_cores', 0)} | {pimcomp_instr.get('total_instructions', 0)} | {pimcomp_instr.get('op_counts', {}).get('send', 0)} | {pimcomp_instr.get('op_counts', {}).get('recv', 0)} | {pimcomp_instr.get('op_counts', {}).get('mvmul', 0)} |",
            "",
            "### Raptor Op Distribution",
            "",
            "| Op | Count | Share |",
            "|---|---:|---:|",
            *format_op_table(raptor_instr.get("op_counts", {}), raptor_instr.get("total_instructions", 0)),
            "",
            "### PIMCOMP Op Distribution",
            "",
            "| Op | Count | Share |",
            "|---|---:|---:|",
            *format_op_table(pimcomp_instr.get("op_counts", {}), pimcomp_instr.get("total_instructions", 0)),
            "",
            "## Step Timings",
            "",
            "| Step | Status | Duration (s) | Return code |",
            "|---|---|---:|---:|",
        ]
    )
    for step in steps:
        lines.append(
            f"| {step.name} | {step.status.upper()} | {step.duration_sec:.3f} | "
            f"{step.returncode if step.returncode is not None else ''} |"
        )
    failed_steps = [step for step in steps if step.status != "passed"]
    if failed_steps:
        lines.extend(["", "### Failed Step Details", ""])
        for step in failed_steps:
            lines.extend(
                [
                    f"#### {step.name}",
                    "",
                    f"- Command: `{step.command}`",
                    f"- Error: `{step.error or 'n/a'}`",
                ]
            )
            if step.output_tail:
                lines.extend(["", "```text", step.output_tail, "```"])
            lines.append("")

    if raptor_pass_timings:
        lines.extend(["", "## Raptor Pass Timings", "", "| Pass | Duration (s) |", "|---|---:|"])
        for name, duration in raptor_pass_timings.items():
            lines.append(f"| {name} | {duration:.4f} |")
    report_path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", required=True, type=Path)
    parser.add_argument("--out-dir", required=True, type=Path)
    parser.add_argument("--raptor-path", default=REPO / "build_release/Release/bin/onnx-mlir", type=Path)
    parser.add_argument("--onnx-include-dir", default=REPO / "onnx-mlir/include", type=Path)
    parser.add_argument("--pimcomp-dir", default=REPO / "third_party/PIMCOMP-NN", type=Path)
    parser.add_argument(
        "--pimcomp-config",
        type=Path,
        help="PIMCOMP hardware config (default: <pimcomp-dir>/config.json).",
    )
    parser.add_argument("--pim-simulator-dir", default=REPO / "backend-simulators/pim/pim-simulator", type=Path)
    parser.add_argument("--pimsim-nn-build-dir", default=REPO / "backend-simulators/pim/pimsim-nn/build", type=Path)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--threshold", type=float, default=1e-3)
    parser.add_argument("--rtol", type=float, default=1e-5)
    parser.add_argument("--timeout-seconds", type=float, default=3600.0)
    parser.add_argument("--core-count", type=int)
    parser.add_argument("--crossbar-count", type=int)
    parser.add_argument("--crossbar-size", type=int)
    parser.add_argument("--mesh-rows", type=int)
    parser.add_argument("--mesh-cols", type=int)
    parser.add_argument("--pimsim-time-ms", type=int, default=1000)
    parser.add_argument("--pimsim-mode", choices=["latency", "throughput"], default="latency")
    parser.add_argument("--pimcomp-pipeline", choices=["element", "batch"])
    parser.add_argument("--pimcomp-model-name", help="Use a PIMCOMP built-in model name such as vgg16.")
    parser.add_argument(
        "--pimcomp-replication",
        choices=["balance", "W0H0", "uniform", "GA"],
        default="balance",
    )
    parser.add_argument(
        "--reuse-raptor-report",
        type=Path,
        help="Reuse Raptor artifacts and results from an existing comparison_report.json.",
    )
    parser.add_argument(
        "--reuse-pimcomp-dir",
        type=Path,
        help="Reuse a directory containing PIMCOMP SimulationInfo.gz, VerificationInfo.json, and MappingResult.txt.",
    )
    parser.add_argument("--skip-pimsim-nn", action="store_true")
    parser.add_argument("--verbose-raptor-compile", action="store_true")
    parser.add_argument("--raptor-extra-arg", action="append", default=[])
    parser.add_argument(
        "--fail-on-error",
        action="store_true",
        help="Return a non-zero status if a stage or semantic validation fails.",
    )
    args = parser.parse_args()
    if args.pimcomp_pipeline is None:
        args.pimcomp_pipeline = "element" if args.pimsim_mode == "latency" else "batch"

    model_path = args.model.resolve()
    args.pimcomp_dir = args.pimcomp_dir.resolve()
    args.pimcomp_config = (
        args.pimcomp_config.resolve()
        if args.pimcomp_config is not None
        else args.pimcomp_dir / "config.json"
    )
    out_dir = args.out_dir.resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    failures: list[dict[str, str]] = []
    steps: list[StepRecord] = []
    hardware: dict[str, int] = {
        "mesh_rows": 0,
        "mesh_cols": 0,
        "crossbar_count": 0,
        "crossbar_size": 0,
        "core_count": 0,
    }
    inputs_desc: list[tuple[int, str, int, list[int]]] = []
    outputs_desc: list[tuple[int, str, int, list[int]]] = []
    arrays_in_order: list[np.ndarray] = []
    runtime_inputs: list[np.ndarray] = []

    runner_path: Path | None = None
    reference_dir: Path | None = None
    raptor_pim_dir: Path | None = None
    raptor_pimsim_dir: Path | None = None
    raptor_pass_timings: dict[str, float] = {}
    verification_info: Path | None = None
    simulation_info: Path | None = None
    pimcomp_export_dir: Path | None = None
    pimcomp_model_path: Path | None = None
    pimsim_config: Path | None = None
    reuse_raptor = args.reuse_raptor_report is not None

    raptor_validation = skipped_validation("Raptor validation did not run")
    pimcomp_validation = failed_validation("PIMCOMP validation did not run")
    raptor_perf: dict[str, Any] = skipped_perf("pimsim-nn Raptor did not run")
    pimcomp_perf: dict[str, Any] = skipped_perf("pimsim-nn PIMCOMP did not run")
    raptor_instr: dict[str, Any] = empty_instruction_summary("Raptor instruction parsing did not run")
    pimcomp_instr: dict[str, Any] = empty_instruction_summary("PIMCOMP instruction parsing did not run")

    loaded_hardware = try_stage(failures, "Load hardware configuration", load_effective_hardware, args)
    if loaded_hardware is not None:
        hardware = loaded_hardware

    model_io = try_stage(failures, "Load model inputs", load_model_inputs, model_path, args.seed)
    if model_io is not None:
        inputs_desc, outputs_desc, arrays_in_order = model_io
        runtime_inputs = arrays_in_order

    if reuse_raptor and model_io is not None:
        reuse_report_path = args.reuse_raptor_report.resolve()
        with open(reuse_report_path, "r", encoding="utf-8") as f:
            reused = json.load(f)
        reused_hardware = reused["hardware"]
        if reused_hardware != hardware:
            raise ValueError(f"Reused Raptor hardware differs: {reused_hardware} != {hardware}")
        reference_dir = Path(reused["paths"]["reference_outputs"])
        raptor_pim_dir = Path(reused["paths"]["raptor_pim"])
        arrays_in_order = load_saved_inputs(
            inputs_desc,
            reference_dir.parent / "inputs",
        )
        runtime_inputs = arrays_in_order
        raptor_validation = CompareResult(**reused["raptor_validation"])
        raptor_perf = reused["raptor_performance"]
        raptor_instr = reused["raptor_instruction_summary"]
        raptor_pass_timings = reused["raptor_pass_timings"]
        print_step("Reuse Raptor")
        print(f"  Report: {reuse_report_path}")

    expected_runner_path = out_dir / "runner/build/runner"

    if not reuse_raptor:
        reference_compile = try_stage(
            failures,
            "Compile reference",
            compile_reference,
            args,
            model_path,
            out_dir,
            steps,
        )
        if reference_compile is not None:
            runner_path = reference_compile
        else:
            if expected_runner_path.exists():
                runner_path = expected_runner_path
                print(
                    "\n"
                    + Style.BRIGHT
                    + Fore.YELLOW
                    + "[Continue]"
                    + Style.RESET_ALL
                    + f" Reusing partial runner: {runner_path}"
                )

    if not reuse_raptor and runner_path is not None and runner_path.exists() and model_io is not None:
        generated_reference = try_stage(
            failures,
            "Run reference",
            generate_reference_outputs,
            runner_path,
            runner_path.parent,
            model_path,
            arrays_in_order,
            steps,
            args,
            out_dir,
        )
        if generated_reference is not None:
            reference_dir = generated_reference
    elif not reuse_raptor:
        record_failure(
            failures,
            "Skip reference outputs",
            "Reference outputs were skipped because the native runner or model inputs are not available.",
        )

    if not reuse_raptor and model_path.exists() and hardware["core_count"] > 0:
        compiled_raptor = try_stage(
            failures,
            "Compile Raptor PIM",
            compile_raptor_target,
            model_path,
            out_dir / "raptor",
            hardware,
            args,
            steps,
        )
        if compiled_raptor is not None:
            raptor_pim_dir, raptor_pass_timings = compiled_raptor
    elif not reuse_raptor:
        record_failure(
            failures,
            "Skip Raptor PIM compile",
            "Raptor PIM compile was skipped because the ONNX model or hardware configuration is not available.",
        )

    if not reuse_raptor and raptor_pim_dir is not None:
        wrote_inputs = try_stage_success(
            failures,
            "Write Raptor inputs",
            write_inputs_to_memory_bin,
            raptor_pim_dir / "memory.bin",
            raptor_pim_dir / "config.json",
            runtime_inputs,
        )
        if wrote_inputs and reference_dir is not None and outputs_desc:
            validation = try_stage(
                failures,
                "Functional Validation Raptor",
                run_functional_validation,
                "Functional Validation Raptor",
                raptor_pim_dir,
                raptor_pim_dir / "config.json",
                out_dir / "simulation/out.bin",
                outputs_desc,
                reference_dir,
                steps,
                args,
            )
            raptor_validation = validation if validation is not None else failed_validation("Raptor validation failed")
        elif reference_dir is None:
            raptor_validation = skipped_validation("Reference outputs are not available")
        elif not outputs_desc:
            raptor_validation = skipped_validation("Output descriptors are not available")
        else:
            raptor_validation = skipped_validation("Raptor input materialization failed")
    elif not reuse_raptor:
        raptor_validation = skipped_validation("Raptor PIM compilation did not produce a PIM directory")

    pimcomp_model_path = try_stage(
        failures,
        "Prepare PIMCOMP model",
        prepare_pimcomp_model,
        model_path,
        out_dir / "pimcomp/model",
    )

    if args.reuse_pimcomp_dir is not None:
        reused_pimcomp_dir = args.reuse_pimcomp_dir.resolve()
        copied_pimcomp = try_stage_success(
            failures,
            "Reuse PIMCOMP outputs",
            copy_pimcomp_outputs,
            reused_pimcomp_dir,
            out_dir / "pimcomp",
        )
        if copied_pimcomp:
            verification_info = out_dir / "pimcomp/VerificationInfo.json"
            simulation_info = out_dir / "pimcomp/SimulationInfo.gz"
            print_step("Reuse PIMCOMP")
            print(f"  Directory: {reused_pimcomp_dir}")
    else:
        compiled_pimcomp = try_stage(
            failures,
            "Compile PIMCOMP",
            compile_pimcomp,
            args,
            pimcomp_model_path,
            out_dir / "pimcomp",
            steps,
        ) if pimcomp_model_path is not None else None
        if compiled_pimcomp is not None:
            verification_info, simulation_info = compiled_pimcomp

    if verification_info is not None and simulation_info is not None and model_io is not None:
        exported = try_stage(
            failures,
            "Export PIMCOMP for Functional Validation",
            export_pimcomp_for_rust,
            pimcomp_model_path,
            verification_info,
            simulation_info,
            runtime_inputs,
            out_dir / "pimcomp/exported",
        )
        if exported is not None:
            pimcomp_export_dir = exported
    elif verification_info is None or simulation_info is None:
        record_failure(
            failures,
            "Export PIMCOMP for Functional Validation",
            "PIMCOMP functional export failed because PIMCOMP did not produce VerificationInfo.json and SimulationInfo.gz.",
        )
    else:
        record_failure(
            failures,
            "Export PIMCOMP for Functional Validation",
            "PIMCOMP functional export failed because model inputs are not available.",
        )

    if pimcomp_export_dir is not None and reference_dir is not None and outputs_desc:
        validation = try_stage(
            failures,
            "Functional Validation PIMCOMP",
            run_functional_validation,
            "Functional Validation PIMCOMP",
            pimcomp_export_dir,
            pimcomp_export_dir / "config.json",
            out_dir / "simulation/pimcomp.out.bin",
            outputs_desc,
            reference_dir,
            steps,
            args,
            channel_last=True,
        )
        pimcomp_validation = validation if validation is not None else failed_validation("PIMCOMP validation failed")
    elif pimcomp_export_dir is None:
        pimcomp_validation = failed_validation("PIMCOMP functional export is not available")
    elif reference_dir is None:
        pimcomp_validation = failed_validation("Reference outputs are not available")
    else:
        pimcomp_validation = failed_validation("Output descriptors are not available")

    if not args.skip_pimsim_nn and hardware["core_count"] > 0:
        written_config = try_stage(
            failures,
            "Select pimsim-nn config",
            select_pimsim_config,
            args,
            hardware,
        )
        if written_config is not None:
            pimsim_config = written_config
    elif not args.skip_pimsim_nn:
        record_failure(
            failures,
            "Skip pimsim-nn config",
            "pimsim-nn config was skipped because the hardware configuration is not available.",
        )

    if args.skip_pimsim_nn:
        if not reuse_raptor:
            raptor_perf = skipped_perf("Skipped by --skip-pimsim-nn")
        pimcomp_perf = skipped_perf("Skipped by --skip-pimsim-nn")
    elif pimsim_config is None:
        if not reuse_raptor:
            raptor_perf = skipped_perf("pimsim-nn config is not available")
        pimcomp_perf = skipped_perf("pimsim-nn config is not available")
    else:
        if not reuse_raptor and raptor_pim_dir is not None:
            raptor_pimsim_dir = try_stage(
                failures,
                "Export Raptor for pimsim-nn",
                export_raptor_latency_artifact,
                raptor_pim_dir,
                out_dir / "raptor/pimsim_nn",
            )
            if raptor_pimsim_dir is not None:
                perf = try_stage(
                    failures,
                    "Non-Functional Simulation Raptor",
                    run_pimsim_nn,
                    "Non-Functional Simulation Raptor",
                    raptor_pimsim_dir,
                    pimsim_config,
                    steps,
                    args,
                )
                raptor_perf = perf if perf is not None else failed_perf("pimsim-nn Raptor failed")
            else:
                raptor_perf = failed_perf("Raptor pimsim-nn export failed")
        elif not reuse_raptor:
            raptor_perf = skipped_perf("Raptor PIM directory is not available")

        if simulation_info is not None:
            pimcomp_pimsim_dir = try_stage(
                failures,
                "Export PIMCOMP for pimsim-nn",
                export_pimcomp_for_pimsim_nn,
                simulation_info,
                out_dir / "pimcomp/pimsim_nn",
            )
            if pimcomp_pimsim_dir is not None:
                perf = try_stage(
                    failures,
                    "Non-Functional Simulation PIMCOMP",
                    run_pimsim_nn,
                    "Non-Functional Simulation PIMCOMP",
                    pimcomp_pimsim_dir,
                    pimsim_config,
                    steps,
                    args,
                )
                pimcomp_perf = perf if perf is not None else failed_perf("pimsim-nn PIMCOMP failed")
            else:
                pimcomp_perf = failed_perf("PIMCOMP pimsim-nn export failed")
        else:
            pimcomp_perf = skipped_perf("PIMCOMP SimulationInfo.gz is not available")

    if not reuse_raptor and raptor_pim_dir is not None and raptor_pim_dir.exists():
        parsed = try_stage(failures, "Parse Raptor instructions", parse_raptor_instructions, raptor_pim_dir)
        raptor_instr = parsed if parsed is not None else empty_instruction_summary(error="Failed to parse Raptor instructions")
    elif not reuse_raptor:
        raptor_instr = empty_instruction_summary("Raptor PIM directory is not available")

    if simulation_info is not None and simulation_info.exists():
        parsed = try_stage(failures, "Parse PIMCOMP instructions", parse_pimcomp_instructions, simulation_info)
        pimcomp_instr = parsed if parsed is not None else empty_instruction_summary(error="Failed to parse PIMCOMP instructions")
    else:
        pimcomp_instr = empty_instruction_summary("PIMCOMP SimulationInfo.gz is not available")

    report_path = out_dir / "pimcomp/comparison_report.md"
    write_report(
        report_path,
        model_path=model_path,
        hardware=hardware,
        steps=steps,
        failures=failures,
        raptor_validation=raptor_validation,
        pimcomp_validation=pimcomp_validation,
        raptor_perf=raptor_perf,
        pimcomp_perf=pimcomp_perf,
        raptor_instr=raptor_instr,
        pimcomp_instr=pimcomp_instr,
        raptor_pass_timings=raptor_pass_timings,
        pimsim_mode=args.pimsim_mode,
        pimcomp_pipeline=args.pimcomp_pipeline,
        pimcomp_replication=args.pimcomp_replication,
    )

    json_report = {
        "model": str(model_path),
        "hardware": hardware,
        "pimsim_mode": args.pimsim_mode,
        "pimcomp_pipeline": args.pimcomp_pipeline,
        "pimcomp_replication": args.pimcomp_replication,
        "reused_raptor_report": optional_path(args.reuse_raptor_report.resolve()) if reuse_raptor else None,
        "failures": failures,
        "steps": [asdict(step) for step in steps],
        "raptor_validation": asdict(raptor_validation),
        "pimcomp_validation": asdict(pimcomp_validation),
        "raptor_performance": raptor_perf,
        "pimcomp_performance": pimcomp_perf,
        "raptor_instruction_summary": raptor_instr,
        "pimcomp_instruction_summary": pimcomp_instr,
        "raptor_pass_timings": raptor_pass_timings,
        "paths": {
            "reference_outputs": optional_path(reference_dir),
            "raptor_pim": optional_path(raptor_pim_dir),
            "raptor_pimsim_nn": optional_path(raptor_pimsim_dir),
            "pimcomp_simulation_info": optional_path(simulation_info),
            "pimcomp_exported_pim": optional_path(pimcomp_export_dir),
            "pimsim_config": optional_path(pimsim_config),
            "report_markdown": str(report_path),
        },
    }
    json_path = out_dir / "pimcomp/comparison_report.json"
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(json_report, f, indent=2)
        f.write("\n")

    failed_steps = any(step.status != "passed" for step in steps)
    functional_failure = any(
        result.status == "done" and not result.passed
        for result in (raptor_validation, pimcomp_validation)
    )
    failed = bool(failures or failed_steps or functional_failure)
    result = "FAIL" if args.fail_on_error and failed else "DONE" if failed else "PASS"
    color = Fore.RED if result == "FAIL" else Fore.YELLOW if result == "DONE" else Fore.GREEN
    print("\n" + Style.BRIGHT + f"Result: {color}{result}" + Style.RESET_ALL)
    print(f"  Report: {report_path}")
    print(f"  JSON:   {json_path}")
    if failures or failed_steps:
        print(
            Fore.YELLOW
            + f"  Completed with {len(failures)} recorded failure/skipped stage(s)."
            + Style.RESET_ALL
        )

    if args.fail_on_error and failed:
        raise SystemExit(1)


if __name__ == "__main__":
    main()