"""
Lane-based PE rewrite tests.

Verifies frame-lanes.AC3, frame-lanes.AC4, frame-lanes.AC5, and frame-lanes.AC8:
- AC3.1: FrameOp.ALLOC_SHARED assigns next free lane from parent frame
- AC3.2: FrameOp.FREE_LANE removes tag_store entry, clears lane data, keeps frame
- AC3.3: FrameOp.FREE on shared frame returns lane if frame still in use
- AC3.4: FrameOp.ALLOC unchanged — allocates fresh frame, assigns lane 0
- AC3.5: FrameAllocated event gains lane field
- AC3.6: ALLOC_SHARED with all lanes occupied emits TokenRejected
- AC4: ALLOC_REMOTE reads fref+2 for data-driven ALLOC_SHARED vs ALLOC
- AC5.1: FREE_FRAME opcode uses smart FREE behaviour on shared frames
- AC8.1: Two act_ids sharing a frame have independent matching
- AC8.2: ALLOC_SHARED with exhausted lanes emits TokenRejected
- AC8.3: FREE on shared frame preserves other lanes' data
- AC8.4: ALLOC_REMOTE emits ALLOC_SHARED when fref+2 is non-zero
- AC8.5: ALLOC_REMOTE emits ALLOC when fref+2 is zero (backwards compatible)
- AC8.6: Full loop pipelining scenario — two iterations concurrent on different lanes
"""

import pytest
import simpy

from cm_inst import (
    ArithOp, FrameDest, FrameOp, Instruction, Port, TokenKind, OutputStyle,
    RoutingOp,
)
from emu.events import (
    FrameAllocated, FrameFreed, TokenReceived, TokenRejected, Matched, Emitted,
)
from emu.pe import ProcessingElement
from emu.types import PEConfig
from tokens import DyadToken, FrameControlToken


def inject_and_run(env, pe, token):
    """Helper: inject token and run simulation."""
    def _put():
        yield pe.input_store.put(token)
    env.process(_put())
    env.run()


class TestAllocShared:
    """AC3.1: ALLOC_SHARED assigns next free lane from parent frame."""

    def test_alloc_shared_basic(self):
        """Parent allocates frame, child allocates shared lane."""
        env = simpy.Environment()
        events = []
        config = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Parent ALLOC
        fct_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct_parent)

        parent_frame_id, parent_lane = pe.tag_store[0]
        assert parent_lane == 0, "Parent should allocate lane 0"

        # Child ALLOC_SHARED with parent_act_id=0
        fct_child = FrameControlToken(
            target=0, act_id=1, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_child)

        child_frame_id, child_lane = pe.tag_store[1]
        assert child_frame_id == parent_frame_id, "Child should share parent's frame"
        assert child_lane == 1, "Child should allocate lane 1"
        assert child_lane != parent_lane, "Child lane should differ from parent"

        # Verify FrameAllocated event for child
        frame_allocated = [e for e in events if isinstance(e, FrameAllocated)]
        assert len(frame_allocated) >= 2, "Should have 2 FrameAllocated events"
        assert frame_allocated[0].lane == 0, "Parent allocated lane 0"
        assert frame_allocated[1].lane == 1, "Child allocated lane 1"

    def test_alloc_shared_multiple_lanes(self):
        """Multiple children allocate different lanes from same parent frame."""
        env = simpy.Environment()
        events = []
        config = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Parent ALLOC
        fct_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct_parent)
        parent_frame_id, _parent_lane = pe.tag_store[0]

        # Child 1 ALLOC_SHARED
        fct_child1 = FrameControlToken(
            target=0, act_id=1, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_child1)
        _child1_frame_id, child1_lane = pe.tag_store[1]

        # Child 2 ALLOC_SHARED
        fct_child2 = FrameControlToken(
            target=0, act_id=2, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_child2)
        _child2_frame_id, child2_lane = pe.tag_store[2]

        # All should share same frame
        assert pe.tag_store[0][0] == parent_frame_id
        assert pe.tag_store[1][0] == parent_frame_id
        assert pe.tag_store[2][0] == parent_frame_id

        # Lanes should differ: 0, 1, 2
        assert child1_lane != 0, "Child1 lane should not be 0"
        assert child2_lane != 0, "Child2 lane should not be 0"
        assert child1_lane != child2_lane, "Child1 and child2 lanes should differ"

    def test_alloc_shared_invalid_parent(self):
        """ALLOC_SHARED with non-existent parent emits TokenRejected."""
        env = simpy.Environment()
        events = []
        config = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Try ALLOC_SHARED with non-existent parent_act_id=999
        fct = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC_SHARED, payload=999
        )
        inject_and_run(env, pe, fct)

        rejected = [e for e in events if isinstance(e, TokenRejected)]
        assert len(rejected) > 0, "Should have TokenRejected event"
        assert "not in tag store" in rejected[0].reason, "Reason should mention tag_store"

        # Parent should not be in tag_store
        assert 999 not in pe.tag_store

    def test_alloc_shared_self_referential_guard(self):
        """ALLOC_SHARED with act_id already in tag_store emits TokenRejected."""
        env = simpy.Environment()
        events = []
        config = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # First ALLOC to establish act_id=0 in tag_store
        fct_alloc = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct_alloc)
        assert 0 in pe.tag_store, "act_id=0 should be in tag_store after ALLOC"
        frame_id_0, lane_0 = pe.tag_store[0]

        # Now try ALLOC_SHARED with act_id=0 and payload=1 (parent_act_id=1)
        # This should be rejected because act_id=0 already exists
        fct_alloc_parent = FrameControlToken(
            target=0, act_id=1, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct_alloc_parent)
        assert 1 in pe.tag_store, "act_id=1 should be in tag_store after ALLOC"

        events.clear()
        fct_shared = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC_SHARED, payload=1
        )
        inject_and_run(env, pe, fct_shared)

        rejected = [e for e in events if isinstance(e, TokenRejected)]
        assert len(rejected) > 0, "Should have TokenRejected event"
        assert "already in tag store" in rejected[0].reason, "Reason should mention already in tag store"

        # Frame and lane should be unchanged
        assert pe.tag_store[0] == (frame_id_0, lane_0), "act_id=0 state should be unchanged"


