from __future__ import annotations

import argparse
import copy
import csv
import json
import os
from pathlib import Path
import sys

import numpy as np
import onnx
import pytest


TOOLS = Path(__file__).resolve().parents[1] / "tools/pim"
sys.path.insert(0, str(TOOLS / "ablation"))
sys.path.insert(0, str(TOOLS / "pimcomp/compare"))

import run_ablation as ablation
import run_pimcomp_models as suite
import compare_raptor_pimcomp_model as comparator
from raptor_validation.validate_one import validate_outputs


@pytest.mark.parametrize("mode", ["latency", "throughput"])
def test_comparison_report_energy_is_per_inference(tmp_path, mode):
    report = tmp_path / "comparison_report.md"
    validation = comparator.CompareResult(True, {"output": 1e-7}, relative_l2_errors={"output": 1e-6})
    metrics = {"average_power_mw": 2, "average_energy_pj": 30,
               "latency_ms": 4, "average_latency_ms": 4,
               "throughput": 250, "output_count": 6}
    comparator.write_report(
        report, model_path=tmp_path / "model.onnx", hardware={}, steps=[], failures=[],
        raptor_validation=validation, pimcomp_validation=validation,
        raptor_perf=metrics, pimcomp_perf=metrics, raptor_instr={}, pimcomp_instr={},
        raptor_pass_timings={}, pimsim_mode=mode, pimcomp_pipeline="batch",
        pimcomp_replication="weight")
    text = report.read_text()
    assert "Average power (mW) | Energy per inference (pJ/inference)" in text
    assert "Avg latency" not in text
    assert "Avg energy" not in text
    assert metrics["average_energy_pj"] == 30
    assert "Raptor relative L2 | Pimcomp relative L2" in text
    assert "1.000000e-06 | 1.000000e-06" in text


@pytest.mark.parametrize("zero_output", [False, True])
def test_comparison_frontends_reject_missing_small_outputs(tmp_path, zero_output):
    expected = np.array([[1e-4, -2e-4]], dtype=np.float32)
    actual = np.zeros_like(expected) if zero_output else expected * np.float32(1 + 1e-6)
    np.savetxt(tmp_path / "output0_output.csv", expected.reshape(-1), delimiter=",")
    actual.tofile(tmp_path / "output.bin")
    outputs = [(0, "output", 1, [1, 2])]
    result = comparator.compare_simulator_outputs(
        tmp_path / "output.bin", outputs, tmp_path, threshold=1e-3, rtol=1e-4)
    assert result.passed == (not zero_output)
    assert result.relative_l2_errors["output"] == pytest.approx(1 if zero_output else 1e-6, rel=.1)
    assert validate_outputs([actual], tmp_path, outputs, 1e-3, 1e-4, False) == (not zero_output)


def test_comparison_preserves_double_reference_precision(tmp_path):
    expected = np.array([1 + 2**-40], dtype=np.float64)
    np.savetxt(tmp_path / "output0_output.csv", expected, delimiter=",")
    expected.tofile(tmp_path / "output.bin")
    outputs = [(0, "output", 11, [1])]
    assert comparator.compare_simulator_outputs(
        tmp_path / "output.bin", outputs, tmp_path, threshold=0, rtol=0).passed
    assert validate_outputs([expected], tmp_path, outputs, 0, 0, False)


def test_relative_agreement_only_reports_possible_numerical_sensitivity(tmp_path):
    expected = np.array([1e6, 0], dtype=np.float32)
    actual = np.array([1e6, .01], dtype=np.float32)
    np.savetxt(tmp_path / "output0_output.csv", expected, delimiter=",")
    actual.tofile(tmp_path / "output.bin")
    result = comparator.compare_simulator_outputs(
        tmp_path / "output.bin", [(0, "output", 1, [2])], tmp_path, threshold=1e-3, rtol=1e-4)
    assert not result.passed
    assert result.relative_l2_errors["output"] < 1e-4
    assert "possible numerical sensitivity, not proof of correctness" in result.error


def test_comparison_rejects_missing_output_arrays(tmp_path, monkeypatch):
    outputs = [(0, "output", 1, [1])]
    monkeypatch.setattr(comparator, "parse_pim_simulator_outputs", lambda *args: [])
    with pytest.raises(ValueError, match="output count"):
        comparator.compare_simulator_outputs(tmp_path / "missing.bin", outputs, tmp_path, threshold=1e-3, rtol=1e-4)
    with pytest.raises(ValueError, match="output count"):
        validate_outputs([], tmp_path, outputs, 1e-3, 1e-4, False)


