from __future__ import annotations

import argparse
import copy
import csv
import importlib.util
import json
import os
from pathlib import Path
import sys

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


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: 100 if variant.name == "baseline" else value
                        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 rows[0]["latency"] == "+0.00%"
        assert rows[1]["latency"] == change


@pytest.mark.parametrize("custom_input", [False, True])
def test_ablation_plots_follow_the_selected_experiment(tmp_path, monkeypatch, custom_input):
    spec = importlib.util.spec_from_file_location("ablation_plots", TOOLS / "ablation/generate_plots.py")
    plots = importlib.util.module_from_spec(spec)
    monkeypatch.setitem(sys.modules, spec.name, plots)
    spec.loader.exec_module(plots)
    expected = tmp_path / "experiment/results_ablation.csv" if custom_input else (
        ablation.DEFAULT_RESULTS_PATH
    )
    monkeypatch.setattr(sys, "argv", ["generate_plots.py", *(["--input", str(expected)] if custom_input else [])])
    result = plots.Result("baseline", "arch-a", "vgg8", "latency", 1, "PASS", {})
    inputs, outputs = [], []
    monkeypatch.setattr(plots, "configure_style", lambda: None)
    monkeypatch.setattr(plots, "load_results", lambda path: inputs.append(path) or [result])
    monkeypatch.setattr(plots, "draw_metric", lambda rows, metric, label, output: outputs.append(output))
    plots.main()
    assert inputs == [expected]
    assert outputs == [expected.parent / "plots"] * 3


@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,
    )
    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": [],
        "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("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)


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({"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_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 rows[1]["functional_validation"] == "PASS"
    assert all(rows[1][metric] == "NA" for _, metric in ablation.PERCENTAGE_FIELDS)