"""Tests for auto-placement in the placement pass.

Tests verify:
- or1-asm.AC10.1: Unplaced nodes are assigned to PEs without exceeding limits
- or1-asm.AC10.2: Explicitly placed nodes are not moved by auto-placement
- or1-asm.AC10.3: Connected nodes prefer co-location on same PE (locality heuristic)
- or1-asm.AC10.4: Program too large for available PEs produces error with per-PE utilization breakdown
"""

from asm.place import place
from asm.ir import IRGraph, IRNode, IREdge, SystemConfig, SourceLoc, IRRegion, RegionKind
import asm.ir
from asm.errors import ErrorCategory
from cm_inst import ArithOp, LogicOp, Port


class TestBasicAutoPlacement:
    """AC10.1: Unplaced nodes are assigned to PEs without exceeding limits."""

    def test_four_unplaced_nodes_two_pes(self):
        """Four unplaced nodes with SystemConfig(pe_count=2) get assigned."""
        nodes = {
            "&a": IRNode(name="&a", opcode=ArithOp.ADD, pe=None, loc=SourceLoc(1, 1)),
            "&b": IRNode(name="&b", opcode=ArithOp.SUB, pe=None, loc=SourceLoc(2, 1)),
            "&c": IRNode(name="&c", opcode=ArithOp.INC, pe=None, loc=SourceLoc(3, 1)),
            "&d": IRNode(name="&d", opcode=ArithOp.DEC, pe=None, loc=SourceLoc(4, 1)),
        }
        system = SystemConfig(pe_count=2, sm_count=1, iram_capacity=64, frame_count=4)
        graph = IRGraph(nodes, system=system)
        result = place(graph)

        # Should have no errors
        assert len(result.errors) == 0, f"Expected no errors, got: {[e.message for e in result.errors]}"

        # All nodes should have PE assignments
        for node_name in nodes.keys():
            assert result.nodes[node_name].pe is not None, f"{node_name} still unplaced"
            assert 0 <= result.nodes[node_name].pe < 2, f"{node_name} on invalid PE"

    def test_monadic_node_placement(self):
        """Monadic nodes take up only 1 IRAM slot."""
        # CONST is monadic (RoutingOp.CONST)
        from cm_inst import RoutingOp
        nodes = {
            "&c1": IRNode(name="&c1", opcode=RoutingOp.CONST, pe=None, const=5, loc=SourceLoc(1, 1)),
            "&c2": IRNode(name="&c2", opcode=RoutingOp.CONST, pe=None, const=10, loc=SourceLoc(2, 1)),
        }
        system = SystemConfig(pe_count=1, sm_count=1, iram_capacity=64, frame_count=4)
        graph = IRGraph(nodes, system=system)
        result = place(graph)

        assert len(result.errors) == 0
        for node in result.nodes.values():
            assert node.pe == 0


class TestExplicitPreserved:
    """AC10.2: Explicitly placed nodes are not moved by auto-placement."""

    def test_mixed_placed_and_unplaced(self):
        """Explicitly placed nodes keep their PE; unplaced ones get assigned."""
        nodes = {
            "&placed0": IRNode(name="&placed0", opcode=ArithOp.ADD, pe=0, loc=SourceLoc(1, 1)),
            "&placed1": IRNode(name="&placed1", opcode=ArithOp.SUB, pe=1, loc=SourceLoc(2, 1)),
            "&unplaced1": IRNode(name="&unplaced1", opcode=ArithOp.INC, pe=None, loc=SourceLoc(3, 1)),
            "&unplaced2": IRNode(name="&unplaced2", opcode=ArithOp.DEC, pe=None, loc=SourceLoc(4, 1)),
        }
        system = SystemConfig(pe_count=2, sm_count=1, iram_capacity=64, frame_count=4)
        graph = IRGraph(nodes, system=system)
        result = place(graph)

        assert len(result.errors) == 0
        # Explicitly placed nodes keep their PE
        assert result.nodes["&placed0"].pe == 0
        assert result.nodes["&placed1"].pe == 1
        # Unplaced nodes get assigned
        assert result.nodes["&unplaced1"].pe is not None
        assert result.nodes["&unplaced2"].pe is not None

    def test_all_explicit_placement(self):
        """All nodes explicitly placed - no auto-placement needed."""
        nodes = {
            "&a": IRNode(name="&a", opcode=ArithOp.ADD, pe=0, loc=SourceLoc(1, 1)),
            "&b": IRNode(name="&b", opcode=ArithOp.SUB, pe=1, loc=SourceLoc(2, 1)),
        }
        system = SystemConfig(pe_count=2, sm_count=1, iram_capacity=64, frame_count=4)
        graph = IRGraph(nodes, system=system)
        result = place(graph)

        assert len(result.errors) == 0
        assert result.nodes["&a"].pe == 0
        assert result.nodes["&b"].pe == 1