@pytest.mark.parametrize("raptor_pass,pimcomp_pass,winner", [
    (True, True, "pimcomp"), (True, False, "NA"),
    (False, True, "NA"), (False, False, "NA"),
])
def test_timing_winner_requires_both_functional_passes(tmp_path, raptor_pass, pimcomp_pass, winner):
    report = suite.result_dir(tmp_path, "vgg8", "arch-a", "latency", 1) / "pimcomp/comparison_report.json"
    report.parent.mkdir(parents=True)
    data = {"validation_policy": suite.VALIDATION_POLICY}
    for side, passed, latency in [("raptor", raptor_pass, 10), ("pimcomp", pimcomp_pass, 1)]:
        data[side + "_validation"] = {"status": "done", "passed": passed}
        data[side + "_performance"] = {"latency_ms": latency, "output_count": 1}
    report.write_text(json.dumps(data))
    output = suite.write_results_csv(tmp_path, "arch-a", ["vgg8"], (("latency", 1, "element"),))
    with output.open() as stream:
        row = next(csv.DictReader(stream))
    assert row["better_compiler"] == winner
    assert row["speedup"] == "10.00"  # Diagnostic timing remains available.


def test_ablation_results_stay_in_independent_output_roots(tmp_path):
    variants = ablation.selected_variants(["no-plan"])
    outputs = []
    for name, value in (("first", 80), ("second", 60)):
        root = tmp_path / name
        for variant in variants:
            path = ablation.variant_output_dir(root, variant) / "results_comparison.csv"
            path.parent.mkdir(parents=True)
            row = {"arch": "arch-a", "model": "vgg8", "mode": "latency",
                   "raptor_pipeline": "1", "raptor_functional_validation": "PASS"}
            row.update({field: value if variant.name == "baseline" else 100
                        for field, _ in ablation.PERCENTAGE_FIELDS})
            with path.open("w", newline="") as stream:
                writer = csv.DictWriter(stream, fieldnames=list(row))
                writer.writeheader()
                writer.writerow(row)
        output, failures = ablation.aggregate_results(root, variants)
        assert failures == []
        assert output == root / "results_ablation.csv"
        outputs.append(output)
    for output, change in zip(outputs, ("+20.00%", "+40.00%")):
        with output.open(newline="") as stream:
            rows = list(csv.DictReader(stream))
        assert len(rows) == 1
        assert rows[0]["comparison"] == "plan"
        assert rows[0]["baseline_variant"] == "no-plan"
        assert rows[0]["latency_gain_percent"] == change


@pytest.mark.parametrize("options,expected", [
    ([], {("baseline", "latency", "1"), ("no-plan", "latency", "1"),
          ("baseline", "throughput", "4"), ("no-plan", "throughput", "4"),
          ("no-sync", "throughput", "4")}),
    (["--mode", "latency"], {("baseline", "latency", "1"), ("no-plan", "latency", "1")}),
    (["--pipeline=8"], {(variant, "throughput", "8") for variant in ("baseline", "no-plan", "no-sync")}),
    (["--mode", "latency", "throughput", "--pipeline", "2"],
     {(variant, "throughput", "2") for variant in ("baseline", "no-plan", "no-sync")}),
])
def test_ablation_runs_only_required_variants(options, expected, capsys):
    import shlex

    assert ablation.main(["--dry-run", *options]) == 0
    commands = [shlex.split(line) for line in capsys.readouterr().out.splitlines()]
    actual = set()
    for command in commands:
        variant = command[command.index("--ablation-variant") + 1]
        actual.add((variant, *ablation.case_key(command)))
        assert "--raptor-only" in command
        assert ("--raptor-extra-arg=--pim-disable-synchronization" in command) == (variant == "no-sync")
        assert ("--raptor-extra-arg=--pim-disable-spatial-planning" in command) == (variant == "no-plan")
    assert actual == expected
    assert len(commands) == len(expected)


def test_ablation_rejects_combined_disabled_variant():
    with pytest.raises(SystemExit):
        ablation.parse_args(["--variant", "no-sync-no-plan"])


def test_ablation_reports_gains_against_each_disabled_feature(tmp_path):
    variants = ablation.selected_variants(None)
    for variant in variants:
        path = ablation.variant_output_dir(tmp_path, variant) / "results_comparison.csv"
        path.parent.mkdir(parents=True)
        rows = []
        for mode, pipeline in (("latency", "1"), ("throughput", "4")):
            if variant.name == "no-sync" and mode == "latency":
                continue
            rows.append({"arch": "arch-b", "model": "vgg8", "mode": mode,
                         "raptor_pipeline": pipeline,
                         "raptor_functional_validation": "FAIL" if variant.name == "no-sync" else "PASS",
                         "raptor_latency_ms": {"baseline": "8", "no-plan": "10", "no-sync": "6.4"}[variant.name],
                         "raptor_throughput_samples_s": {"baseline": "125", "no-plan": "100", "no-sync": "156.25"}[variant.name],
                         "raptor_power_mw": "100", "raptor_energy_pj": "100"})
        with path.open("w", newline="") as stream:
            writer = csv.DictWriter(stream, fieldnames=list(rows[0]))
            writer.writeheader()
            writer.writerows(rows)
    output, failures = ablation.aggregate_results(tmp_path, variants)
    assert not failures
    with output.open(newline="") as stream:
        rows = list(csv.DictReader(stream))
    assert len(rows) == 3
    latency, plan, sync = rows
    assert (latency["comparison"], latency["latency_gain_percent"]) == ("plan", "+20.00%")
    assert (plan["comparison"], plan["throughput_gain_percent"]) == ("plan", "+25.00%")
    assert (sync["comparison"], sync["throughput_gain_percent"]) == ("sync", "-20.00%")
    assert plan["baseline_variant"] == "no-plan"
    assert sync["baseline_variant"] == "no-sync"
    assert sync["baseline_functional_validation"] == "FAIL"
    assert sync["functional_validation"] == "PASS"


