import argparse
from collections import defaultdict, deque
import json
import os
from pathlib import Path
import struct
import subprocess
import sys
import tempfile
from types import SimpleNamespace


REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO / "validation/tools/pim/pimcomp/correctness/synchronization"))
import global_memory as sync


def read_events(artifact):
    streams = {}
    for path in artifact.glob("core_*.pim"):
        core = int(path.stem.split("_")[1])
        registers, events = {}, []
        data = path.read_bytes()
        assert data[:4] == b"PIMB"
        for pc, (op, rd, rs, flags, arg, selection, offset, size) in enumerate(
                struct.iter_unpack("<BBBBiiii", data[12:])):
            if op == 1:
                registers[rd] = arg
            elif op in (25, 26):
                register, bit = (rs, 2) if op == 25 else (rd, 1)
                address = registers[register] + (offset if selection & bit else 0)
                events.append((pc, op, address, size))
            elif op in (29, 30, 31, 32):
                events.append((pc, op, arg, selection, offset))
        if events:
            streams[core] = events
    return streams


def check_lifetimes(streams):
    graph, signals, waits = defaultdict(list), defaultdict(list), {}
    sends, receives, stores, loads = defaultdict(list), defaultdict(list), {}, defaultdict(list)
    for core, events in streams.items():
        for iteration in range(3):
            nodes = [(core, iteration, event[0]) for event in events]
            for left, right in zip(nodes, nodes[1:]):
                graph[left].append(right)
            if iteration:
                graph[(core, iteration - 1, events[-1][0])].append(nodes[0])
            for node, event in zip(nodes, events):
                pc, op, *args = event
                if op == 32:
                    target, register, _ = args
                    assert 0 <= register < 32, (core, event)
                    signals[target, iteration, register].append(node)
                elif op == 31:
                    _, register, expected = args
                    key = core, iteration, register
                    assert key not in waits and 0 <= register < 32
                    waits[key] = node, expected
                elif op in (29, 30):
                    peer = args[0]
                    key = (core, peer, iteration) if op == 29 else (peer, core, iteration)
                    (sends if op == 29 else receives)[key].append(node)
                elif op == 26:
                    key = iteration, *args
                    assert key not in stores, ("duplicate slot store", key)
                    stores[key] = node
                elif op == 25:
                    loads[iteration, *args].append(node)
    assert signals.keys() == waits.keys()
    for key, (wait, expected) in waits.items():
        assert len(signals[key]) == expected, (key, expected, signals[key])
        for signal in signals[key]:
            graph[signal].append(wait)
    assert sends.keys() == receives.keys()
    for key, send_nodes in sends.items():
        assert len(send_nodes) == len(receives[key])
        for send, receive in zip(send_nodes, receives[key]):
            graph[send].append(receive)
    indegree = defaultdict(int)
    for source, successors in list(graph.items()):
        indegree[source] += 0
        for successor in successors:
            indegree[successor] += 1
    pending = deque(node for node, degree in indegree.items() if degree == 0)
    visited = 0
    while pending:
        visited += 1
        for successor in graph[pending.popleft()]:
            indegree[successor] -= 1
            if indegree[successor] == 0:
                pending.append(successor)
    assert visited == len(indegree), "communication happens-before cycle"

    def reaches(source, target):
        pending, seen = [source], {source}
        while pending:
            for successor in graph[pending.pop()]:
                if successor == target:
                    return True
                if successor not in seen:
                    seen.add(successor)
                    pending.append(successor)
        return False

    shared, stalled_reader = 0, None
    for (iteration, address, size), writer in stores.items():
        readers = loads.get((iteration, address, size), [])
        if not readers:
            continue
        shared += iteration == 0 and len({reader[0] for reader in readers}) > 1
        if len({reader[0] for reader in readers}) > 1:
            stalled_reader = readers[-1]
        for reader in readers:
            assert reaches(writer, reader), ("missing READY", writer, reader)
            if iteration < 2:
                assert reaches(reader, stores[iteration + 1, address, size]), (
                    "unsafe overwrite", reader, address)
    assert shared, "fixture must exercise a publication with multiple physical readers"
    return shared, stalled_reader


def run_functional(artifact, inputs, references, outputs, out, policy, seed, reader=None):
    out.mkdir()
    command = [str(sync.RUST_BINARY), "--folder", str(artifact),
        "--output", str(out / "output.bin"), "--dump",
        sync.compare.build_dump_ranges(artifact / "config.json", outputs),
        "--mode", "throughput", "--batch-size", "8", "--input-dir", str(inputs[0].parent),
        "--batch-output-dir", str(out / "iterations"),
        "--diagnostic-schedule-policy", policy, "--diagnostic-schedule-seed", str(seed)]
    if reader:
        command += ["--diagnostic-target-stall", f"{reader[0] + 1}:{reader[2]}:1:10000"]
    result = subprocess.run(command, capture_output=True, text=True, timeout=120)
    assert result.returncode == 0, result.stdout + result.stderr
    for index, reference in enumerate(references[:8]):
        result = sync.compare.compare_simulator_outputs(
            out / "iterations" / f"output_{index:06d}.bin", outputs, reference,
            threshold=1e-3, rtol=1e-5)
        assert result.passed, result


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--compiler", required=True)
    args = parser.parse_args()
    with tempfile.TemporaryDirectory(prefix="raptor-publication-") as temporary:
        root = Path(temporary)
        model = REPO / "validation/networks/pimcomp_models/vgg8/vgg8-mnist-reconstructed.onnx"
        options = SimpleNamespace(timeout=120, threshold=1e-3, rtol=1e-5)
        samples, inputs, _, _ = sync.make_inputs(model, 8, 0, root)
        references = sync.make_references(model, samples, root, options)
        _, outputs = sync.compare.onnx_io(model)
        for arch, cores, weights in (("arch-a", 168, 96), ("arch-b", 138, 128)):
            out = root / arch
            out.mkdir()
            command = [args.compiler, str(model), "-o", str(out / "model"),
                "--maccel=PIM", "--EmitPimCodegen", "--crossbar-size=128",
                f"--crossbar-count={weights}", f"--core-count={cores}", "--pipeline=4",
                f"--pim-target-config={sync.CONFIG_ROOT / arch / 'throughput_config_1000ms.json'}",
                "--mlir-disable-threading", "--pim-emit-json", "--pim-detect-communication-deadlock"]
            compiled = subprocess.run(command, capture_output=True, text=True,
                timeout=120, env=dict(os.environ, OMP_NUM_THREADS="1"))
            assert compiled.returncode == 0, compiled.stdout + compiled.stderr
            artifact = out / "pim"
            streams = read_events(artifact)
            shared, reader = check_lifetimes(streams)
            for policy, seed in (("greedy", 0), ("randomized", 1), ("randomized", 7)):
                run_functional(artifact, inputs, references, outputs,
                    out / f"{policy}-{seed}", policy, seed)
            run_functional(artifact, inputs, references, outputs,
                out / "stalled", "randomized", 11, reader)
            print(f"{arch}: {shared} shared publications; three-iteration lifetime proof; eight inputs; randomized and stalled execution passed")


if __name__ == "__main__":
    main()