"""Spatial selection lowers one winner and rejects incomplete comparisons."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import shutil
import subprocess
import tempfile

import onnx
from onnx import TensorProto, helper

REPO = Path(__file__).resolve().parents[2]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--compiler", required=True)
    args = parser.parse_args()
    with tempfile.TemporaryDirectory(prefix="raptor-spatial-selection-") as temporary:
        root = Path(temporary)
        guard = root / "static-audit.so"
        subprocess.run([os.environ.get("CC", "cc"), "-shared", "-fPIC",
            str(REPO / "test/PIM/StaticSchedulingAudit.c"), "-o", str(guard)], check=True)
        env = dict(os.environ, LD_PRELOAD=str(guard), OMP_NUM_THREADS="1")
        iteration = 0

        def compile(source, pipeline=2, emit="EmitPimCodegen", extra=(), succeeds=True):
            nonlocal iteration
            output = root / str(iteration)
            iteration += 1
            output.mkdir()
            command = [args.compiler, str(source), "-o", str(output / "model"), "--maccel=PIM",
                f"--{emit}", f"--pipeline={pipeline}", "--core-count=138", "--crossbar-count=128",
                "--crossbar-size=128", "--mlir-disable-threading", "--pim-detect-communication-deadlock",
                f"--pim-target-config={REPO}/validation/pimsim_configs/pimcomp/arch-b/throughput_config_1000ms.json"]
            command.extend(extra)
            result = subprocess.run(command, env=env, capture_output=True, text=True, timeout=120)
            assert (result.returncode == 0) == succeeds, result.stdout + result.stderr
            path = output / "reports/pipeline_selection.json"
            report = json.loads(path.read_text()) if path.exists() else None
            program = {str(p.relative_to(output / "pim")): hashlib.sha256(p.read_bytes()).hexdigest()
                for p in (output / "pim").rglob("*") if p.is_file()}
            if "--pim-export-spatial-dataflow=all" in extra:
                for stage in ("spatial3_scheduled_no_comm", "spatial4_scheduled"):
                    for suffix in ("nodes.csv", "edges.csv"):
                        assert (output / f"reports/{stage}.{suffix}").stat().st_size > 0
                assert (output / "dialects/spatial4_scheduled.mlir").exists()
                assert (output / "dialects/pim1_buff.mlir").exists()
            shutil.rmtree(output)
            return report, program, result.stdout + result.stderr

        vgg = REPO / "validation/networks/pimcomp_models/vgg8/vgg8-mnist-reconstructed.onnx"
        selected, program, _ = compile(vgg)
        count = selected["generated_candidates"]
        assert count > 1 and selected["strategy"] == "spatial-queue-aware"
        assert selected["lowered_candidates"] == 1
        assert len(selected["backend_attempts"]) == 1
        assert selected["backend_attempts"][0]["candidate"] == selected["selection"]
        assert not selected["prediction_error_bound_available"]
        assert selected["selection_margin"] == 0.021
        for candidate in selected["candidates"]:
            assert candidate["cost_available"] and candidate["convergence_proved"]
            assert candidate["lowering_supported"] and candidate["eligible"]
            assert "request_order" not in candidate["spatial_cost"]
            assert "without_local_copy_work" not in candidate["spatial_cost"]
        diagnostic, diagnostic_program, _ = compile(vgg, extra=(
            "--pim-export-spatial-dataflow=all", "--pim-dump-ir", "--pim-periodic-diagnostics"))
        assert diagnostic_program == program and diagnostic["selection"] == selected["selection"]

        # A single schedule needs no optional timing contract.
        source = root / "softmax.onnx"
        tensor = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 8])
        output = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 8])
        model = helper.make_model(helper.make_graph([helper.make_node("Softmax", ["x"], ["y"], axis=1)],
            "softmax", [tensor], [output]), opset_imports=[helper.make_opsetid("", 18)])
        model.ir_version = 10
        onnx.save(model, source)
        selected, program, _ = compile(source)
        assert selected["single_candidate_analysis_skipped"] and selected["lowered_candidates"] == 1
        assert "spatial_cost" not in selected["candidates"][0]

        # An unknown estimate cannot silently exclude any alternative.
        model = onnx.load(vgg)
        output = model.graph.output[0].name
        internal = output + "_before_softmax"
        for node in model.graph.node:
            for index, value in enumerate(node.output):
                if value == output:
                    node.output[index] = internal
            for index, value in enumerate(node.input):
                if value == output:
                    node.input[index] = internal
        model.graph.node.append(helper.make_node("Softmax", [internal], [output], axis=-1))
        source = root / "vgg-softmax.onnx"
        onnx.save(model, source)
        report, program, diagnostics = compile(source, succeeds=False)
        assert report["generated_candidates"] > 1 and report["status"] == "incomplete-costs"
        assert "cannot compare the complete candidate set" in diagnostics
        assert "unsupported abstract operation" in diagnostics
        assert not program and "backend_attempts" not in report
        # Latency PEFT and earlier emission do not invoke candidate costing.
        report, program, _ = compile(vgg, pipeline=1)
        assert report is None and program
        report, program, stderr = compile(vgg, emit="EmitSpatial")
        assert report is None and not program
        assert "intermediate emission retains the original schedule" in stderr


if __name__ == "__main__":
    main()