@pytest.mark.parametrize("variant", [None, "baseline", "no-plan", "no-sync"])
def test_unconverged_fast_throughput_is_available_in_comparison_and_ablation(tmp_path, variant):
    raw = """Fast throughput convergence failed before sim_time; using legacy full-duration throughput result
Chip Simulation Result:
  - output count:       22 samples
  - throughput:         22 samples/s
  - average latency:    45.45454545 ms
  - average power:      300 mW
  - average energy:     10000000000 pJ/it
"""
    metrics = comparator.parse_pimsim_nn_metrics(raw)
    assert metrics["fast_throughput_converged"] is False
    assert metrics["throughput_measurement"] == "full_duration_fallback"
    assert "unconverged" in comparator.perf_status(metrics)
    report = suite.result_dir(tmp_path, "vgg8", "arch-b", "throughput", 4, variant) / "pimcomp/comparison_report.json"
    report.parent.mkdir(parents=True)
    data = {"validation_policy": suite.VALIDATION_POLICY}
    for side in ("raptor", "pimcomp"):
        data[side + "_validation"] = {"status": "done", "passed": True}
        # Older reports can contain only the raw simulator output.
        data[side + "_performance"] = {"raw_output": raw, "fast_throughput_converged": False}
    report.write_text(json.dumps(data))
    output = suite.write_results_csv(tmp_path, "arch-b", ["vgg8"],
                                     (("throughput", 4, "batch"),), ablation_variant=variant)
    with output.open(newline="") as stream:
        row = next(csv.DictReader(stream))
    for side in ("raptor", "pimcomp"):
        assert float(row[side + "_throughput_samples_s"]) == 22
        assert float(row[side + "_latency_ms"]) == pytest.approx(1000 / 22)
    if variant is not None:
        enabled = ablation.variant_output_dir(tmp_path, ablation.VARIANT_BY_NAME["baseline"]) / "results_comparison.csv"
        enabled.parent.mkdir(parents=True, exist_ok=True)
        enabled.write_text(output.read_text())
        aggregate, failures = ablation.aggregate_results(tmp_path, ablation.selected_variants([variant]))
        assert not failures
        with aggregate.open(newline="") as stream:
            for gain in csv.DictReader(stream):
                assert gain["throughput_gain_percent"] == "+0.00%"


def test_fast_throughput_convergence_metadata():
    metrics = comparator.parse_pimsim_nn_metrics("Fast throughput converged; stopping early\nthroughput: 100 samples/s")
    assert metrics["fast_throughput_converged"] is True
    assert metrics["throughput_measurement"] == "converged_windows"
    assert metrics["throughput"] == 100


@pytest.fixture
def recorded_failure(tmp_path):
    model, functional, config = (tmp_path / name for name in ("model.onnx", "functional.onnx", "config.json"))
    for path in (model, functional, config):
        path.write_text(path.name)
    args = argparse.Namespace(
        pimcomp_config=config, pimsim_mode="throughput", pimsim_time_ms=1000,
        pimcomp_pipeline="batch", pimcomp_replication="GA", pimcomp_model_name=None,
        seed=0, batch_size=64, threshold=1e-3, rtol=1e-4, no_fast=False,
        skip_pimsim_nn=False, pimcomp_dir=tmp_path / "pimcomp",
        pimsim_nn_build_dir=tmp_path / "oracle", pim_simulator_dir=tmp_path / "functional",
        raptor_path=tmp_path / "raptor", raptor_extra_arg=["--pipeline=2"],
    )
    hardware = {"core_count": 168, "crossbar_size": 128}
    context = comparator.pimcomp_reuse_context(args, model, functional, hardware)
    report = {
        "pimcomp_model_source": "original_onnx", "pimcomp_reuse_context": context,
        "pimcomp_validation": {"passed": False, "max_diffs": {}, "status": "failed",
                               "error": "Unsupported model"},
        "pimcomp_performance": {"skipped": True, "reason": "Compilation failed"},
        "pimcomp_instruction_summary": comparator.empty_instruction_summary("No artifact"),
        "paths": {}, "steps": [], "artifact_signatures": {"pimcomp": {}},
        "failures": [{"stage": "Compile Pimcomp frontend", "error": "Unsupported model"}],
    }
    path = tmp_path / "comparison_report.json"
    path.write_text(json.dumps(report))
    return path, report, args, model, functional, hardware


def test_failed_pimcomp_report_needs_no_simulation_artifacts(recorded_failure):
    path, report, *_ = recorded_failure
    reused = comparator.load_reused_pimcomp_report(path, report["pimcomp_reuse_context"])
    assert reused == report
    assert not reused["pimcomp_validation"]["passed"]
    failures = []
    comparator.restore_side_records(reused, "pimcomp", failures, [])
    assert failures == report["failures"]


