"""Tests for CPU-specific optimization passes."""

from __future__ import annotations

from sagellm_backend.graph.graph import Graph, Node, OpType
from sagellm_backend.graph.passes.constant_folding import ConstantFoldingPass
from sagellm_backend.graph.passes.dead_node_elimination import DeadNodeEliminationPass
from sagellm_backend.graph.passes.identity_elimination import IdentityEliminationPass
from sagellm_backend.graph.passes.operator_fusion import OperatorFusionPass
from sagellm_backend.graph.passes.static_precompute import StaticPatternPrecomputePass


class TestDeadNodeEliminationPass:
    """Tests for DeadNodeEliminationPass."""

    def test_pass_name(self) -> None:
        """Test pass name."""
        pass_obj = DeadNodeEliminationPass()
        assert pass_obj.name == "dead_node_elimination"

    def test_no_dead_nodes(self) -> None:
        """Test graph with no dead nodes."""
        pass_obj = DeadNodeEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        n2 = Node(node_id="n2", op_type=OpType.OUTPUT)
        graph.add_node(n1)
        graph.add_node(n2)
        graph.add_edge("n1", "n2")
        graph.outputs = ["n2"]

        result = pass_obj.run(graph)

        assert result.modified is False
        assert len(graph.nodes) == 2

    def test_remove_single_dead_node(self) -> None:
        """Test removing a single dead node."""
        pass_obj = DeadNodeEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        n2 = Node(node_id="n2", op_type=OpType.OUTPUT)
        n3 = Node(node_id="n3", op_type=OpType.ADD)  # Dead - not connected to output

        graph.add_node(n1)
        graph.add_node(n2)
        graph.add_node(n3)
        graph.add_edge("n1", "n2")
        graph.outputs = ["n2"]

        result = pass_obj.run(graph)

        assert result.modified is True
        assert result.num_changes == 1
        assert len(graph.nodes) == 2
        assert "n3" not in graph.nodes

    def test_remove_dead_subgraph(self) -> None:
        """Test removing a dead subgraph."""
        pass_obj = DeadNodeEliminationPass()
        graph = Graph()

        # Live path: n1 -> n2 -> n3
        # Dead path: n4 -> n5 -> n6
        live_nodes = [
            Node(node_id="n1", op_type=OpType.INPUT),
            Node(node_id="n2", op_type=OpType.ADD),
            Node(node_id="n3", op_type=OpType.OUTPUT),
        ]
        dead_nodes = [
            Node(node_id="n4", op_type=OpType.INPUT),
            Node(node_id="n5", op_type=OpType.MUL),
            Node(node_id="n6", op_type=OpType.OUTPUT),
        ]

        for node in live_nodes + dead_nodes:
            graph.add_node(node)

        graph.add_edge("n1", "n2")
        graph.add_edge("n2", "n3")
        graph.add_edge("n4", "n5")
        graph.add_edge("n5", "n6")

        # Only n3 is marked as output
        graph.outputs = ["n3"]

        result = pass_obj.run(graph)

        assert result.modified is True
        assert result.num_changes == 3  # n4, n5, n6 removed
        assert len(graph.nodes) == 3
        assert "n4" not in graph.nodes
        assert "n5" not in graph.nodes
        assert "n6" not in graph.nodes

    def test_keep_control_flow_referenced_nodes(self) -> None:
        """Test control-flow referenced branch nodes are retained."""
        pass_obj = DeadNodeEliminationPass()
        graph = Graph()

        cond = Node(node_id="cond", op_type=OpType.INPUT)
        branch_then = Node(node_id="then_branch", op_type=OpType.ADD)
        branch_else = Node(node_id="else_branch", op_type=OpType.MUL)
        merge = Node(node_id="merge", op_type=OpType.ADD)
        ctl = Node(
            node_id="if1",
            op_type=OpType.CONDITIONAL,
            attrs={
                "condition": "cond",
                "then_branch": "then_branch",
                "else_branch": "else_branch",
                "merge": "merge",
            },
        )
        out = Node(node_id="out", op_type=OpType.OUTPUT)
        dead = Node(node_id="dead", op_type=OpType.ADD)

        for node in [cond, branch_then, branch_else, merge, ctl, out, dead]:
            graph.add_node(node)

        graph.add_edge("cond", "if1")
        graph.add_edge("if1", "out")
        graph.outputs = ["out"]

        result = pass_obj.run(graph)

        assert result.modified is True
        assert "dead" not in graph.nodes
        assert "then_branch" in graph.nodes
        assert "else_branch" in graph.nodes
        assert "merge" in graph.nodes

    def test_keep_all_paths_to_outputs(self) -> None:
        """Test that all paths to outputs are kept."""
        pass_obj = DeadNodeEliminationPass()
        graph = Graph()

        # Diamond pattern: n1 splits to n2 and n3, both feed to n4
        #     n1
        #    /  \\
        #   n2  n3
        #    \\  /
        #     n4
        nodes = [
            Node(node_id="n1", op_type=OpType.INPUT),
            Node(node_id="n2", op_type=OpType.ADD),
            Node(node_id="n3", op_type=OpType.MUL),
            Node(node_id="n4", op_type=OpType.OUTPUT),
        ]

        for node in nodes:
            graph.add_node(node)

        graph.add_edge("n1", "n2")
        graph.add_edge("n1", "n3")
        graph.add_edge("n2", "n4")
        graph.add_edge("n3", "n4")
        graph.outputs = ["n4"]

        result = pass_obj.run(graph)

        # No nodes should be removed
        assert result.modified is False
        assert len(graph.nodes) == 4