class TestLaneExhaustion:
    """AC3.6, AC8.2: Lane exhaustion and TokenRejected."""

    def test_alloc_shared_exhausts_all_lanes(self):
        """Allocate all lanes, then ALLOC_SHARED fails with TokenRejected."""
        env = simpy.Environment()
        events = []
        config = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Parent ALLOC uses lane 0
        fct_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct_parent)

        # Allocate lanes 1, 2, 3
        for i in range(1, 4):
            fct = FrameControlToken(
                target=0, act_id=i, op=FrameOp.ALLOC_SHARED, payload=0
            )
            inject_and_run(env, pe, fct)
            assert i in pe.tag_store, f"Child {i} should be allocated"

        # Try to allocate one more (all lanes exhausted)
        fct_fail = FrameControlToken(
            target=0, act_id=4, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_fail)

        rejected = [e for e in events if isinstance(e, TokenRejected)]
        assert len(rejected) > 0, "Should have TokenRejected event"
        assert "no free lanes" in rejected[0].reason, "Reason should be 'no free lanes'"

        # act_id=4 should not be in tag_store
        assert 4 not in pe.tag_store, "Failed allocation should not add to tag_store"

    def test_lane_exhaustion_with_multiple_frames(self):
        """Lane exhaustion is per-frame; different frames have independent lanes."""
        env = simpy.Environment()
        events = []
        config = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Frame 1: Parent 0 allocates lane 0
        fct1 = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct1)
        frame1_id, _lane = pe.tag_store[0]

        # Frame 2: Parent 10 allocates lane 0
        fct2 = FrameControlToken(
            target=0, act_id=10, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct2)
        frame2_id, _lane = pe.tag_store[10]

        assert frame1_id != frame2_id, "Should allocate different frames"

        # Frame 1: Exhaust all lanes
        for i in range(1, 4):
            fct = FrameControlToken(
                target=0, act_id=i, op=FrameOp.ALLOC_SHARED, payload=0
            )
            inject_and_run(env, pe, fct)

        # Frame 2: Can still allocate more lanes (independent)
        for i in range(11, 14):
            fct = FrameControlToken(
                target=0, act_id=i, op=FrameOp.ALLOC_SHARED, payload=10
            )
            inject_and_run(env, pe, fct)
            assert i in pe.tag_store, f"Frame2 child {i} should be allocated"


class TestFreeLane:
    """AC3.2: FREE_LANE clears lane data, keeps frame, returns lane to pool."""

    def test_free_lane_basic(self):
        """FREE_LANE removes act_id from tag_store, clears lane data, keeps frame."""
        env = simpy.Environment()
        events = []
        config = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Parent ALLOC
        fct_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct_parent)
        parent_frame_id, _parent_lane = pe.tag_store[0]

        # Child ALLOC_SHARED
        fct_child = FrameControlToken(
            target=0, act_id=1, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_child)
        _child_frame_id, child_lane = pe.tag_store[1]

        # FREE_LANE for child
        fct_free = FrameControlToken(
            target=0, act_id=1, op=FrameOp.FREE_LANE, payload=0
        )
        inject_and_run(env, pe, fct_free)

        # Child should be removed from tag_store
        assert 1 not in pe.tag_store, "Child should be removed from tag_store"

        # Parent should still be present
        assert 0 in pe.tag_store, "Parent should still be in tag_store"

        # Frame should NOT be in free_frames (still used by parent)
        assert parent_frame_id not in pe.free_frames, "Frame should not be free"

        # FrameFreed event should have frame_freed=False
        frame_freed = [e for e in events if isinstance(e, FrameFreed)]
        assert len(frame_freed) > 0, "Should have FrameFreed event"
        assert frame_freed[-1].frame_freed == False, "frame_freed should be False"
        assert frame_freed[-1].lane == child_lane, "Event should report correct lane"

    def test_free_lane_returns_lane_to_pool(self):
        """After FREE_LANE, freed lane can be reused by ALLOC_SHARED."""
        env = simpy.Environment()
        events = []
        config = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Parent ALLOC
        fct_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct_parent)
        parent_frame_id, _parent_lane = pe.tag_store[0]

        # Child 1 ALLOC_SHARED (lane 1)
        fct_child1 = FrameControlToken(
            target=0, act_id=1, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_child1)
        _child1_frame_id, child1_lane = pe.tag_store[1]
        assert child1_lane == 1

        # FREE_LANE child 1
        fct_free = FrameControlToken(
            target=0, act_id=1, op=FrameOp.FREE_LANE, payload=0
        )
        inject_and_run(env, pe, fct_free)

        # Child 2 ALLOC_SHARED (should get lane 1 again)
        fct_child2 = FrameControlToken(
            target=0, act_id=2, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_child2)
        _child2_frame_id, child2_lane = pe.tag_store[2]

        # Lane 1 should be reused for child 2
        assert child2_lane == 1, "Freed lane 1 should be reused"