@pytest.mark.parametrize("tool", ["compiler", "oracle", "exporter"])
def test_reuse_checks_tool_content(recorded_failure, tool):
    path, report, args, model, functional, hardware = recorded_failure
    original = report["pimcomp_reuse_context"]["files"][tool]["path"]
    tool_path = Path(original)
    tool_path.parent.mkdir(parents=True, exist_ok=True)
    tool_path.write_bytes(b"changed executable or exporter")
    expected = comparator.pimcomp_reuse_context(args, model, functional, hardware)
    with pytest.raises(ValueError, match="Cannot reuse Pimcomp report"):
        comparator.load_reused_pimcomp_report(path, expected)


@pytest.mark.parametrize("mutation", ["instructions", "weights", "missing"])
def test_reuse_checks_artifact_content(recorded_failure, mutation):
    path, report, *_ = recorded_failure
    artifact = path.parent / "artifact"
    artifact.mkdir()
    (artifact / "core_0.json").write_text('[{"op":"nop"}]')
    (artifact / "weights.bin").write_bytes(b"weights")
    report["paths"]["pimcomp_exported_pim"] = str(artifact)
    report["artifact_signatures"]["pimcomp"] = comparator.artifact_signatures(report, "pimcomp")
    path.write_text(json.dumps(report))
    assert comparator.load_reused_pimcomp_report(path, report["pimcomp_reuse_context"])
    if mutation == "missing":
        (artifact / "core_0.json").unlink()
    elif mutation == "weights":
        (artifact / "weights.bin").write_bytes(b"wrong weights")
    else:
        (artifact / "core_0.json").write_text('[{"op":"wait"}]')
    with pytest.raises(ValueError, match="artifacts differ"):
        comparator.load_reused_pimcomp_report(path, report["pimcomp_reuse_context"])


def test_raptor_context_covers_pipeline_and_ignores_pimcomp_pipeline(recorded_failure):
    _, _, args, model, functional, hardware = recorded_failure
    expected = comparator.experiment_context(args, model, functional, hardware, "raptor")
    args.pimcomp_pipeline = "element"
    assert expected == comparator.experiment_context(args, model, functional, hardware, "raptor")
    args.raptor_extra_arg = ["--pipeline=8"]
    assert expected != comparator.experiment_context(args, model, functional, hardware, "raptor")


def test_reuse_rejects_previous_validation_policy(recorded_failure):
    path, report, *_ = recorded_failure
    context = copy.deepcopy(report["pimcomp_reuse_context"])
    del report["pimcomp_reuse_context"]["validation_policy"]
    path.write_text(json.dumps(report))
    with pytest.raises(ValueError, match="Cannot reuse Pimcomp report"):
        comparator.load_reused_pimcomp_report(path, context)


def test_raptor_reuse_rejects_previous_validation_policy(recorded_failure, monkeypatch):
    path, _, _, model, _, _ = recorded_failure
    path.write_text('{"validation_policy": "legacy"}')
    monkeypatch.setattr(comparator, "load_effective_hardware", lambda args: {})
    monkeypatch.setattr(comparator, "load_model_inputs", lambda *args: ([], [], []))
    monkeypatch.setattr(sys, "argv", [
        "compare_raptor_pimcomp_model.py", "--model", str(model),
        "--out-dir", str(path.parent / "out"), "--pimcomp-config", str(path.parent / "config.json"),
        "--reuse-raptor-report", str(path),
    ])
    with pytest.raises(ValueError, match="validation policy differs"):
        comparator.main()


@pytest.mark.parametrize("tolerances", [None, {"atol": 0.01, "rtol": 0.0001},
                                      {"atol": 0.001, "rtol": 0.01}])
def test_raptor_reuse_checks_actual_tolerances(recorded_failure, monkeypatch, tolerances):
    path, _, _, model, _, _ = recorded_failure
    path.write_text(json.dumps({"validation_policy": comparator.VALIDATION_POLICY,
                                "validation_tolerances": tolerances}))
    monkeypatch.setattr(comparator, "load_effective_hardware", lambda args: {})
    monkeypatch.setattr(comparator, "load_model_inputs", lambda *args: ([], [], []))
    monkeypatch.setattr(sys, "argv", [
        "compare_raptor_pimcomp_model.py", "--model", str(model),
        "--out-dir", str(path.parent / "out"), "--pimcomp-config", str(path.parent / "config.json"),
        "--reuse-raptor-report", str(path),
    ])
    with pytest.raises(ValueError, match="validation tolerances differ"):
        comparator.main()