class TestConstantFoldingPass:
    """Tests for Co τnstantFoldingPass."""

    def test_pass_name(self) -> None:
        """Test pass name."""
        pass_obj = ConstantFoldingPass()
        assert pass_obj.name == "constant_folding"

    def test_no_foldable_operations(self) -> None:
        """Test graph with no foldable operations."""
        pass_obj = ConstantFoldingPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        n2 = Node(node_id="n2", op_type=OpType.ADD)
        graph.add_node(n1)
        graph.add_node(n2)
        graph.add_edge("n1", "n2")

        result = pass_obj.run(graph)

        assert result.modified is False

    def test_fold_constant_add(self) -> None:
        """Test folding constant addition."""
        pass_obj = ConstantFoldingPass()
        graph = Graph()

        # Two constants feeding into an ADD
        c1 = Node(node_id="c1", op_type=OpType.CONSTANT, attrs={"value": 2})
        c2 = Node(node_id="c2", op_type=OpType.CONSTANT, attrs={"value": 3})
        add = Node(node_id="add1", op_type=OpType.ADD)

        graph.add_node(c1)
        graph.add_node(c2)
        graph.add_node(add)
        graph.add_edge("c1", "add1")
        graph.add_edge("c2", "add1")

        result = pass_obj.run(graph)

        assert result.modified is True
        assert result.num_changes == 1

        # ADD should be converted to CONSTANT
        assert graph.nodes["add1"].op_type == OpType.CONSTANT
        assert graph.nodes["add1"].metadata["folded"] is True
        assert graph.nodes["add1"].attrs["value"] == 5.0
        assert graph.nodes["add1"].metadata["original_op"] == "add"

    def test_fold_multiple_operations(self) -> None:
        """Test folding multiple constant operations."""
        pass_obj = ConstantFoldingPass()
        graph = Graph()

        # Multiple constant operations
        c1 = Node(node_id="c1", op_type=OpType.CONSTANT, attrs={"value": 2})
        c2 = Node(node_id="c2", op_type=OpType.CONSTANT, attrs={"value": 3})
        c3 = Node(node_id="c3", op_type=OpType.CONSTANT, attrs={"value": 4})
        add1 = Node(node_id="add1", op_type=OpType.ADD)
        mul1 = Node(node_id="mul1", op_type=OpType.MUL)

        graph.add_node(c1)
        graph.add_node(c2)
        graph.add_node(c3)
        graph.add_node(add1)
        graph.add_node(mul1)

        graph.add_edge("c1", "add1")
        graph.add_edge("c2", "add1")
        graph.add_edge("c3", "mul1")
        graph.add_edge("add1", "mul1")

        # Run once - should fold add1 and mul1 transitively
        result1 = pass_obj.run(graph)
        assert result1.modified is True
        assert graph.nodes["add1"].attrs["value"] == 5.0
        assert graph.nodes["mul1"].attrs["value"] == 20.0

    def test_fold_reshape_with_static_shape(self) -> None:
        """Test folding reshape on static constant tensor."""
        pass_obj = ConstantFoldingPass()
        graph = Graph()

        c1 = Node(node_id="c1", op_type=OpType.CONSTANT, attrs={"value": [1, 2, 3, 4]})
        reshape = Node(node_id="reshape1", op_type=OpType.RESHAPE, attrs={"shape": [2, 2]})

        graph.add_node(c1)
        graph.add_node(reshape)
        graph.add_edge("c1", "reshape1")

        result = pass_obj.run(graph)

        assert result.modified is True
        assert graph.nodes["reshape1"].op_type == OpType.CONSTANT
        assert graph.nodes["reshape1"].attrs["value"] == [[1, 2], [3, 4]]

    def test_skip_already_constant(self) -> None:
        """Test that constant nodes are not folded again."""
        pass_obj = ConstantFoldingPass()
        graph = Graph()

        c1 = Node(node_id="c1", op_type=OpType.CONSTANT)
        graph.add_node(c1)

        result = pass_obj.run(graph)

        assert result.modified is False