class TestIndependentMatching:
    """AC8.1: Two act_ids sharing a frame have independent matching."""

    def test_independent_matching_same_offset(self):
        """L operand for act_id 0 does not interfere with L for act_id 1."""
        env = simpy.Environment()
        events = []
        config = PEConfig(
            frame_count=4, lane_count=4, matchable_offsets=4, on_event=events.append
        )
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Parent ALLOC
        fct_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct_parent)
        parent_frame_id, _parent_lane = pe.tag_store[0]

        # Child ALLOC_SHARED
        fct_child = FrameControlToken(
            target=0, act_id=1, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_child)
        _child_frame_id, child_lane = pe.tag_store[1]

        # Install dyadic instruction at offset 0
        inst = Instruction(
            opcode=ArithOp.ADD,
            output=OutputStyle.SINK,
            has_const=False,
            dest_count=0,
            wide=False,
            fref=0,
        )
        pe.iram[0] = inst

        # Send L operand for act_id=0
        tok_l_0 = DyadToken(
            target=0, offset=0, act_id=0, data=5, port=Port.L
        )
        inject_and_run(env, pe, tok_l_0)

        # Should have 1 TokenReceived, 0 Matched (waiting for R)
        matched = [e for e in events if isinstance(e, Matched)]
        assert len(matched) == 0, "Should not match yet (waiting for R)"

        # Send L operand for act_id=1 at same offset
        tok_l_1 = DyadToken(
            target=0, offset=0, act_id=1, data=7, port=Port.L
        )
        inject_and_run(env, pe, tok_l_1)

        # Should still have 0 Matched (both waiting for R)
        matched = [e for e in events if isinstance(e, Matched)]
        assert len(matched) == 0, "Both should be waiting for R"

        # Send R for act_id=0
        tok_r_0 = DyadToken(
            target=0, offset=0, act_id=0, data=3, port=Port.R
        )
        inject_and_run(env, pe, tok_r_0)

        # Should now have 1 Matched for act_id=0
        matched = [e for e in events if isinstance(e, Matched)]
        assert len(matched) == 1, "Should have 1 match for act_id=0"
        assert matched[0].act_id == 0, "Match should be for act_id=0"
        assert matched[0].left == 5, "Left should be 5"
        assert matched[0].right == 3, "Right should be 3"

        # Send R for act_id=1
        tok_r_1 = DyadToken(
            target=0, offset=0, act_id=1, data=2, port=Port.R
        )
        inject_and_run(env, pe, tok_r_1)

        # Should now have 2 Matched
        matched = [e for e in events if isinstance(e, Matched)]
        assert len(matched) == 2, "Should have 2 matches total"
        m1 = [m for m in matched if m.act_id == 1][0]
        assert m1.left == 7, "act_id=1 left should be 7"
        assert m1.right == 2, "act_id=1 right should be 2"

    def test_independent_matching_different_offsets(self):
        """Different offsets per lane maintain independence."""
        env = simpy.Environment()
        events = []
        config = PEConfig(
            frame_count=4, lane_count=4, matchable_offsets=4, on_event=events.append
        )
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Parent ALLOC
        fct_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct_parent)

        # Child ALLOC_SHARED
        fct_child = FrameControlToken(
            target=0, act_id=1, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_child)

        # Install dyadic instructions at offsets 0 and 1
        inst0 = Instruction(
            opcode=ArithOp.ADD, output=OutputStyle.SINK,
            has_const=False, dest_count=0, wide=False, fref=0
        )
        inst1 = Instruction(
            opcode=ArithOp.SUB, output=OutputStyle.SINK,
            has_const=False, dest_count=0, wide=False, fref=0
        )
        pe.iram[0] = inst0
        pe.iram[1] = inst1

        # Send L for act_id=0 at offset 0
        tok_l_0_off0 = DyadToken(
            target=0, offset=0, act_id=0, data=10, port=Port.L
        )
        inject_and_run(env, pe, tok_l_0_off0)

        # Send L for act_id=1 at offset 1
        tok_l_1_off1 = DyadToken(
            target=0, offset=1, act_id=1, data=20, port=Port.L
        )
        inject_and_run(env, pe, tok_l_1_off1)

        # Neither should match yet
        matched = [e for e in events if isinstance(e, Matched)]
        assert len(matched) == 0, "No matches yet"

        # Send R for act_id=0 at offset 0
        tok_r_0_off0 = DyadToken(
            target=0, offset=0, act_id=0, data=5, port=Port.R
        )
        inject_and_run(env, pe, tok_r_0_off0)

        # Should match for offset 0
        matched = [e for e in events if isinstance(e, Matched)]
        assert len(matched) == 1, "Should have 1 match"
        assert matched[0].offset == 0, "Match should be at offset 0"

        # Send R for act_id=1 at offset 1
        tok_r_1_off1 = DyadToken(
            target=0, offset=1, act_id=1, data=15, port=Port.R
        )
        inject_and_run(env, pe, tok_r_1_off1)

        # Should match for offset 1
        matched = [e for e in events if isinstance(e, Matched)]
        assert len(matched) == 2, "Should have 2 matches"
        m1 = [m for m in matched if m.offset == 1][0]
        assert m1.act_id == 1, "Offset 1 match should be act_id=1"