def test_failed_raptor_compilation_is_not_reported_as_skipped(tmp_path, monkeypatch):
    model = tmp_path / "model.onnx"
    model.write_bytes(b"test model")
    runner = tmp_path / "runner"
    runner.touch()
    monkeypatch.setattr(comparator, "load_effective_hardware", lambda args: {"core_count": 1})
    monkeypatch.setattr(comparator, "load_model_inputs", lambda *args:
        ([(0, "input", 1, [1])], [(0, "output", 1, [1])], [np.zeros(1, np.float32)]))
    monkeypatch.setattr(comparator, "ensure_reference_runner", lambda *args: runner)
    monkeypatch.setattr(comparator, "generate_reference_outputs", lambda *args, **kwargs: tmp_path)

    def fail_compile(*args):
        raise RuntimeError("unsupported test operation")

    monkeypatch.setattr(comparator, "compile_raptor_target", fail_compile)
    out = tmp_path / "out"
    monkeypatch.setattr(sys, "argv", [
        "compare_raptor_pimcomp_model.py", "--model", str(model), "--out-dir", str(out),
        "--raptor-only", "--skip-pimsim-nn", "--fail-on-error",
    ])
    with pytest.raises(SystemExit) as exit_code:
        comparator.main()
    assert exit_code.value.code == 1
    report = json.loads((out / "pimcomp/comparison_report.json").read_text())
    assert report["raptor_validation"]["status"] == "failed"
    assert not report["raptor_validation"]["passed"]
    assert "unsupported test operation" in report["failures"][0]["error"]


@pytest.mark.parametrize("setting,value", [
    ("pimsim_mode", "latency"), ("pimcomp_pipeline", "element"),
    ("pimsim_time_ms", 2000), ("pimcomp_replication", "balance"),
    ("seed", 1), ("batch_size", 128), ("no_fast", True),
    ("skip_pimsim_nn", True), ("threshold", 0.1), ("rtol", 0.01),
])
def test_reuse_rejects_different_experiment_settings(recorded_failure, setting, value):
    path, report, args, model, functional, hardware = recorded_failure
    setattr(args, setting, value)
    expected = comparator.pimcomp_reuse_context(args, model, functional, hardware)
    with pytest.raises(ValueError, match="Cannot reuse Pimcomp report"):
        comparator.load_reused_pimcomp_report(path, expected)


@pytest.mark.parametrize("changed", ["model", "functional", "config", "hardware", "model_path"])
def test_reuse_rejects_changed_inputs_and_architecture(recorded_failure, changed):
    path, report, args, model, functional, hardware = recorded_failure
    if changed == "hardware":
        hardware["core_count"] = 138
    elif changed == "model_path":
        replacement = model.with_name("other.onnx")
        replacement.write_bytes(model.read_bytes())
        model = replacement
    else:
        {"model": model, "functional": functional, "config": args.pimcomp_config}[changed].write_text("changed")
    expected = comparator.pimcomp_reuse_context(args, model, functional, hardware)
    with pytest.raises(ValueError, match="Cannot reuse Pimcomp report"):
        comparator.load_reused_pimcomp_report(path, expected)


@pytest.mark.parametrize("damage", ["json", "legacy", "validation", "performance", "summary", "steps"])
def test_reuse_rejects_incomplete_reports(recorded_failure, damage):
    path, report, *_ = recorded_failure
    expected = copy.deepcopy(report["pimcomp_reuse_context"])
    if damage == "json":
        path.write_text("{")
    else:
        field = {"legacy": "pimcomp_reuse_context", "validation": "pimcomp_validation",
                 "performance": "pimcomp_performance", "summary": "pimcomp_instruction_summary",
                 "steps": "steps"}[damage]
        report.pop(field)
        path.write_text(json.dumps(report))
    with pytest.raises(ValueError, match="Rerun Pimcomp"):
        comparator.load_reused_pimcomp_report(path, expected)


@pytest.mark.parametrize("only,fresh", [(None, True), (None, False), ("pimcomp", True), ("raptor", False)])
def test_suite_reruns_anchor_and_shares_only_new_reports(tmp_path, monkeypatch, only, fresh):
    root = tmp_path / "suite"
    for _, pipeline, _ in suite.COMPARISONS[1:]:
        report = suite.result_dir(root, "vgg8", "arch-a", "throughput", pipeline) / "pimcomp/comparison_report.json"
        report.parent.mkdir(parents=True)
        report.write_text("{}")
        os.utime(report, ns=(1, 1))
    batches = []

    def run_jobs(jobs, max_workers, log_dir, offset):
        assert max_workers == 1
        batches.append([command for _, command, _ in jobs])
        for _, command, _ in jobs:
            output = Path(command[command.index("--out-dir") + 1]) / "pimcomp/comparison_report.json"
            if fresh or len(batches) > 1:
                output.parent.mkdir(parents=True, exist_ok=True)
                output.write_text('{"new": true}')
        return [], offset + len(jobs)

    summary = tmp_path / "results.csv"
    summary.write_text("summary\n")
    monkeypatch.setattr(suite, "validate_pimcomp_source", lambda: None)
    monkeypatch.setattr(suite, "run", lambda *a, **k: 0)
    monkeypatch.setattr(suite, "run_comparison_jobs", run_jobs)
    monkeypatch.setattr(suite, "write_results_csv", lambda *a: summary)
    monkeypatch.setattr(suite, "comparison_passed", lambda report: True)
    monkeypatch.setattr(sys, "argv", [
        "run_pimcomp_models.py", "--out-dir", str(root), "--models", "vgg8",
        "--archs", "arch-a", "--mode", "throughput", "--jobs", "1",
        *(["--only", only] if only else []),
    ])
    assert suite.main() == 0
    if only == "raptor":
        assert len(batches[0]) == 3
        assert all("--reuse-pimcomp-report" in command for command in batches[0])
    else:
        assert len(batches[0]) == 1
        assert "--reuse-pimcomp-report" not in batches[0][0]
        assert all(("--reuse-pimcomp-report" in command) == fresh for command in batches[1])
        assert ("--reuse-raptor-report" in batches[0][0]) == (only == "pimcomp")