class TestStaticPatternPrecomputePass:
    """Tests for StaticPatternPrecomputePass."""

    def test_pass_name(self) -> None:
        """Test pass name."""
        pass_obj = StaticPatternPrecomputePass()
        assert pass_obj.name == "static_precompute"

    def test_precompute_rope(self) -> None:
        """Test RoPE precompute conversion to constant."""
        pass_obj = StaticPatternPrecomputePass()
        graph = Graph()

        rope = Node(
            node_id="rope1",
            op_type=OpType.PLACEHOLDER,
            attrs={"precompute_pattern": "rope", "seq_len": 4, "dim": 8},
        )
        graph.add_node(rope)

        result = pass_obj.run(graph)

        assert result.modified is True
        assert graph.nodes["rope1"].op_type == OpType.CONSTANT
        value = graph.nodes["rope1"].attrs["value"]
        assert value["type"] == "rope"
        assert len(value["cos"]) == 4
        assert len(value["cos"][0]) == 4

    def test_precompute_alibi(self) -> None:
        """Test ALiBi precompute conversion to constant."""
        pass_obj = StaticPatternPrecomputePass()
        graph = Graph()

        alibi = Node(
            node_id="alibi1",
            op_type=OpType.PLACEHOLDER,
            attrs={"precompute_pattern": "alibi", "num_heads": 2, "seq_len": 3},
        )
        graph.add_node(alibi)

        result = pass_obj.run(graph)

        assert result.modified is True
        value = graph.nodes["alibi1"].attrs["value"]
        assert value["type"] == "alibi"
        assert len(value["slopes"]) == 2
        assert len(value["bias"]) == 2
        assert len(value["bias"][0]) == 3

    def test_precompute_static_mask(self) -> None:
        """Test static mask precompute conversion to constant."""
        pass_obj = StaticPatternPrecomputePass()
        graph = Graph()

        mask = Node(
            node_id="mask1",
            op_type=OpType.PLACEHOLDER,
            attrs={"precompute_pattern": "static_mask", "seq_len": 3, "causal": True},
        )
        graph.add_node(mask)

        result = pass_obj.run(graph)

        assert result.modified is True
        value = graph.nodes["mask1"].attrs["value"]
        assert value["type"] == "static_mask"
        assert value["mask"][0][1] == float("-inf")
        assert value["mask"][1][0] == 0.0