class TestLocalityHeuristic:
    """AC10.3: Connected nodes prefer co-location on same PE."""

    def test_two_connected_unplaced_nodes(self):
        """Two connected unplaced nodes end up on the same PE."""
        nodes = {
            "&a": IRNode(name="&a", opcode=ArithOp.ADD, pe=None, loc=SourceLoc(1, 1)),
            "&b": IRNode(name="&b", opcode=ArithOp.SUB, pe=None, loc=SourceLoc(2, 1)),
        }
        edges = [
            IREdge(source="&a", dest="&b", port=Port.L, loc=SourceLoc(3, 1)),
        ]
        system = SystemConfig(pe_count=4, sm_count=1, iram_capacity=64, frame_count=4)
        graph = IRGraph(nodes, edges=edges, system=system)
        result = place(graph)

        assert len(result.errors) == 0
        # Both nodes should be on the same PE
        assert result.nodes["&a"].pe == result.nodes["&b"].pe

    def test_cluster_of_three_interconnected_nodes(self):
        """Cluster of 3 interconnected nodes end up on the same PE."""
        nodes = {
            "&a": IRNode(name="&a", opcode=ArithOp.ADD, pe=None, loc=SourceLoc(1, 1)),
            "&b": IRNode(name="&b", opcode=ArithOp.SUB, pe=None, loc=SourceLoc(2, 1)),
            "&c": IRNode(name="&c", opcode=ArithOp.INC, pe=None, loc=SourceLoc(3, 1)),
        }
        edges = [
            IREdge(source="&a", dest="&b", port=Port.L, loc=SourceLoc(4, 1)),
            IREdge(source="&b", dest="&c", port=Port.L, loc=SourceLoc(5, 1)),
            IREdge(source="&a", dest="&c", port=Port.R, loc=SourceLoc(6, 1)),
        ]
        system = SystemConfig(pe_count=4, sm_count=1, iram_capacity=64, frame_count=4)
        graph = IRGraph(nodes, edges=edges, system=system)
        result = place(graph)

        assert len(result.errors) == 0
        # All three nodes should be on the same PE
        pe_a = result.nodes["&a"].pe
        pe_b = result.nodes["&b"].pe
        pe_c = result.nodes["&c"].pe
        assert pe_a == pe_b == pe_c

    def test_locality_with_mixed_placed_unplaced(self):
        """Unplaced node prefers PE of its connected placed neighbour."""
        nodes = {
            "&placed": IRNode(name="&placed", opcode=ArithOp.ADD, pe=1, loc=SourceLoc(1, 1)),
            "&unplaced": IRNode(name="&unplaced", opcode=ArithOp.SUB, pe=None, loc=SourceLoc(2, 1)),
        }
        edges = [
            IREdge(source="&placed", dest="&unplaced", port=Port.L, loc=SourceLoc(3, 1)),
        ]
        system = SystemConfig(pe_count=4, sm_count=1, iram_capacity=64, frame_count=4)
        graph = IRGraph(nodes, edges=edges, system=system)
        result = place(graph)

        assert len(result.errors) == 0
        # Unplaced node should prefer PE1 (where placed node is)
        assert result.nodes["&unplaced"].pe == 1