class TestSmartFree:
    """AC3.3, AC8.3: Smart FREE on shared frames preserves data and manages lanes."""

    def test_free_on_shared_frame_preserves_other_lanes(self):
        """FREE on act_id=0 when act_id=1 uses frame; lane 1 data preserved."""
        env = simpy.Environment()
        events = []
        config = PEConfig(
            frame_count=4, lane_count=4, matchable_offsets=4, on_event=events.append
        )
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Parent ALLOC
        fct_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct_parent)
        parent_frame_id, _parent_lane = pe.tag_store[0]

        # Child ALLOC_SHARED
        fct_child = FrameControlToken(
            target=0, act_id=1, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_child)
        _child_frame_id, child_lane = pe.tag_store[1]

        # Install instruction
        inst = Instruction(
            opcode=ArithOp.ADD, output=OutputStyle.SINK,
            has_const=False, dest_count=0, wide=False, fref=0
        )
        pe.iram[0] = inst

        # Store L operand on child's lane
        tok_l_1 = DyadToken(
            target=0, offset=0, act_id=1, data=7, port=Port.L
        )
        inject_and_run(env, pe, tok_l_1)

        # Verify child's match slot has data
        frame_id, lane = pe.tag_store[1]
        assert pe.match_data[frame_id][0][lane] == 7, "Child lane should have L operand"
        assert pe.presence[frame_id][0][lane] == True, "Child presence should be set"

        # FREE parent
        fct_free_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.FREE, payload=0
        )
        inject_and_run(env, pe, fct_free_parent)

        # Parent should be removed, child should still be present
        assert 0 not in pe.tag_store, "Parent should be removed"
        assert 1 in pe.tag_store, "Child should still be present"

        # Frame should NOT be in free_frames
        assert parent_frame_id not in pe.free_frames, "Frame should not be free"

        # Child's match data should be preserved
        assert pe.match_data[frame_id][0][lane] == 7, "Child data should be preserved"
        assert pe.presence[frame_id][0][lane] == True, "Child presence should be preserved"

        # FrameFreed event should have frame_freed=False
        frame_freed = [e for e in events if isinstance(e, FrameFreed)]
        assert any(e.frame_freed == False for e in frame_freed), "Should have frame_freed=False"

    def test_free_last_lane_returns_frame(self):
        """FREE on last act_id using frame returns frame to free_frames."""
        env = simpy.Environment()
        events = []
        config = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Parent ALLOC
        fct_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct_parent)
        parent_frame_id, _parent_lane = pe.tag_store[0]

        # Child ALLOC_SHARED
        fct_child = FrameControlToken(
            target=0, act_id=1, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_child)

        # FREE child
        fct_free_child = FrameControlToken(
            target=0, act_id=1, op=FrameOp.FREE_LANE, payload=0
        )
        inject_and_run(env, pe, fct_free_child)

        # Frame should still not be free (parent still using it)
        assert parent_frame_id not in pe.free_frames, "Frame should not be free yet"

        # FREE parent
        fct_free_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.FREE, payload=0
        )
        inject_and_run(env, pe, fct_free_parent)

        # Now frame should be free
        assert parent_frame_id in pe.free_frames, "Frame should be free"

        # tag_store should be empty
        assert len(pe.tag_store) == 0, "tag_store should be empty"

        # lane_free entry should be cleaned up
        assert parent_frame_id not in pe.lane_free, "lane_free entry should be cleaned"

        # FrameFreed event should have frame_freed=True
        frame_freed = [e for e in events if isinstance(e, FrameFreed)]
        assert any(e.frame_freed == True for e in frame_freed), "Should have frame_freed=True"

    def test_alloc_unchanged_allocates_fresh_frame(self):
        """Regular ALLOC still works: allocates fresh frame, lane 0."""
        env = simpy.Environment()
        events = []
        config = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # First ALLOC
        fct1 = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct1)
        frame_id_0, lane_0 = pe.tag_store[0]
        assert lane_0 == 0, "First ALLOC should assign lane 0"

        # Second ALLOC (different frame)
        fct2 = FrameControlToken(
            target=0, act_id=10, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct2)
        frame_id_10, lane_10 = pe.tag_store[10]
        assert lane_10 == 0, "Second ALLOC should assign lane 0"

        # Frames should be different
        assert frame_id_0 != frame_id_10, "Different ALLOC should get different frames"

    def test_data_preservation_across_free_lanes(self):
        """Match data on one lane not affected by FREE of another lane."""
        env = simpy.Environment()
        events = []
        config = PEConfig(
            frame_count=4, lane_count=4, matchable_offsets=4, on_event=events.append
        )
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Parent ALLOC
        fct_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct_parent)
        frame_id, _parent_lane = pe.tag_store[0]

        # Child 1 ALLOC_SHARED
        fct_child1 = FrameControlToken(
            target=0, act_id=1, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_child1)
        _frame_id_1, lane_1 = pe.tag_store[1]

        # Child 2 ALLOC_SHARED
        fct_child2 = FrameControlToken(
            target=0, act_id=2, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_child2)
        _frame_id_2, lane_2 = pe.tag_store[2]

        # Install instruction
        inst = Instruction(
            opcode=ArithOp.ADD, output=OutputStyle.SINK,
            has_const=False, dest_count=0, wide=False, fref=0
        )
        pe.iram[0] = inst

        # Store L operand on lane 1
        tok_l_1 = DyadToken(
            target=0, offset=0, act_id=1, data=7, port=Port.L
        )
        inject_and_run(env, pe, tok_l_1)

        # Store L operand on lane 2
        tok_l_2 = DyadToken(
            target=0, offset=0, act_id=2, data=11, port=Port.L
        )
        inject_and_run(env, pe, tok_l_2)

        # FREE lane 1
        fct_free_1 = FrameControlToken(
            target=0, act_id=1, op=FrameOp.FREE_LANE, payload=0
        )
        inject_and_run(env, pe, fct_free_1)

        # Lane 2's data should be untouched
        assert pe.match_data[frame_id][0][lane_2] == 11, "Lane 2 data should be preserved"
        assert pe.presence[frame_id][0][lane_2] == True, "Lane 2 presence should be preserved"

        # Lane 1 should be cleared
        assert pe.match_data[frame_id][0][lane_1] is None, "Lane 1 data should be cleared"
        assert pe.presence[frame_id][0][lane_1] == False, "Lane 1 presence should be cleared"