class TestIdentityEliminationPass:
    """Tests for IdentityEliminationPass."""

    def test_pass_name(self) -> None:
        """Test pass name."""
        pass_obj = IdentityEliminationPass()
        assert pass_obj.name == "identity_elimination"

    def test_no_identity_operations(self) -> None:
        """Test graph with no identity operations."""
        pass_obj = IdentityEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        n2 = Node(node_id="n2", op_type=OpType.ADD)
        graph.add_node(n1)
        graph.add_node(n2)
        graph.add_edge("n1", "n2")

        result = pass_obj.run(graph)

        assert result.modified is False

    def test_eliminate_identity_reshape(self) -> None:
        """Test eliminating identity reshape."""
        pass_obj = IdentityEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        reshape = Node(
            node_id="reshape1",
            op_type=OpType.RESHAPE,
            attrs={"is_identity": True},
        )
        n3 = Node(node_id="n3", op_type=OpType.OUTPUT)

        graph.add_node(n1)
        graph.add_node(reshape)
        graph.add_node(n3)
        graph.add_edge("n1", "reshape1")
        graph.add_edge("reshape1", "n3")

        result = pass_obj.run(graph)

        assert result.modified is True
        assert result.num_changes == 1
        assert "reshape1" not in graph.nodes

        # n1 should connect directly to n3
        assert "n3" in graph.nodes["n1"].outputs
        assert "n1" in graph.nodes["n3"].inputs

    def test_eliminate_identity_transpose(self) -> None:
        """Test eliminating identity transpose."""
        pass_obj = IdentityEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        transpose = Node(
            node_id="transpose1",
            op_type=OpType.TRANSPOSE,
            attrs={"perm": [0, 1, 2]},  # Identity permutation
        )
        n3 = Node(node_id="n3", op_type=OpType.OUTPUT)

        graph.add_node(n1)
        graph.add_node(transpose)
        graph.add_node(n3)
        graph.add_edge("n1", "transpose1")
        graph.add_edge("transpose1", "n3")

        result = pass_obj.run(graph)

        assert result.modified is True
        assert "transpose1" not in graph.nodes

    def test_eliminate_add_zero(self) -> None:
        """Test eliminating ADD with zero."""
        pass_obj = IdentityEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        add = Node(
            node_id="add1",
            op_type=OpType.ADD,
            attrs={"add_zero": True},
        )
        n3 = Node(node_id="n3", op_type=OpType.OUTPUT)

        graph.add_node(n1)
        graph.add_node(add)
        graph.add_node(n3)
        graph.add_edge("n1", "add1")
        graph.add_edge("add1", "n3")

        result = pass_obj.run(graph)

        assert result.modified is True
        assert "add1" not in graph.nodes

    def test_eliminate_mul_one(self) -> None:
        """Test eliminating MUL with one."""
        pass_obj = IdentityEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        mul = Node(
            node_id="mul1",
            op_type=OpType.MUL,
            attrs={"mul_one": True},
        )
        n3 = Node(node_id="n3", op_type=OpType.OUTPUT)

        graph.add_node(n1)
        graph.add_node(mul)
        graph.add_node(n3)
        graph.add_edge("n1", "mul1")
        graph.add_edge("mul1", "n3")

        result = pass_obj.run(graph)

        assert result.modified is True
        assert "mul1" not in graph.nodes

    def test_eliminate_chain_of_identities(self) -> None:
        """Test eliminating multiple identity operations in sequence."""
        pass_obj = IdentityEliminationPass()
        graph = Graph()

        n1 = Node(node_id="n1", op_type=OpType.INPUT)
        id1 = Node(node_id="id1", op_type=OpType.RESHAPE, attrs={"is_identity": True})
        id2 = Node(node_id="id2", op_type=OpType.ADD, attrs={"add_zero": True})
        n4 = Node(node_id="n4", op_type=OpType.OUTPUT)

        graph.add_node(n1)
        graph.add_node(id1)
        graph.add_node(id2)
        graph.add_node(n4)

        graph.add_edge("n1", "id1")
        graph.add_edge("id1", "id2")
        graph.add_edge("id2", "n4")

        # Run multiple times to eliminate chain
        result1 = pass_obj.run(graph)
        assert result1.modified is True
        assert result1.num_changes >= 1

        _ = pass_obj.run(graph)
        # May eliminate more or be done
        assert len(graph.nodes) <= 4