class TestOverflow:
    """AC10.4: Program too large produces error with per-PE utilization breakdown."""

    def test_200_nodes_overflow_error(self):
        """200 unplaced nodes with limited resources produce overflow error."""
        # Create 200 monadic nodes (each takes 1 IRAM slot)
        nodes = {}
        for i in range(200):
            nodes[f"&n{i}"] = IRNode(
                name=f"&n{i}",
                opcode=ArithOp.INC,
                pe=None,
                loc=SourceLoc(i + 1, 1),
            )

        system = SystemConfig(pe_count=2, sm_count=1, iram_capacity=64, frame_count=4)
        graph = IRGraph(nodes, system=system)
        result = place(graph)

        # Should have error(s) about placement failure
        assert len(result.errors) > 0
        error = result.errors[0]
        assert error.category == ErrorCategory.PLACEMENT
        assert "Cannot place" in error.message or "full" in error.message.lower()
        # Error message should include per-PE utilization info
        assert "PE0" in error.message or "IRAM" in error.message

    def test_overflow_error_includes_breakdown(self):
        """Overflow error includes per-PE slot utilization breakdown."""
        # 130 monadic nodes: on 2 PEs with 64 IRAM capacity each, this overflows
        nodes = {}
        for i in range(130):
            nodes[f"&n{i}"] = IRNode(
                name=f"&n{i}",
                opcode=ArithOp.INC,
                pe=None,
                loc=SourceLoc(i + 1, 1),
            )

        system = SystemConfig(pe_count=2, sm_count=1, iram_capacity=64, frame_count=4)
        graph = IRGraph(nodes, system=system)
        result = place(graph)

        assert len(result.errors) > 0
        error_msg = result.errors[0].message
        # Should mention PE utilization
        assert any(x in error_msg for x in ["PE", "IRAM", "full", "capacity"])


class TestFunctionScopedNodesPlacement:
    """Verify auto-placed nodes inside function scopes receive PE assignments."""

    def test_function_scoped_node_gets_pe_assignment(self):
        """Nodes inside function regions should get PE assignments after place()."""
        # Create a graph with a function region containing unplaced nodes
        func_nodes = {
            "$main.&add": IRNode(
                name="$main.&add",
                opcode=ArithOp.ADD,
                pe=None,
                loc=SourceLoc(1, 1),
            ),
            "$main.&inc": IRNode(
                name="$main.&inc",
                opcode=ArithOp.INC,
                pe=None,
                loc=SourceLoc(2, 1),
            ),
        }

        func_region = IRGraph(nodes=func_nodes, edges=[], regions=[], data_defs=[], errors=[])

        main_region = asm.ir.IRRegion(
            tag="$main",
            kind=asm.ir.RegionKind.FUNCTION,
            body=func_region,
            loc=SourceLoc(0, 1),
        )

        graph = IRGraph(
            nodes={},
            edges=[],
            regions=[main_region],
            data_defs=[],
            system=SystemConfig(pe_count=2, sm_count=1, iram_capacity=64, frame_count=4),
            errors=[],
        )

        result = place(graph)

        # Verify no errors
        assert len(result.errors) == 0, f"Expected no errors, got: {[e.message for e in result.errors]}"

        # Verify nodes inside the function region received PE assignments
        assert len(result.regions) == 1
        result_func = result.regions[0]
        assert result_func.tag == "$main"

        # Check that nodes in the function body have PE assignments
        for node_name, node in result_func.body.nodes.items():
            assert node.pe is not None, f"Node {node_name} in function scope still has pe=None"
            assert 0 <= node.pe < 2, f"Node {node_name} has invalid PE {node.pe}"