class TestAllocRemoteDataDriven:
    """AC8.4, AC8.5: ALLOC_REMOTE reads fref+2 for data-driven ALLOC_SHARED vs ALLOC."""

    def test_alloc_remote_emits_alloc_shared_when_parent_nonzero(self):
        """AC8.4: ALLOC_REMOTE emits ALLOC_SHARED when fref+2 is non-zero."""
        env = simpy.Environment()
        events = []
        output_store = simpy.Store(env)

        # PE0: source of ALLOC_REMOTE
        config0 = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe0 = ProcessingElement(env=env, pe_id=0, config=config0)
        pe0.route_table[1] = output_store  # Capture emitted token

        # Allocate a frame for act_id=0 on PE0
        fct_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe0, fct_parent)
        frame_id, _lane = pe0.tag_store[0]

        # Set up ALLOC_REMOTE instruction with fref pointing to frame constants
        # fref+0: target PE=1, fref+1: target act_id=5, fref+2: parent act_id=3
        inst = Instruction(
            opcode=RoutingOp.ALLOC_REMOTE,
            output=OutputStyle.SINK,  # Not used for ALLOC_REMOTE
            has_const=False,
            dest_count=0,
            wide=False,
            fref=10,
        )
        pe0.iram[0] = inst

        # Load frame slots with constants
        pe0.frames[frame_id][10] = 1       # target PE
        pe0.frames[frame_id][11] = 5       # target act_id
        pe0.frames[frame_id][12] = 3       # parent act_id (non-zero = ALLOC_SHARED)

        # Send MonadToken to trigger ALLOC_REMOTE
        tok = DyadToken(
            target=0, offset=0, act_id=0, data=0, port=Port.L
        )
        inject_and_run(env, pe0, tok)

        # Verify FrameControlToken was emitted with ALLOC_SHARED
        assert len(output_store.items) > 0, "Should have emitted a token"
        emitted = output_store.items[0]
        assert isinstance(emitted, FrameControlToken), "Should emit FrameControlToken"
        assert emitted.op == FrameOp.ALLOC_SHARED, "Should emit ALLOC_SHARED"
        assert emitted.payload == 3, "Payload should be parent act_id=3"
        assert emitted.target == 1, "Should target PE 1"
        assert emitted.act_id == 5, "Should target act_id 5"

    def test_alloc_remote_emits_alloc_when_parent_zero(self):
        """AC8.5: ALLOC_REMOTE emits ALLOC when fref+2 is zero (backwards compatible)."""
        env = simpy.Environment()
        events = []
        output_store = simpy.Store(env)

        # PE0: source of ALLOC_REMOTE
        config0 = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe0 = ProcessingElement(env=env, pe_id=0, config=config0)
        pe0.route_table[1] = output_store  # Capture emitted token

        # Allocate a frame for act_id=0 on PE0
        fct_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe0, fct_parent)
        frame_id, _lane = pe0.tag_store[0]

        # Set up ALLOC_REMOTE instruction
        # fref+0: target PE=1, fref+1: target act_id=5, fref+2: parent act_id=0
        inst = Instruction(
            opcode=RoutingOp.ALLOC_REMOTE,
            output=OutputStyle.SINK,
            has_const=False,
            dest_count=0,
            wide=False,
            fref=10,
        )
        pe0.iram[0] = inst

        # Load frame slots with constants
        pe0.frames[frame_id][10] = 1       # target PE
        pe0.frames[frame_id][11] = 5       # target act_id
        pe0.frames[frame_id][12] = 0       # parent act_id (zero = ALLOC)

        # Send MonadToken to trigger ALLOC_REMOTE
        tok = DyadToken(
            target=0, offset=0, act_id=0, data=0, port=Port.L
        )
        inject_and_run(env, pe0, tok)

        # Verify FrameControlToken was emitted with ALLOC (not ALLOC_SHARED)
        assert len(output_store.items) > 0, "Should have emitted a token"
        emitted = output_store.items[0]
        assert isinstance(emitted, FrameControlToken), "Should emit FrameControlToken"
        assert emitted.op == FrameOp.ALLOC, "Should emit ALLOC"
        assert emitted.payload == 0, "Payload should be 0 for ALLOC"
        assert emitted.target == 1, "Should target PE 1"
        assert emitted.act_id == 5, "Should target act_id 5"

    def test_alloc_remote_fref_plus_2_missing_defaults_to_zero(self):
        """ALLOC_REMOTE gracefully handles fref+2 outside frame bounds (defaults to 0)."""
        env = simpy.Environment()
        events = []
        output_store = simpy.Store(env)

        # PE0: source of ALLOC_REMOTE
        config0 = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe0 = ProcessingElement(env=env, pe_id=0, config=config0)
        pe0.route_table[1] = output_store

        # Allocate frame
        fct_parent = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe0, fct_parent)
        frame_id, _lane = pe0.tag_store[0]

        # Set up ALLOC_REMOTE with fref pointing near end of frame
        inst = Instruction(
            opcode=RoutingOp.ALLOC_REMOTE,
            output=OutputStyle.SINK,
            has_const=False,
            dest_count=0,
            wide=False,
            fref=62,  # frame_slots defaults to 64, so fref+2=64 is outside
        )
        pe0.iram[0] = inst

        # Load only fref+0 and fref+1 (fref+2 is beyond frame bounds)
        pe0.frames[frame_id][62] = 1
        pe0.frames[frame_id][63] = 7

        # Send MonadToken
        tok = DyadToken(
            target=0, offset=0, act_id=0, data=0, port=Port.L
        )
        inject_and_run(env, pe0, tok)

        # Should emit ALLOC (not ALLOC_SHARED) because fref+2 is missing/falsy
        assert len(output_store.items) > 0, "Should have emitted a token"
        emitted = output_store.items[0]
        assert emitted.op == FrameOp.ALLOC, "Should emit ALLOC when fref+2 is missing"