def test_missing_pimcomp_report_is_not_reused(tmp_path):
    assert not suite.pimcomp_report_ready(tmp_path / "output", None)


@pytest.mark.parametrize("model", ["yolo11n", "googlenet"])
def test_comparison_uses_same_model_for_reference_and_execution(tmp_path, monkeypatch, model):
    commands = []
    monkeypatch.setattr(suite, "run", lambda command, **kwargs: commands.append(command) or 0)
    monkeypatch.setattr(sys, "argv", [
        "run_pimcomp_models.py", "--out-dir", str(tmp_path), "--models", model,
        "--archs", "arch-b", "--mode", "throughput", "--pipeline", "8",
        "--raptor-only", "--dry-run",
    ])
    assert suite.main() == 0
    reference, comparison = commands
    expected = str(suite.MODELS[model])
    assert reference[reference.index("--model") + 1] == expected
    assert comparison[comparison.index("--model") + 1] == expected
    assert comparison[comparison.index("--functional-model") + 1] == expected
    assert Path(comparison[comparison.index("--common-dir") + 1]).name == "common-full-model"


def test_yolo_uses_full_model_without_moving_comparison_artifacts():
    assert suite.MODELS["yolo11n"] == suite.REPO / "validation/networks/yolo11n/depth_51/yolo11n_depth_51.onnx"
    assert suite.model_dir(None, "yolo11n") == suite.SUITE / "yolo11n"
    assert suite.common_dir(None, "yolo11n") == suite.SUITE / "yolo11n/artifacts/common-full-model"


def test_googlenet_removes_only_lrn():
    original = onnx.load(suite.SUITE / "googlenet/googlenet-12.onnx")
    actual = onnx.load(suite.MODELS["googlenet"])
    removed = {node.output[0]: node.input[0] for node in original.graph.node if node.op_type == "LRN"}
    assert len(removed) == 2
    retained = []
    for node in original.graph.node:
        if node.op_type == "LRN":
            continue
        for i, value in enumerate(node.input):
            while value in removed:
                value = removed[value]
            node.input[i] = value
        retained.append(copy.deepcopy(node))
    del original.graph.node[:]
    original.graph.node.extend(retained)
    assert actual.SerializeToString() == original.SerializeToString()
    assert sum(node.op_type == "Softmax" for node in actual.graph.node) == 1
    assert [output.name for output in actual.graph.output] == ["prob_1"]
    onnx.checker.check_model(actual)
    assert suite.common_dir(None, "googlenet").name == "common-full-model"


@pytest.mark.parametrize("model,counts", [
    ("yolo11n", {"Softmax": 2, "Split": 11, "Resize": 2}),
    ("resnet18", {"BatchNormalization": 20, "Conv": 20}),
    ("resnet34", {"BatchNormalization": 36, "Conv": 36}),
    ("vgg8", {"Conv": 6, "Gemm": 2, "Relu": 7, "MaxPool": 3}),
])
def test_comparison_models_retain_required_operations(model, counts):
    graph = onnx.load(suite.MODELS[model]).graph
    for operation, expected in counts.items():
        assert sum(node.op_type == operation for node in graph.node) == expected


def test_missing_pimcomp_report_is_rejected_at_load(recorded_failure):
    path, report, *_ = recorded_failure
    with pytest.raises(ValueError, match="Cannot reuse Pimcomp report"):
        comparator.load_reused_pimcomp_report(path.with_name("missing.json"), report["pimcomp_reuse_context"])


def test_canonical_default_result_locations():
    assert suite.DEFAULT_RESULTS_PATH == TOOLS / "pimcomp/compare/results_comparison.csv"
    assert ablation.DEFAULT_RESULTS_PATH == TOOLS / "ablation/results_ablation.csv"


def test_comparison_default_and_custom_aggregate_paths(tmp_path, monkeypatch):
    canonical = tmp_path / "scripts/results_comparison.csv"
    artifact_root = tmp_path / "models"
    monkeypatch.setattr(suite, "DEFAULT_RESULTS_PATH", canonical)
    monkeypatch.setattr(suite, "SUITE", artifact_root)
    assert suite.write_results_csv(None, "arch-a", []) == canonical
    original = canonical.read_bytes()
    for name in ("first", "second"):
        custom = tmp_path / name
        assert suite.write_results_csv(custom, "arch-a", []) == custom / "results_comparison.csv"
    assert canonical.read_bytes() == original
    variant = suite.write_results_csv(None, "arch-a", [], ablation_variant="baseline")
    assert variant == artifact_root / "ablation/baseline/results_comparison.csv"
    suite.clean_artifacts(None, [])
    assert canonical.read_bytes() == original