class TestOperatorFusionPass:
    """Tests for OperatorFusionPass."""

    def test_pass_name(self) -> None:
        """Test pass name."""
        pass_obj = OperatorFusionPass()
        assert pass_obj.name == "operator_fusion"

    def test_fuse_matmul_add(self) -> None:
        """Test MATMUL + ADD fusion."""
        pass_obj = OperatorFusionPass()
        graph = Graph()

        x = Node(node_id="x", op_type=OpType.INPUT)
        w = Node(node_id="w", op_type=OpType.INPUT)
        b = Node(node_id="b", op_type=OpType.CONSTANT)
        matmul = Node(node_id="matmul1", op_type=OpType.MATMUL)
        add = Node(node_id="add1", op_type=OpType.ADD)
        out = Node(node_id="out", op_type=OpType.OUTPUT)

        for node in [x, w, b, matmul, add, out]:
            graph.add_node(node)

        graph.add_edge("x", "matmul1")
        graph.add_edge("w", "matmul1")
        graph.add_edge("matmul1", "add1")
        graph.add_edge("b", "add1")
        graph.add_edge("add1", "out")
        graph.outputs = ["out"]

        result = pass_obj.run(graph)

        assert result.modified is True
        assert result.num_changes == 1
        assert "matmul1" not in graph.nodes
        assert graph.nodes["add1"].op_type == OpType.FUSION
        assert graph.nodes["add1"].attrs["fused_pattern"] == "matmul_add"
        assert graph.nodes["add1"].metadata["fused"] is True
        assert "x" in graph.nodes["add1"].inputs
        assert "w" in graph.nodes["add1"].inputs
        assert "b" in graph.nodes["add1"].inputs

    def test_fuse_layernorm_linear(self) -> None:
        """Test LAYER_NORM + MATMUL(Linear) fusion."""
        pass_obj = OperatorFusionPass()
        graph = Graph()

        x = Node(node_id="x", op_type=OpType.INPUT)
        w = Node(node_id="w", op_type=OpType.INPUT)
        ln = Node(node_id="ln1", op_type=OpType.LAYER_NORM)
        linear = Node(node_id="linear1", op_type=OpType.MATMUL)
        out = Node(node_id="out", op_type=OpType.OUTPUT)

        for node in [x, w, ln, linear, out]:
            graph.add_node(node)

        graph.add_edge("x", "ln1")
        graph.add_edge("ln1", "linear1")
        graph.add_edge("w", "linear1")
        graph.add_edge("linear1", "out")
        graph.outputs = ["out"]

        result = pass_obj.run(graph)

        assert result.modified is True
        assert "ln1" not in graph.nodes
        assert graph.nodes["linear1"].op_type == OpType.FUSION
        assert graph.nodes["linear1"].attrs["fused_pattern"] == "layernorm_linear"
        assert "x" in graph.nodes["linear1"].inputs
        assert "w" in graph.nodes["linear1"].inputs

    def test_fuse_activation_elementwise(self) -> None:
        """Test activation + elementwise fusion."""
        pass_obj = OperatorFusionPass()
        graph = Graph()

        x = Node(node_id="x", op_type=OpType.INPUT)
        y = Node(node_id="y", op_type=OpType.INPUT)
        act = Node(node_id="act1", op_type=OpType.SILU)
        elem = Node(node_id="mul1", op_type=OpType.MUL)
        out = Node(node_id="out", op_type=OpType.OUTPUT)

        for node in [x, y, act, elem, out]:
            graph.add_node(node)

        graph.add_edge("x", "act1")
        graph.add_edge("act1", "mul1")
        graph.add_edge("y", "mul1")
        graph.add_edge("mul1", "out")
        graph.outputs = ["out"]

        result = pass_obj.run(graph)

        assert result.modified is True
        assert "act1" not in graph.nodes
        assert graph.nodes["mul1"].op_type == OpType.FUSION
        assert graph.nodes["mul1"].attrs["fused_pattern"] == "activation_elementwise"
        assert "x" in graph.nodes["mul1"].inputs
        assert "y" in graph.nodes["mul1"].inputs

    def test_skip_when_producer_has_multiple_consumers(self) -> None:
        """Test no fusion when producer has multiple consumers."""
        pass_obj = OperatorFusionPass()
        graph = Graph()

        x = Node(node_id="x", op_type=OpType.INPUT)
        w = Node(node_id="w", op_type=OpType.INPUT)
        b = Node(node_id="b", op_type=OpType.CONSTANT)
        matmul = Node(node_id="matmul1", op_type=OpType.MATMUL)
        add = Node(node_id="add1", op_type=OpType.ADD)
        mul = Node(node_id="mul1", op_type=OpType.MUL)
        out = Node(node_id="out", op_type=OpType.OUTPUT)

        for node in [x, w, b, matmul, add, mul, out]:
            graph.add_node(node)

        graph.add_edge("x", "matmul1")
        graph.add_edge("w", "matmul1")
        graph.add_edge("matmul1", "add1")
        graph.add_edge("matmul1", "mul1")
        graph.add_edge("b", "add1")
        graph.add_edge("add1", "out")
        graph.outputs = ["out"]

        result = pass_obj.run(graph)

        assert result.modified is False
        assert graph.nodes["matmul1"].op_type == OpType.MATMUL
        assert graph.nodes["add1"].op_type == OpType.ADD

    def test_skip_non_cpu_or_dtype_mismatch(self) -> None:
        """Test fusable pattern is skipped for non-CPU or dtype mismatch."""
        pass_obj = OperatorFusionPass()
        graph = Graph()

        x = Node(node_id="x", op_type=OpType.INPUT, device="cuda", dtype="float16")
        w = Node(node_id="w", op_type=OpType.INPUT, device="cuda", dtype="float16")
        matmul = Node(node_id="matmul1", op_type=OpType.MATMUL, device="cuda", dtype="float16")
        add = Node(node_id="add1", op_type=OpType.ADD, device="cuda", dtype="float32")
        out = Node(node_id="out", op_type=OpType.OUTPUT, device="cuda", dtype="float16")

        for node in [x, w, matmul, add, out]:
            graph.add_node(node)

        graph.add_edge("x", "matmul1")
        graph.add_edge("w", "matmul1")
        graph.add_edge("matmul1", "add1")
        graph.add_edge("add1", "out")
        graph.outputs = ["out"]

        result = pass_obj.run(graph)

        assert result.modified is False
        assert graph.nodes["matmul1"].op_type == OpType.MATMUL
        assert graph.nodes["add1"].op_type == OpType.ADD