class TestFreeFrameOpcode:
    """AC5.1: FREE_FRAME opcode uses smart FREE behaviour on shared frames."""

    def test_free_frame_opcode_shared_frame_partial_free(self):
        """FREE_FRAME smart free: partial frame free when other lanes remain."""
        env = simpy.Environment()
        events = []
        config = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Pre-allocate frame with two act_ids on different lanes
        # This simulates ALLOC for act_id=0 and ALLOC_SHARED for act_id=1
        frame_id = 0
        pe.frames[frame_id] = [None] * pe.frame_slots
        pe.tag_store[0] = (frame_id, 0)  # act_id=0 on lane 0
        pe.tag_store[1] = (frame_id, 1)  # act_id=1 on lane 1
        pe.lane_free[frame_id] = {2, 3}  # Lanes 2 and 3 are free
        # Remove frame_id from free_frames (it's in use)
        if frame_id in pe.free_frames:
            pe.free_frames.remove(frame_id)

        # Install FREE_FRAME instruction
        inst = Instruction(
            opcode=RoutingOp.FREE_FRAME,
            output=OutputStyle.SINK,
            has_const=False,
            dest_count=0,
            wide=False,
            fref=0,
        )
        pe.iram[0] = inst

        # Send MonadToken for act_id=0 to trigger FREE_FRAME
        tok = DyadToken(
            target=0, offset=0, act_id=0, data=0, port=Port.L
        )
        inject_and_run(env, pe, tok)

        # Verify act_id=0 is removed from tag_store
        assert 0 not in pe.tag_store, "act_id=0 should be removed from tag_store"

        # Verify act_id=1 is still in tag_store
        assert 1 in pe.tag_store, "act_id=1 should still be in tag_store"

        # Verify frame is NOT returned to free_frames (still in use by act_id=1)
        assert frame_id not in pe.free_frames, "Frame should not be in free_frames"

        # Verify FrameFreed event has frame_freed=False
        frame_freed = [e for e in events if isinstance(e, FrameFreed)]
        assert any(e.frame_freed == False for e in frame_freed), \
            "Should have FrameFreed event with frame_freed=False"
        last_frame_freed = [e for e in frame_freed if e.act_id == 0][-1]
        assert last_frame_freed.frame_freed == False, "Frame should not be marked as freed"

    def test_free_frame_opcode_shared_frame_full_free(self):
        """FREE_FRAME smart free: full frame free when last lane is freed."""
        env = simpy.Environment()
        events = []
        config = PEConfig(frame_count=4, lane_count=4, on_event=events.append)
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # Pre-allocate frame with two act_ids
        frame_id = 0
        pe.frames[frame_id] = [None] * pe.frame_slots
        pe.tag_store[0] = (frame_id, 0)  # act_id=0 on lane 0
        pe.tag_store[1] = (frame_id, 1)  # act_id=1 on lane 1
        pe.lane_free[frame_id] = {2, 3}
        if frame_id in pe.free_frames:
            pe.free_frames.remove(frame_id)

        # Install FREE_FRAME instruction
        inst = Instruction(
            opcode=RoutingOp.FREE_FRAME,
            output=OutputStyle.SINK,
            has_const=False,
            dest_count=0,
            wide=False,
            fref=0,
        )
        pe.iram[0] = inst

        # First: free act_id=0
        tok0 = DyadToken(
            target=0, offset=0, act_id=0, data=0, port=Port.L
        )
        inject_and_run(env, pe, tok0)

        # Verify frame still not free
        assert frame_id not in pe.free_frames, "Frame should not be free after first FREE_FRAME"
        assert 1 in pe.tag_store, "act_id=1 should still be present"

        # Second: free act_id=1 (last lane on frame)
        tok1 = DyadToken(
            target=0, offset=0, act_id=1, data=0, port=Port.L
        )
        inject_and_run(env, pe, tok1)

        # Verify frame is now freed
        assert frame_id in pe.free_frames, "Frame should be in free_frames after last FREE_FRAME"
        assert 1 not in pe.tag_store, "act_id=1 should be removed from tag_store"

        # Verify tag_store is empty
        assert len(pe.tag_store) == 0, "tag_store should be empty"

        # Verify lane_free is cleaned up
        assert frame_id not in pe.lane_free, "lane_free entry should be deleted"

        # Verify FrameFreed event has frame_freed=True
        frame_freed = [e for e in events if isinstance(e, FrameFreed)]
        last_frame_freed = [e for e in frame_freed if e.act_id == 1][-1]
        assert last_frame_freed.frame_freed == True, \
            "Last FREE_FRAME should emit FrameFreed with frame_freed=True"