@pytest.mark.parametrize("custom", [False, True])
def test_ablation_cli_selects_canonical_or_custom_output(tmp_path, monkeypatch, custom):
    artifact_root = tmp_path / "models"
    canonical = tmp_path / "scripts/results_ablation.csv"
    monkeypatch.setattr(ablation, "DEFAULT_OUT_DIR", artifact_root)
    monkeypatch.setattr(ablation, "DEFAULT_RESULTS_PATH", canonical)
    monkeypatch.setattr(ablation.subprocess, "Popen", lambda *a, **k: argparse.Namespace(wait=lambda: 0))
    calls = []
    monkeypatch.setattr(ablation, "aggregate_results", lambda root, *a, **k: (calls.append((root, k["output"])) or None, []))
    command = ["--variant", "baseline", "--jobs", "1"]
    if custom:
        command += ["--out-dir", str(tmp_path / "custom")]
    assert ablation.main(command) == 0
    assert calls == [(tmp_path / "custom", None) if custom else (artifact_root, canonical)]


@pytest.mark.parametrize("value", [None, "", "NA", "0", "0.00", "-1", "nan", "inf"])
def test_ablation_percentages_reject_unavailable_or_invalid_values(value):
    assert ablation.percentage_difference(value, "100") == "NA"
    assert ablation.percentage_difference("100", value) == "NA"


def test_tiny_positive_measurements_and_extreme_changes_are_preserved():
    value = 1e-12
    assert float(suite.format_value(value)) == value
    assert ablation.percentage_difference(suite.format_value(value), "100") == "-100.00%"
    assert ablation.percentage_difference("1000000", "1") == "+99999900.00%"


@pytest.mark.parametrize("status", [{"error": "timeout"}, {"skipped": True}, {"unsupported": True}])
def test_failed_performance_never_uses_stale_numeric_fields(status):
    metrics = {"latency_ms": 10, "throughput": 100, "average_power_mw": 20,
               "average_energy_pj": 30, **status}
    assert all(value is None for value in suite.performance_values(metrics).values())


@pytest.mark.parametrize("value", [None, "NA", 0, -1, float("nan"), float("inf")])
def test_comparison_aggregation_rejects_invalid_measurements(value):
    metrics = suite.performance_values({"latency_ms": value, "throughput": value,
                                       "average_power_mw": value, "average_energy_pj": value})
    assert all(number is None for number in metrics.values())


@pytest.mark.parametrize("primary", [None, "NA", 0, -1, float("nan"), float("inf"), -float("inf")])
def test_latency_uses_valid_fallback_after_invalid_primary(primary):
    metrics = suite.performance_values({"latency_ms": primary, "average_latency_ms": "0.0063"})
    assert metrics["latency"] == 0.0063


def test_latency_preserves_candidate_priority():
    metrics = suite.performance_values({"latency_ms": "0.014", "average_latency_ms": 0.015})
    assert metrics["latency"] == 0.014


@pytest.mark.parametrize("primary,fallback", [
    (None, "NA"), ("NA", None), (0, -1), (-1, 0),
    (float("nan"), float("inf")), (float("inf"), -float("inf")),
])
def test_latency_is_missing_when_all_candidates_are_invalid(primary, fallback):
    metrics = suite.performance_values({"latency_ms": primary, "average_latency_ms": fallback})
    assert metrics["latency"] is None


@pytest.mark.parametrize("invalid", [None, "NA", 0, -1, float("nan"), float("inf"), -float("inf")])
@pytest.mark.parametrize("first_valid", range(5))
def test_latency_priority_across_structured_and_parsed_metrics(monkeypatch, invalid, first_valid):
    candidates = [invalid if index < first_valid else index + 1 for index in range(4)]
    parsed = dict(zip(("latency_ms", "average_latency_ms"), candidates[2:]))
    monkeypatch.setattr(suite, "parse_pimsim_nn_metrics", lambda raw: parsed)
    performance = dict(zip(("latency_ms", "average_latency_ms"), candidates[:2]))
    assert suite.performance_values(performance)["latency"] == (first_valid + 1 if first_valid < 4 else None)


def test_structured_average_latency_precedes_parsed_latency(monkeypatch):
    monkeypatch.setattr(suite, "parse_pimsim_nn_metrics", lambda raw: {"latency_ms": 1})
    assert suite.performance_values({"average_latency_ms": 2})["latency"] == 2


@pytest.mark.parametrize("field,metric", [
    ("throughput", "throughput"), ("average_power_mw", "power"), ("average_energy_pj", "energy"),
])
@pytest.mark.parametrize("structured,parsed,expected", [
    (10, 20, 10), (None, 20, 20), ("NA", 20, 20), (0, 20, 20), (-1, 20, 20),
    (float("nan"), 20, 20), (float("inf"), 20, 20), (-float("inf"), 20, 20),
    (0, float("nan"), None),
])
def test_measurement_priority_across_sources(monkeypatch, field, metric, structured, parsed, expected):
    monkeypatch.setattr(suite, "parse_pimsim_nn_metrics", lambda raw: {field: parsed})
    assert suite.performance_values({field: structured})[metric] == expected