class TestLoopPipelining:
    """AC8.6: Full loop pipelining integration test with multiple lanes."""

    def test_full_loop_pipelining_scenario(self):
        """
        Complete loop pipelining lifecycle: two iterations of a dyadic instruction
        running concurrently on different lanes, both producing correct results.

        Simulates:
        1. ALLOC(act_id=0)                          → frame, lane 0
        2. Setup: write destination to frame
        3. Iteration 1: inject L and R DyadTokens for act_id=0
        4. ALLOC_SHARED(act_id=1, parent=0)         → same frame, lane 1
        5. Iteration 2: inject L and R DyadTokens for act_id=1
        6. Both iterations match independently, both produce correct results
        7. FREE(act_id=0) → lane 0 freed, frame stays
        8. FREE(act_id=1) → last lane, frame returned to free list
        """
        env = simpy.Environment()
        events = []
        config = PEConfig(
            frame_count=4, lane_count=4, matchable_offsets=4,
            on_event=events.append
        )
        pe = ProcessingElement(env=env, pe_id=0, config=config)

        # 1. ALLOC(act_id=0) → frame, lane 0
        fct_alloc_0 = FrameControlToken(
            target=0, act_id=0, op=FrameOp.ALLOC, payload=0
        )
        inject_and_run(env, pe, fct_alloc_0)

        # Verify act_id=0 is allocated
        assert 0 in pe.tag_store, "act_id=0 should be in tag_store"
        frame_id, lane_0 = pe.tag_store[0]
        assert lane_0 == 0, "First ALLOC should assign lane 0"

        # Verify FrameAllocated event for iteration 1
        frame_allocated = [e for e in events if isinstance(e, FrameAllocated)]
        assert len(frame_allocated) >= 1, "Should have FrameAllocated event"
        assert frame_allocated[0].frame_id == frame_id, "Event should report correct frame_id"
        assert frame_allocated[0].lane == 0, "Event should report lane 0"

        # 2. Setup: write destination to frame at slot 8
        dest = FrameDest(
            target_pe=1, offset=0, act_id=0, port=Port.L,
            token_kind=TokenKind.MONADIC
        )
        pe.frames[frame_id][8] = dest

        # Set up route to capture output
        pe.route_table[1] = simpy.Store(env)

        # 3. Install ADD instruction at IRAM offset 0
        inst = Instruction(
            opcode=ArithOp.ADD,
            output=OutputStyle.INHERIT,
            has_const=False,
            dest_count=1,
            wide=False,
            fref=8,
        )
        pe.iram[0] = inst

        # 4. ALLOC_SHARED(act_id=1, parent=0) → same frame, lane 1
        fct_alloc_shared = FrameControlToken(
            target=0, act_id=1, op=FrameOp.ALLOC_SHARED, payload=0
        )
        inject_and_run(env, pe, fct_alloc_shared)

        # Verify act_id=1 is allocated on same frame, different lane
        assert 1 in pe.tag_store, "act_id=1 should be in tag_store"
        frame_id_1, lane_1 = pe.tag_store[1]
        assert frame_id_1 == frame_id, "Both should share same frame"
        assert lane_1 == 1, "Second allocation should assign lane 1"
        assert lane_1 != lane_0, "Lanes should be different"

        # Verify FrameAllocated event for iteration 2
        frame_allocated = [e for e in events if isinstance(e, FrameAllocated)]
        assert len(frame_allocated) >= 2, "Should have 2 FrameAllocated events"
        assert frame_allocated[1].frame_id == frame_id, "Event should report correct frame_id"
        assert frame_allocated[1].lane == 1, "Event should report lane 1"

        # 5. Inject iteration 1 operands (act_id=0, lane 0)
        tok_l_0 = DyadToken(
            target=0, offset=0, act_id=0, data=100, port=Port.L
        )
        inject_and_run(env, pe, tok_l_0)

        tok_r_0 = DyadToken(
            target=0, offset=0, act_id=0, data=200, port=Port.R
        )
        inject_and_run(env, pe, tok_r_0)

        # Verify Matched event for iteration 1
        matched = [e for e in events if isinstance(e, Matched)]
        assert len(matched) >= 1, "Should have Matched event for iteration 1"
        match_0 = [m for m in matched if m.act_id == 0][-1]
        assert match_0.left == 100, "Iteration 1 left operand should be 100"
        assert match_0.right == 200, "Iteration 1 right operand should be 200"
        assert match_0.offset == 0, "Iteration 1 offset should be 0"

        # Verify output token with correct data (100+200=300)
        emitted = [e for e in events if isinstance(e, Emitted)]
        assert len(emitted) >= 1, "Should have Emitted event for iteration 1"
        out_tok_0 = emitted[-1].token
        assert out_tok_0.data == 300, "Iteration 1 output should be 300 (100+200)"
        assert out_tok_0.target == 1, "Output should route to target_pe=1"

        # 6. Inject iteration 2 operands (act_id=1, lane 1)
        tok_l_1 = DyadToken(
            target=0, offset=0, act_id=1, data=1000, port=Port.L
        )
        inject_and_run(env, pe, tok_l_1)

        tok_r_1 = DyadToken(
            target=0, offset=0, act_id=1, data=2000, port=Port.R
        )
        inject_and_run(env, pe, tok_r_1)

        # Verify Matched event for iteration 2
        matched = [e for e in events if isinstance(e, Matched)]
        assert len(matched) >= 2, "Should have Matched events for both iterations"
        match_1 = [m for m in matched if m.act_id == 1][-1]
        assert match_1.left == 1000, "Iteration 2 left operand should be 1000"
        assert match_1.right == 2000, "Iteration 2 right operand should be 2000"
        assert match_1.offset == 0, "Iteration 2 offset should be 0"

        # Verify output token with correct data (1000+2000=3000)
        emitted = [e for e in events if isinstance(e, Emitted)]
        assert len(emitted) >= 2, "Should have Emitted events for both iterations"
        out_tok_1 = emitted[-1].token
        assert out_tok_1.data == 3000, "Iteration 2 output should be 3000 (1000+2000)"
        assert out_tok_1.target == 1, "Output should route to target_pe=1"

        # Interleaved verification: confirm independent lanes
        matches_by_id = {}
        for m in matched:
            if m.act_id not in matches_by_id:
                matches_by_id[m.act_id] = []
            matches_by_id[m.act_id].append(m)

        assert 0 in matches_by_id, "Should have match for iteration 1 (act_id=0)"
        assert 1 in matches_by_id, "Should have match for iteration 2 (act_id=1)"
        assert matches_by_id[0][-1].left == 100, "Iteration 1 left should be 100"
        assert matches_by_id[1][-1].left == 1000, "Iteration 2 left should be 1000"

        # 7. FREE(act_id=0) → lane 0 freed, frame stays
        fct_free_0 = FrameControlToken(
            target=0, act_id=0, op=FrameOp.FREE, payload=0
        )
        inject_and_run(env, pe, fct_free_0)

        # Verify act_id=0 removed, act_id=1 still present
        assert 0 not in pe.tag_store, "act_id=0 should be removed from tag_store"
        assert 1 in pe.tag_store, "act_id=1 should still be in tag_store"

        # Verify frame not returned (still used by act_id=1)
        assert frame_id not in pe.free_frames, "Frame should not be in free_frames"

        # Verify FrameFreed event with frame_freed=False
        frame_freed = [e for e in events if isinstance(e, FrameFreed)]
        freed_0 = [f for f in frame_freed if f.act_id == 0][-1]
        assert freed_0.frame_freed == False, "frame_freed should be False (not last lane)"
        assert freed_0.lane == lane_0, "Event should report lane 0"

        # 8. FREE(act_id=1) → last lane, frame returned to free list
        fct_free_1 = FrameControlToken(
            target=0, act_id=1, op=FrameOp.FREE, payload=0
        )
        inject_and_run(env, pe, fct_free_1)

        # Verify act_id=1 removed from tag_store
        assert 1 not in pe.tag_store, "act_id=1 should be removed from tag_store"

        # Verify tag_store is now empty
        assert len(pe.tag_store) == 0, "tag_store should be empty"

        # Verify frame returned to free_frames
        assert frame_id in pe.free_frames, "Frame should be in free_frames"

        # Verify lane_free entry cleaned up
        assert frame_id not in pe.lane_free, "lane_free entry should be deleted"

        # Verify FrameFreed event with frame_freed=True
        frame_freed = [e for e in events if isinstance(e, FrameFreed)]
        freed_1 = [f for f in frame_freed if f.act_id == 1][-1]
        assert freed_1.frame_freed == True, "frame_freed should be True (last lane)"
        assert freed_1.lane == lane_1, "Event should report lane 1"

        # Summary: verify AC8.6 acceptance criteria
        # Both iterations produce mathematically correct results
        assert matches_by_id[0][-1].left + matches_by_id[0][-1].right == 300, \
            "Iteration 1 arithmetic correct"
        assert matches_by_id[1][-1].left + matches_by_id[1][-1].right == 3000, \
            "Iteration 2 arithmetic correct"

        # Both iterations ran on SAME frame (verified at allocation, re-confirmed)
        assert frame_id_1 == frame_id, "Both iterations ran on same frame"

        # Both iterations used DIFFERENT lanes
        assert lane_0 != lane_1, "Iterations used different lanes"
        assert lane_0 == 0 and lane_1 == 1, "Lanes are 0 and 1 respectively"

        # Freeing one iteration preserved the other
        frame_freed_events = [e for e in events if isinstance(e, FrameFreed)]
        assert len(frame_freed_events) >= 2, "Should have 2 FrameFreed events"

        # Freeing the last iteration returned the frame
        assert frame_id in pe.free_frames, "Frame returned to pool after last FREE"