@pytest.mark.parametrize("status", [
    {"error": "timeout"}, {"skipped": True}, {"unsupported": True},
    {"output_count": 0},
])
def test_parsed_measurements_respect_simulation_status(monkeypatch, status):
    parsed = {"latency_ms": 1, "throughput": 2, "average_power_mw": 3, "average_energy_pj": 4}
    monkeypatch.setattr(suite, "parse_pimsim_nn_metrics", lambda raw: parsed)
    assert suite.performance_values(status) == {
        "latency": None, "throughput": None, "power": 3 if "output_count" in status else None,
        "energy": None,
    }


def test_parsed_zero_outputs_preserve_only_measured_power(monkeypatch):
    parsed = {"output_count": 0, "latency_ms": 1, "throughput": 2,
              "average_power_mw": 3, "average_energy_pj": 4}
    monkeypatch.setattr(suite, "parse_pimsim_nn_metrics", lambda raw: parsed)
    assert suite.performance_values({}) == {"latency": None, "throughput": None, "power": 3, "energy": None}


@pytest.mark.parametrize("value,expected", [
    (1e-12, "1e-12"), (0.00499999, "0.00499999"), (0.005, "0.005"),
    (0.00500001, "0.00500001"), (0.0063, "0.0063"), (0.014, "0.014"),
    (0.015, "0.015"), (1.23456789, "1.23457"), (123.456789, "123.457"),
    (123456789, "1.23457e+08"), (1e20, "1e+20"),
])
def test_measurements_use_six_significant_digits(value, expected):
    assert suite.format_value(value) == expected


@pytest.mark.parametrize("value", [None, 0, -1, float("nan"), float("inf"), -float("inf")])
def test_format_value_preserves_missing_measurements(value):
    assert suite.format_value(value) == "NA"


def test_zero_outputs_preserve_only_measured_average_power():
    metrics = suite.performance_values({"output_count": 0, "average_latency_ms": 10,
                                       "throughput": 100, "average_energy_pj": 30,
                                       "average_power_mw": 20})
    assert metrics == {"latency": None, "throughput": None, "energy": None, "power": 20}


def test_functional_pass_survives_unavailable_performance(tmp_path):
    report = suite.result_dir(tmp_path, "vgg8", "arch-a", "throughput", 4) / "pimcomp/comparison_report.json"
    report.parent.mkdir(parents=True)
    report.write_text(json.dumps({"validation_policy": suite.VALIDATION_POLICY,
                                  "raptor_validation": {"status": "done", "passed": True},
                                  "raptor_performance": {"error": "timeout", "throughput": 0}}))
    output = suite.write_results_csv(tmp_path, "arch-a", ["vgg8"], (("throughput", 4, "batch"),))
    with output.open(newline="") as stream:
        row = next(csv.DictReader(stream))
    assert row["raptor_functional_validation"] == "PASS"
    assert all(row[field] == "NA" for field, _ in ablation.PERCENTAGE_FIELDS)


def test_aggregate_rejects_previous_validation_policy(tmp_path):
    report = suite.result_dir(tmp_path, "vgg8", "arch-a", "throughput", 4) / "pimcomp/comparison_report.json"
    report.parent.mkdir(parents=True)
    report.write_text(json.dumps({"raptor_validation": {"status": "done", "passed": True}}))
    with pytest.raises(ValueError, match="different validation policy"):
        suite.write_results_csv(tmp_path, "arch-a", ["vgg8"], (("throughput", 4, "batch"),))
    assert not (tmp_path / "results_comparison.csv").exists()


def test_ablation_aggregate_keeps_pass_with_zero_performance(tmp_path):
    variants = ablation.selected_variants(["no-plan"])
    for variant in variants:
        path = ablation.variant_output_dir(tmp_path, variant) / "results_comparison.csv"
        path.parent.mkdir(parents=True)
        row = {"arch": "arch-a", "model": "googlenet", "mode": "throughput",
               "raptor_pipeline": "4", "raptor_functional_validation": "PASS"}
        row.update({name: "100" if variant.name == "baseline" else "0.00"
                    for name, _ in ablation.PERCENTAGE_FIELDS})
        with path.open("w", newline="") as stream:
            writer = csv.DictWriter(stream, fieldnames=list(row))
            writer.writeheader()
            writer.writerow(row)
    destination = tmp_path / "canonical/results_ablation.csv"
    output, failures = ablation.aggregate_results(tmp_path, variants, output=destination)
    assert not failures and output == destination
    with output.open(newline="") as stream:
        rows = list(csv.DictReader(stream))
    assert len(rows) == 1
    assert rows[0]["functional_validation"] == "PASS"
    assert rows[0]["baseline_functional_validation"] == "PASS"
    assert all(rows[0][f"{metric}_gain_percent"] == "NA" for _, metric in ablation.PERCENTAGE_FIELDS)