"""Test individual pipeline steps/translators."""

import onnx
import ibis
import pytest

from orbital.translate import TRANSLATORS
from orbital.translation.steps.softmax import SoftmaxTranslator
from orbital.translation.steps.imputer import ImputerTranslator
from orbital.translation.steps.argmax import ArgMaxTranslator
from orbital.translation.steps.add import AddTranslator
from orbital.translation.steps.sub import SubTranslator
from orbital.translation.steps.mul import MulTranslator
from orbital.translation.steps.div import DivTranslator
from orbital.translation.steps.identity import IdentityTranslator
from orbital.translation.steps.reshape import ReshapeTranslator
from orbital.translation.steps.matmul import MatMulTranslator
from orbital.translation.steps.cast import CastTranslator, CastLikeTranslator
from orbital.translation.steps.linearclass import LinearClassifierTranslator
from orbital.translation.steps.linearreg import LinearRegressorTranslator
from orbital.translation.steps.scaler import ScalerTranslator
from orbital.translation.steps.onehotencoder import OneHotEncoderTranslator
from orbital.translation.steps.labelencoder import LabelEncoderTranslator
from orbital.translation.steps.where import WhereTranslator
from orbital.translation.steps.zipmap import ZipMapTranslator
from orbital.translation.steps.concat import ConcatTranslator, FeatureVectorizerTranslator
from orbital.translation.steps.gather import GatherTranslator
from orbital.translation.steps.arrayfeatureextractor import ArrayFeatureExtractorTranslator
from orbital.translation.variables import (
    GraphVariables,
    NumericVariablesGroup,
    ValueVariablesGroup,
)
from orbital.translation.optimizer import Optimizer
from orbital.translation.options import TranslationOptions


class TestStepCoverage:
    """Verify that all registered steps have corresponding test classes."""

    def test_all_registered_steps_have_tests(self):
        """Every step registered in TRANSLATORS must have a test class.

        Test classes must follow the naming convention Test{OperationName}Translator.
        This test ensures we don't forget to add tests when implementing new steps.
        """
        import sys

        module = sys.modules[__name__]
        existing_test_classes = {
            name
            for name in dir(module)
            if name.startswith("Test") and isinstance(getattr(module, name), type)
        }

        missing_tests = []
        for operation in sorted(TRANSLATORS.keys()):
            expected_test_class = f"Test{operation}Translator"
            if expected_test_class not in existing_test_classes:
                missing_tests.append((operation, expected_test_class))

        if missing_tests:
            missing_list = "\n".join(
                f"  - {op}: {test_class}" for op, test_class in missing_tests
            )
            pytest.fail(
                f"The following {len(missing_tests)} steps are missing test classes:\n"
                f"{missing_list}\n\n"
                f"Add a test class for each step to test_pipeline_steps.py"
            )


class TestSoftmaxTranslator:
    optimizer = Optimizer(enabled=False)

    def test_softmax_translator_single_input(self):
        """Test SoftmaxTranslator with a single numeric input."""
        table = ibis.memtable({"input": [2.0, 3.0, 4.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = Softmax(input)
            }
        """)

        variables = GraphVariables(table, model)

        translator = SoftmaxTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        # For single input, softmax should return 1.0
        backend = ibis.duckdb.connect()
        computed_value = backend.execute(result)
        assert computed_value == 1.0

    def test_softmax_translator_group_input(self):
        """Test SoftmaxTranslator with a group of numeric inputs."""
        multi_table = ibis.memtable(
            {
                "class_0": [1.0, 2.0, 3.0],
                "class_1": [0.5, 1.5, 2.5],
                "class_2": [2.0, 3.0, 4.0],
            }
        )

        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = Softmax(input)
            }
        """)

        # Use dummy table for GraphVariables since we override the input
        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)

        variables["input"] = NumericVariablesGroup(
            {
                "class_0": multi_table["class_0"],
                "class_1": multi_table["class_1"],
                "class_2": multi_table["class_2"],
            }
        )

        translator = SoftmaxTranslator(
            multi_table,
            model.node[0],
            variables,
            self.optimizer,
            TranslationOptions(),
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        # Should return a NumericVariablesGroup
        assert isinstance(result, NumericVariablesGroup)
        assert len(result) == 3
        assert "class_0" in result
        assert "class_1" in result
        assert "class_2" in result

        # Test that softmax values sum to 1.0 for each row
        backend = ibis.duckdb.connect()

        # backend.execute() returns a pandas Series, so we take the first element
        values = [
            backend.execute(result[class_name])[0]
            for class_name in ["class_0", "class_1", "class_2"]
        ]

        # Verify they sum to approximately 1.0
        total_sum = sum(values)
        assert abs(total_sum - 1.0) < 1e-10, (
            f"Softmax values should sum to 1.0, got {total_sum}"
        )

    def test_softmax_translator_invalid_axis(self):
        """Test that SoftmaxTranslator raises error for unsupported axis."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = Softmax <axis: int = 0> (input)
            }
        """)

        variables = GraphVariables(table, model)

        translator = SoftmaxTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError, match="SoftmaxTranslator supports only axis=-1 or axis=1"
        ):
            translator.process()

    def test_softmax_translator_invalid_input_type(self):
        """Test that SoftmaxTranslator raises error for invalid input type."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = Softmax(input)
            }
        """)

        variables = GraphVariables(table, model)

        # Intentionally set invalid input type to test error handling
        variables["input"] = "invalid_string_input"  # type: ignore[assignment]

        translator = SoftmaxTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError, match="Softmax: The first operand must be a numeric column"
        ):
            translator.process()

    def test_softmax_uses_apply_post_transform(self):
        """Test that SoftmaxTranslator uses the apply_post_transform function."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = Softmax(input)
            }
        """)

        variables = GraphVariables(table, model)

        variables["input"] = NumericVariablesGroup(
            {
                "class_0": ibis.literal(1.0),
                "class_1": ibis.literal(2.0),
            }
        )

        translator = SoftmaxTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, NumericVariablesGroup)


class TestImputerTranslator:
    optimizer = Optimizer(enabled=False)

    def test_imputer_single_column(self):
        """Test ImputerTranslator with a single column input."""
        table = ibis.memtable({"input": [1.0, None, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = ai.onnx.ml.Imputer <imputed_value_floats: floats = [2.0]> (input)
            }
        """)

        variables = GraphVariables(table, model)
        translator = ImputerTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = backend.execute(result)
        assert list(computed) == [1.0, 2.0, 3.0]

    def test_imputer_group_columns(self):
        """Test ImputerTranslator with a group of columns."""
        table = ibis.memtable(
            {
                "col_a": [1.0, None, 3.0],
                "col_b": [None, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = ai.onnx.ml.Imputer <imputed_value_floats: floats = [10.0, 20.0]> (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = ValueVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = ImputerTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["col_a"])) == [1.0, 10.0, 3.0]
        assert list(backend.execute(result["col_b"])) == [20.0, 5.0, 6.0]

    def test_imputer_invalid_imputed_value_type(self):
        """Test ImputerTranslator raises error for non-list imputed values."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = ai.onnx.ml.Imputer <imputed_value_floats: floats = [2.0]> (input)
            }
        """)

        variables = GraphVariables(table, model)

        # Override attributes to test validation
        node = model.node[0]
        translator = ImputerTranslator(
            table, node, variables, self.optimizer, TranslationOptions()
        )
        translator._attributes["imputed_value_floats"] = 2.0  # Invalid: not a list

        with pytest.raises(ValueError, match="imputed_value must be a list or tuple"):
            translator.process()

    def test_imputer_mismatched_column_count(self):
        """Test ImputerTranslator raises error when column count doesn't match."""
        table = ibis.memtable(
            {
                "col_a": [1.0, None, 3.0],
                "col_b": [None, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = ai.onnx.ml.Imputer <imputed_value_floats: floats = [10.0]> (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = ValueVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = ImputerTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="number of imputed values does not match"):
            translator.process()


class TestArgMaxTranslator:
    optimizer = Optimizer(enabled=False)

    def test_argmax_group_columns(self):
        """Test ArgMaxTranslator with a group of columns."""
        table = ibis.memtable(
            {
                "class_0": [1.0, 5.0, 2.0],
                "class_1": [3.0, 2.0, 8.0],
                "class_2": [2.0, 1.0, 3.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (int64[N] output) {
                output = ArgMax <axis: int = 1, keepdims: int = 1> (data)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = NumericVariablesGroup(
            {
                "class_0": table["class_0"],
                "class_1": table["class_1"],
                "class_2": table["class_2"],
            }
        )

        translator = ArgMaxTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = list(backend.execute(result))
        # Row 0: max is class_1 (3.0) -> index 1
        # Row 1: max is class_0 (5.0) -> index 0
        # Row 2: max is class_1 (8.0) -> index 1
        assert computed == [1, 0, 1]

    def test_argmax_single_column(self):
        """Test ArgMaxTranslator raises error for single column input."""
        table = ibis.memtable({"data": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (int64[N] output) {
                output = ArgMax <axis: int = 1> (data)
            }
        """)

        variables = GraphVariables(table, model)

        translator = ArgMaxTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(NotImplementedError, match="can only be applied to a group"):
            translator.process()

    def test_argmax_unsupported_axis(self):
        """Test ArgMaxTranslator raises error for axis != 1."""
        table = ibis.memtable({"data": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (int64[N] output) {
                output = ArgMax <axis: int = 0, keepdims: int = 1> (data)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = NumericVariablesGroup(
            {
                "class_0": table["data"],
                "class_1": table["data"],
            }
        )

        translator = ArgMaxTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(NotImplementedError, match="only supports axis=1"):
            translator.process()

    def test_argmax_unsupported_keepdims(self):
        """Test ArgMaxTranslator raises error for keepdims != 1."""
        table = ibis.memtable({"data": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (int64[N] output) {
                output = ArgMax <axis: int = 1, keepdims: int = 0> (data)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = NumericVariablesGroup(
            {
                "class_0": table["data"],
                "class_1": table["data"],
            }
        )

        translator = ArgMaxTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            NotImplementedError, match="only supports retaining original"
        ):
            translator.process()

    def test_argmax_empty_group(self):
        """Test ArgMaxTranslator raises error for empty group."""
        table = ibis.memtable({"data": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (int64[N] output) {
                output = ArgMax <axis: int = 1, keepdims: int = 1> (data)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = NumericVariablesGroup({})

        translator = ArgMaxTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="requires at least one column"):
            translator.process()

    def test_argmax_single_key_in_group(self):
        """Test ArgMaxTranslator with single key returns that value."""
        table = ibis.memtable({"class_0": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (int64[N] output) {
                output = ArgMax <axis: int = 1, keepdims: int = 1> (data)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = NumericVariablesGroup(
            {
                "class_0": table["class_0"],
            }
        )

        translator = ArgMaxTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables

    def test_argmax_select_last_index(self):
        """Test ArgMaxTranslator with select_last_index=1."""
        table = ibis.memtable(
            {
                "class_0": [3.0, 3.0],
                "class_1": [3.0, 3.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (int64[N] output) {
                output = ArgMax <axis: int = 1, keepdims: int = 1, select_last_index: int = 1> (data)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = NumericVariablesGroup(
            {
                "class_0": table["class_0"],
                "class_1": table["class_1"],
            }
        )

        translator = ArgMaxTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        # Just verify the translator processes without error and exercises the code path
        assert "output" in variables


class TestAddTranslator:
    optimizer = Optimizer(enabled=False)

    def test_add_single_column(self):
        """Test AddTranslator with a single column input."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] add_value = {5.0}>
            {
                output = Add(input, add_value)
            }
        """)

        variables = GraphVariables(table, model)
        translator = AddTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = list(backend.execute(result))
        assert computed == [6.0, 7.0, 8.0]

    def test_add_group_columns(self):
        """Test AddTranslator with a group of columns."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [10.0, 20.0, 30.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[2] add_values = {5.0, 100.0}>
            {
                output = Add(input, add_values)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = AddTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["col_a"])) == [6.0, 7.0, 8.0]
        assert list(backend.execute(result["col_b"])) == [110.0, 120.0, 130.0]

    def test_add_invalid_non_numeric(self):
        """Test AddTranslator raises error for non-numeric operand."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] add_value = {5.0}>
            {
                output = Add(input, add_value)
            }
        """)

        variables = GraphVariables(table, model)
        variables["input"] = "not_a_numeric_value"  # type: ignore[assignment]

        translator = AddTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="first operand must be a numeric value"):
            translator.process()

    def test_add_mismatched_column_count(self):
        """Test AddTranslator raises error when column count doesn't match."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [10.0, 20.0, 30.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] add_values = {5.0}>
            {
                output = Add(input, add_values)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = AddTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="same number of values"):
            translator.process()

    def test_add_single_column_requires_single_value(self):
        """Test AddTranslator raises error when single column given multiple values."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[2] add_values = {5.0, 10.0}>
            {
                output = Add(input, add_values)
            }
        """)

        variables = GraphVariables(table, model)

        translator = AddTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="must contain exactly 1 value"):
            translator.process()

    def test_add_second_operand_not_constant(self):
        """Test AddTranslator raises error when second operand is not a constant."""
        table = ibis.memtable(
            {
                "input": [1.0, 2.0, 3.0],
                "other": [5.0, 5.0, 5.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input, float[N] other) => (float[N] output) {
                output = Add(input, other)
            }
        """)

        variables = GraphVariables(table, model)

        translator = AddTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(NotImplementedError, match="must be a constant list"):
            translator.process()


class TestSubTranslator:
    optimizer = Optimizer(enabled=False)

    def test_sub_single_column(self):
        """Test SubTranslator with a single column input."""
        table = ibis.memtable({"input": [10.0, 20.0, 30.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] sub_value = {3.0}>
            {
                output = Sub(input, sub_value)
            }
        """)

        variables = GraphVariables(table, model)
        translator = SubTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = list(backend.execute(result))
        assert computed == [7.0, 17.0, 27.0]

    def test_sub_group_columns(self):
        """Test SubTranslator with a group of columns."""
        table = ibis.memtable(
            {
                "col_a": [10.0, 20.0, 30.0],
                "col_b": [100.0, 200.0, 300.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[2] sub_values = {3.0, 50.0}>
            {
                output = Sub(input, sub_values)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = SubTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["col_a"])) == [7.0, 17.0, 27.0]
        assert list(backend.execute(result["col_b"])) == [50.0, 150.0, 250.0]

    def test_sub_invalid_non_numeric(self):
        """Test SubTranslator raises error for non-numeric operand."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] sub_value = {5.0}>
            {
                output = Sub(input, sub_value)
            }
        """)

        variables = GraphVariables(table, model)
        variables["input"] = "not_a_numeric_value"  # type: ignore[assignment]

        translator = SubTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="first operand must be a numeric value"):
            translator.process()

    def test_sub_single_column_requires_single_value(self):
        """Test SubTranslator raises error when single column given multiple values."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[2] sub_values = {5.0, 10.0}>
            {
                output = Sub(input, sub_values)
            }
        """)

        variables = GraphVariables(table, model)

        translator = SubTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="must contain exactly 1 value"):
            translator.process()

    def test_sub_mismatched_column_count(self):
        """Test SubTranslator raises error when column count doesn't match."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [10.0, 20.0, 30.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] sub_values = {5.0}>
            {
                output = Sub(input, sub_values)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = SubTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(AssertionError, match="must match the number of fields"):
            translator.process()

    def test_sub_second_operand_not_constant(self):
        """Test SubTranslator raises error when second operand is not a constant."""
        table = ibis.memtable(
            {
                "input": [1.0, 2.0, 3.0],
                "other": [5.0, 5.0, 5.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input, float[N] other) => (float[N] output) {
                output = Sub(input, other)
            }
        """)

        variables = GraphVariables(table, model)

        translator = SubTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(NotImplementedError, match="must be a constant list"):
            translator.process()


class TestMulTranslator:
    optimizer = Optimizer(enabled=False)

    def test_mul_single_column(self):
        """Test MulTranslator with a single column input."""
        table = ibis.memtable({"input": [2.0, 3.0, 4.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] mul_value = {5.0}>
            {
                output = Mul(input, mul_value)
            }
        """)

        variables = GraphVariables(table, model)
        translator = MulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = list(backend.execute(result))
        assert computed == [10.0, 15.0, 20.0]

    def test_mul_group_columns(self):
        """Test MulTranslator with a group of columns."""
        table = ibis.memtable(
            {
                "col_a": [2.0, 3.0, 4.0],
                "col_b": [10.0, 20.0, 30.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[2] mul_values = {3.0, 2.0}>
            {
                output = Mul(input, mul_values)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = MulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["col_a"])) == [6.0, 9.0, 12.0]
        assert list(backend.execute(result["col_b"])) == [20.0, 40.0, 60.0]

    def test_mul_invalid_non_numeric(self):
        """Test MulTranslator raises error for non-numeric operand."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] mul_value = {5.0}>
            {
                output = Mul(input, mul_value)
            }
        """)

        variables = GraphVariables(table, model)
        variables["input"] = "not_a_numeric_value"  # type: ignore[assignment]

        translator = MulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="first operand must be a numeric value"):
            translator.process()

    def test_mul_mismatched_column_count(self):
        """Test MulTranslator raises error when column count doesn't match."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [10.0, 20.0, 30.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] mul_values = {5.0}>
            {
                output = Mul(input, mul_values)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = MulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="same number of values"):
            translator.process()

    def test_mul_single_column_requires_single_value(self):
        """Test MulTranslator raises error when single column given multiple values."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[2] mul_values = {5.0, 10.0}>
            {
                output = Mul(input, mul_values)
            }
        """)

        variables = GraphVariables(table, model)

        translator = MulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="must contain exactly 1 value"):
            translator.process()

    def test_mul_second_operand_not_constant(self):
        """Test MulTranslator raises error when second operand is not a constant."""
        table = ibis.memtable(
            {
                "input": [1.0, 2.0, 3.0],
                "other": [5.0, 5.0, 5.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input, float[N] other) => (float[N] output) {
                output = Mul(input, other)
            }
        """)

        variables = GraphVariables(table, model)

        translator = MulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(NotImplementedError, match="must be a constant list"):
            translator.process()


class TestDivTranslator:
    optimizer = Optimizer(enabled=False)

    def test_div_single_column(self):
        """Test DivTranslator with a single column input."""
        table = ibis.memtable({"input": [10.0, 20.0, 30.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] div_value = {2.0}>
            {
                output = Div(input, div_value)
            }
        """)

        variables = GraphVariables(table, model)
        translator = DivTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = list(backend.execute(result))
        assert computed == [5.0, 10.0, 15.0]

    def test_div_group_columns_with_matching_values(self):
        """Test DivTranslator with group of columns and matching divisor values."""
        table = ibis.memtable(
            {
                "col_a": [10.0, 20.0, 30.0],
                "col_b": [100.0, 200.0, 300.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[2] div_values = {2.0, 10.0}>
            {
                output = Div(input, div_values)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = DivTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["col_a"])) == [5.0, 10.0, 15.0]
        assert list(backend.execute(result["col_b"])) == [10.0, 20.0, 30.0]

    def test_div_group_columns_broadcast_single_value(self):
        """Test DivTranslator with group of columns and single divisor (broadcast)."""
        table = ibis.memtable(
            {
                "col_a": [10.0, 20.0, 30.0],
                "col_b": [100.0, 200.0, 300.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] div_value = {10.0}>
            {
                output = Div(input, div_value)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = DivTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["col_a"])) == [1.0, 2.0, 3.0]
        assert list(backend.execute(result["col_b"])) == [10.0, 20.0, 30.0]

    def test_div_invalid_non_numeric_single(self):
        """Test DivTranslator raises error for non-numeric single operand."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] div_value = {5.0}>
            {
                output = Div(input, div_value)
            }
        """)

        variables = GraphVariables(table, model)
        variables["input"] = "not_a_numeric_value"  # type: ignore[assignment]

        translator = DivTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="first operand must be a numeric value"):
            translator.process()

    def test_div_mismatched_column_count(self):
        """Test DivTranslator raises error when column count doesn't match values."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [10.0, 20.0, 30.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[3] div_values = {5.0, 10.0, 15.0}>
            {
                output = Div(input, div_values)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = DivTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="must match the number of columns"):
            translator.process()

    def test_div_single_column_requires_single_value(self):
        """Test DivTranslator raises error when single column given multiple values."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[2] div_values = {5.0, 10.0}>
            {
                output = Div(input, div_values)
            }
        """)

        variables = GraphVariables(table, model)

        translator = DivTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="must contain only one value"):
            translator.process()

    def test_div_second_operand_not_constant(self):
        """Test DivTranslator raises error when second operand is not a constant."""
        table = ibis.memtable(
            {
                "input": [1.0, 2.0, 3.0],
                "other": [5.0, 5.0, 5.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input, float[N] other) => (float[N] output) {
                output = Div(input, other)
            }
        """)

        variables = GraphVariables(table, model)

        translator = DivTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(NotImplementedError, match="must be a constant list"):
            translator.process()


class TestIdentityTranslator:
    optimizer = Optimizer(enabled=False)

    def test_identity_single_column_passthrough(self):
        """Test IdentityTranslator passes a single column through unchanged."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = Identity(input)
            }
        """)

        variables = GraphVariables(table, model)
        translator = IdentityTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        # Verify output equals input
        backend = ibis.duckdb.connect()
        input_values = list(backend.execute(table["input"]))
        output_values = list(backend.execute(result))
        assert output_values == input_values

    def test_identity_group_columns_passthrough(self):
        """Test IdentityTranslator passes a NumericVariablesGroup through unchanged."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
                "col_c": [7.0, 8.0, 9.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = Identity(input)
            }
        """)

        # Use dummy table for GraphVariables since we override the input
        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
                "col_c": table["col_c"],
            }
        )

        translator = IdentityTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        # Should return the same NumericVariablesGroup
        assert isinstance(result, NumericVariablesGroup)
        assert len(result) == 3
        assert "col_a" in result
        assert "col_b" in result
        assert "col_c" in result

        # Verify all columns preserved with same values
        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["col_a"])) == [1.0, 2.0, 3.0]
        assert list(backend.execute(result["col_b"])) == [4.0, 5.0, 6.0]
        assert list(backend.execute(result["col_c"])) == [7.0, 8.0, 9.0]


class TestReshapeTranslator:
    optimizer = Optimizer(enabled=False)

    def test_reshape_single_column_to_single_column(self):
        """Test ReshapeTranslator with shape=[-1] on single column (passes through)."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <int64[1] shape = {-1}>
            {
                output = Reshape(input, shape)
            }
        """)

        variables = GraphVariables(table, model)
        translator = ReshapeTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        # Verify output equals input (passthrough)
        backend = ibis.duckdb.connect()
        input_values = list(backend.execute(table["input"]))
        output_values = list(backend.execute(result))
        assert output_values == input_values

    def test_reshape_group_to_same_size(self):
        """Test ReshapeTranslator with shape=[-1, N] on N columns (passes through)."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
                "col_c": [7.0, 8.0, 9.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <int64[2] shape = {-1, 3}>
            {
                output = Reshape(input, shape)
            }
        """)

        # Use dummy table for GraphVariables since we override the input
        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
                "col_c": table["col_c"],
            }
        )

        translator = ReshapeTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        # Should return the same NumericVariablesGroup (passthrough)
        assert isinstance(result, NumericVariablesGroup)
        assert len(result) == 3

        # Verify all columns preserved with same values
        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["col_a"])) == [1.0, 2.0, 3.0]
        assert list(backend.execute(result["col_b"])) == [4.0, 5.0, 6.0]
        assert list(backend.execute(result["col_c"])) == [7.0, 8.0, 9.0]

    def test_reshape_requires_integer_shape(self):
        """Test ReshapeTranslator raises error when shape is not integers."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] shape = {-1.0}>
            {
                output = Reshape(input, shape)
            }
        """)

        variables = GraphVariables(table, model)
        translator = ReshapeTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            NotImplementedError, match="requires integer values for the shape"
        ):
            translator.process()

    def test_reshape_cannot_change_row_count(self):
        """Test ReshapeTranslator raises error when shape[0] != -1."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <int64[2] shape = {3, 1}>
            {
                output = Reshape(input, shape)
            }
        """)

        variables = GraphVariables(table, model)
        translator = ReshapeTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            NotImplementedError, match="Reshape can't change the number of rows"
        ):
            translator.process()

    def test_reshape_unsupported_shape(self):
        """Test ReshapeTranslator raises error for unsupported shape combinations."""
        # Test case: shape=[-1, 2] but input has 3 columns
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
                "col_c": [7.0, 8.0, 9.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <int64[2] shape = {-1, 2}>
            {
                output = Reshape(input, shape)
            }
        """)

        # Use dummy table for GraphVariables since we override the input
        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
                "col_c": table["col_c"],
            }
        )

        translator = ReshapeTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="Reshape shape=\\[-1, 2\\] not supported"):
            translator.process()


class TestMatMulTranslator:
    optimizer = Optimizer(enabled=False)

    def test_matmul_single_column_times_1d_weight_vector_length_1(self):
        """Test MatMulTranslator with single column times 1D weight vector (coefficient vector of length 1)."""
        table = ibis.memtable({"input": [2.0, 3.0, 4.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[1] weights = {5.0}>
            {
                output = MatMul(input, weights)
            }
        """)

        variables = GraphVariables(table, model)
        translator = MatMulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = list(backend.execute(result))
        assert computed == [10.0, 15.0, 20.0]

    def test_matmul_group_columns_times_2d_weight_matrix_multiple_outputs(self):
        """Test MatMulTranslator with group of columns times 2D weight matrix (produces multiple outputs)."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
            }
        )
        # Flattened row-major: [[0.5, 1.0, 1.5], [2.0, 2.5, 3.0]]
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[6] weights = {0.5, 1.0, 1.5, 2.0, 2.5, 3.0}>
            {
                output = MatMul(input, weights)
            }
        """)

        model.initializer[0].dims[:] = [2, 3]

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = MatMulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)
        assert len(result) == 3

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["out_0"])) == [8.5, 11.0, 13.5]
        assert list(backend.execute(result["out_1"])) == [11.0, 14.5, 18.0]
        assert list(backend.execute(result["out_2"])) == [13.5, 18.0, 22.5]

    def test_matmul_group_columns_times_1d_weight_vector_single_output(self):
        """Test MatMulTranslator with group of columns times 1D weight vector (produces single output)."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
                "col_c": [7.0, 8.0, 9.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[3] weights = {2.0, 3.0, 4.0}>
            {
                output = MatMul(input, weights)
            }
        """)

        model.initializer[0].dims[:] = [3]

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
                "col_c": table["col_c"],
            }
        )

        translator = MatMulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = list(backend.execute(result))
        assert computed == [42.0, 51.0, 60.0]

    def test_matmul_with_optimizer_enabled(self):
        """Test MatMulTranslator with optimizer folding enabled."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[2] weights = {2.0, 3.0}>
            {
                output = MatMul(input, weights)
            }
        """)

        model.initializer[0].dims[:] = [2]

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = MatMulTranslator(
            table,
            model.node[0],
            variables,
            Optimizer(enabled=True),
            TranslationOptions(),
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = list(backend.execute(result))
        assert computed == [14.0, 19.0, 24.0]

    def test_matmul_error_non_numeric_input_columns(self):
        """Test MatMulTranslator raises error for non-numeric input columns."""
        table = ibis.memtable({"input": ["a", "b", "c"]})
        model = onnx.parser.parse_graph("""
            agraph (string[N] input) => (float[N] output)
            <float[1] weights = {2.0}>
            {
                output = MatMul(input, weights)
            }
        """)

        model.initializer[0].dims[:] = [1]
        variables = GraphVariables(table, model)

        translator = MatMulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="first operand must be a numeric column"):
            translator.process()

    def test_matmul_error_weight_matrix_dimension_mismatch(self):
        """Test MatMulTranslator raises error for weight matrix dimension mismatch."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
            }
        )
        # Weight matrix [3, 2] expects 3 features but input provides only 2
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[6] weights = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}>
            {
                output = MatMul(input, weights)
            }
        """)

        model.initializer[0].dims[:] = [3, 2]

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = MatMulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="Mismatch: number of features"):
            translator.process()

    def test_matmul_error_unsupported_weight_tensor_rank(self):
        """Test MatMulTranslator raises error for unsupported weight tensor shapes (rank > 2)."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[8] weights = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0}>
            {
                output = MatMul(input, weights)
            }
        """)

        model.initializer[0].dims[:] = [2, 2, 2]  # 3D tensor - unsupported

        variables = GraphVariables(table, model)

        translator = MatMulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError, match="coefficient tensor rank > 2 is not supported"
        ):
            translator.process()

    def test_matmul_error_single_column_with_coefficient_vector_longer_than_1(self):
        """Test MatMulTranslator raises error for single column with coefficient vector longer than 1."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[3] weights = {2.0, 3.0, 4.0}>
            {
                output = MatMul(input, weights)
            }
        """)

        # Set weight vector dimensions [3] - but we have single column
        model.initializer[0].dims[:] = [3]

        variables = GraphVariables(table, model)

        translator = MatMulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError,
            match="Expected coefficient vector of length 1 for single operand",
        ):
            translator.process()

    def test_matmul_group_columns_times_2d_weight_matrix_single_output(self):
        """Test MatMulTranslator with group of columns times 2D weight matrix with single output."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[2] weights = {2.0, 3.0}>
            {
                output = MatMul(input, weights)
            }
        """)

        model.initializer[0].dims[:] = [2, 1]

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = MatMulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert not isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        computed = list(backend.execute(result))
        assert computed == [14.0, 19.0, 24.0]

    def test_matmul_single_column_times_2d_weight_matrix_multiple_outputs(self):
        """Test MatMulTranslator with single column times 2D weight matrix [1, N] (produces multiple outputs)."""
        table = ibis.memtable({"input": [2.0, 3.0, 4.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output)
            <float[3] weights = {2.0, 3.0, 4.0}>
            {
                output = MatMul(input, weights)
            }
        """)

        model.initializer[0].dims[:] = [1, 3]

        variables = GraphVariables(table, model)

        translator = MatMulTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)
        assert len(result) == 3

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["out_0"])) == [4.0, 6.0, 8.0]
        assert list(backend.execute(result["out_1"])) == [6.0, 9.0, 12.0]
        assert list(backend.execute(result["out_2"])) == [8.0, 12.0, 16.0]


class TestCastTranslator:
    optimizer = Optimizer(enabled=False)

    def test_cast_single_column_to_float32(self):
        """Test CastTranslator with single column cast to float32."""
        table = ibis.memtable({"input": [1, 2, 3]})  # int values
        model = onnx.parser.parse_graph("""
            agraph (int64[N] input) => (float[N] output) {
                output = Cast <to: int = 1> (input)
            }
        """)

        variables = GraphVariables(table, model)
        translator = CastTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = backend.execute(result)
        assert list(computed) == [1.0, 2.0, 3.0]

    def test_cast_single_column_to_float64(self):
        """Test CastTranslator with single column cast to float64."""
        table = ibis.memtable({"input": [1, 2, 3]})  # int values
        model = onnx.parser.parse_graph("""
            agraph (int64[N] input) => (double[N] output) {
                output = Cast <to: int = 11> (input)
            }
        """)

        variables = GraphVariables(table, model)
        translator = CastTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = backend.execute(result)
        assert list(computed) == [1.0, 2.0, 3.0]

    def test_cast_single_column_to_int64(self):
        """Test CastTranslator with single column cast to int64."""
        table = ibis.memtable({"input": [1.2, 2.7, 3.9]})  # float values
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (int64[N] output) {
                output = Cast <to: int = 7> (input)
            }
        """)

        variables = GraphVariables(table, model)
        translator = CastTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = backend.execute(result)
        assert list(computed) == [1, 3, 4]

    def test_cast_single_column_to_string(self):
        """Test CastTranslator with single column cast to string."""
        table = ibis.memtable({"input": [1, 2, 3]})
        model = onnx.parser.parse_graph("""
            agraph (int64[N] input) => (string[N] output) {
                output = Cast <to: int = 8> (input)
            }
        """)

        variables = GraphVariables(table, model)
        options = TranslationOptions(allow_text_tensors=True)
        translator = CastTranslator(
            table, model.node[0], variables, self.optimizer, options
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = backend.execute(result)
        assert list(computed) == ["1", "2", "3"]

    def test_cast_group_of_columns(self):
        """Test CastTranslator with a group of columns."""
        table = ibis.memtable(
            {
                "col_a": [1, 2, 3],
                "col_b": [4, 5, 6],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (int64[N] input) => (float[N] output) {
                output = Cast <to: int = 1> (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1]}), model)
        variables["input"] = ValueVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = CastTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["col_a"])) == [1.0, 2.0, 3.0]
        assert list(backend.execute(result["col_b"])) == [4.0, 5.0, 6.0]

    def test_skip_cast_to_string_single_column(self):
        """Test CastTranslator skips cast to string when allow_text_tensors=False."""
        table = ibis.memtable({"input": [1, 2, 3]})
        model = onnx.parser.parse_graph("""
            agraph (int64[N] input) => (string[N] output) {
                output = Cast <to: int = 8> (input)
            }
        """)

        variables = GraphVariables(table, model)
        options = TranslationOptions(allow_text_tensors=False)
        translator = CastTranslator(
            table, model.node[0], variables, self.optimizer, options
        )
        translator.process()

        # Should skip the cast and return the original input
        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = backend.execute(result)
        # Should remain int64, not cast to string
        assert list(computed) == [1, 2, 3]

    def test_skip_cast_to_string_column_group(self):
        """Test CastTranslator skips cast to string for column group when allow_text_tensors=False."""
        table = ibis.memtable(
            {
                "col_a": [1, 2, 3],
                "col_b": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (string[N] output) {
                output = Cast <to: int = 8> (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = ValueVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        options = TranslationOptions(allow_text_tensors=False)
        translator = CastTranslator(
            table, model.node[0], variables, self.optimizer, options
        )
        translator.process()

        # Should skip the cast and return the original input group
        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        # Should remain original types, not cast to string
        assert list(backend.execute(result["col_a"])) == [1, 2, 3]
        assert list(backend.execute(result["col_b"])) == [4.0, 5.0, 6.0]

    def test_cast_unsupported_target_type(self):
        """Test CastTranslator raises error for unsupported target type."""
        table = ibis.memtable({"input": [1, 2, 3]})
        model = onnx.parser.parse_graph("""
            agraph (int64[N] input) => (float[N] output) {
                output = Cast <to: int = 999> (input)
            }
        """)

        variables = GraphVariables(table, model)
        translator = CastTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(NotImplementedError, match="Cast: type 999 not supported"):
            translator.process()


class TestCastLikeTranslator:
    optimizer = Optimizer(enabled=False)

    def test_cast_group_to_match_single_column_type(self):
        """Test CastLikeTranslator casts group of columns to match single column type."""
        table = ibis.memtable(
            {
                "col_a": [1, 2, 3],  # int64
                "col_b": [4, 5, 6],  # int64
                "target": [1.0, 2.0, 3.0],  # float64
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (int64[N] input, float[N] target_type) => (float[N] output) {
                output = CastLike(input, target_type)
            }
        """)

        variables = GraphVariables(
            ibis.memtable({"input": [1], "target_type": [1.0]}), model
        )
        variables["input"] = ValueVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )
        variables["target_type"] = table["target"]

        translator = CastLikeTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        # Should match target column type (float64)
        assert list(backend.execute(result["col_a"])) == [1.0, 2.0, 3.0]
        assert list(backend.execute(result["col_b"])) == [4.0, 5.0, 6.0]

    def test_castlike_error_input_is_single_column(self):
        """Test CastLikeTranslator raises error when input is single column."""
        table = ibis.memtable(
            {
                "input": [1, 2, 3],
                "target_type": [1.0, 2.0, 3.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (int64[N] input, float[N] target_type) => (float[N] output) {
                output = CastLike(input, target_type)
            }
        """)

        variables = GraphVariables(table, model)

        translator = CastLikeTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            NotImplementedError,
            match="CastLike currently only supports casting a group of columns",
        ):
            translator.process()

    def test_castlike_error_target_is_group(self):
        """Test CastLikeTranslator raises error when target is a group."""
        table = ibis.memtable(
            {
                "col_a": [1, 2, 3],
                "col_b": [4, 5, 6],
                "target_a": [1.0, 2.0, 3.0],
                "target_b": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (int64[N] input, float[N] target_type) => (float[N] output) {
                output = CastLike(input, target_type)
            }
        """)

        variables = GraphVariables(
            ibis.memtable({"input": [1], "target_type": [1.0]}), model
        )
        variables["input"] = ValueVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )
        variables["target_type"] = ValueVariablesGroup(
            {
                "target_a": table["target_a"],
                "target_b": table["target_b"],
            }
        )

        translator = CastLikeTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            NotImplementedError,
            match="CastLike currently only supports casting to a single column type, not a group",
        ):
            translator.process()


class TestLinearClassifierTranslator:
    optimizer = Optimizer(enabled=False)

    def test_linear_classifier_binary_classification(self):
        """Test LinearClassifierTranslator with binary classification."""
        table = ibis.memtable(
            {
                "feature1": [1.0, 2.0, 3.0],
                "feature2": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (int64[N] Y, float[N,2] Z)
            {
                Y, Z = ai.onnx.ml.LinearClassifier <
                    coefficients: floats = [0.5, 0.5, -0.5, -0.5],
                    intercepts: floats = [1.0, -1.0],
                    classlabels_ints: ints = [0, 1]
                > (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = ValueVariablesGroup(
            {
                "feature1": table["feature1"],
                "feature2": table["feature2"],
            }
        )

        translator = LinearClassifierTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "Y" in variables
        assert "Z" in variables

        predictions = variables.peek_variable("Y")
        scores = variables.peek_variable("Z")

        assert isinstance(scores, ValueVariablesGroup)
        assert "0" in scores
        assert "1" in scores

        backend = ibis.duckdb.connect()
        # Class 0 has higher score so it should be predicted
        assert backend.execute(predictions)[0] == "0"

    def test_linear_classifier_multiclass(self):
        """Test LinearClassifierTranslator with multi-class classification (3+ classes)."""
        table = ibis.memtable(
            {
                "feature1": [1.0, 2.0, 3.0],
                "feature2": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (int64[N] Y, float[N,3] Z)
            {
                Y, Z = ai.onnx.ml.LinearClassifier <
                    coefficients: floats = [1.0, 0.0, 0.0, 1.0, -1.0, -1.0],
                    intercepts: floats = [0.0, 0.5, -0.5],
                    classlabels_ints: ints = [0, 1, 2]
                > (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = ValueVariablesGroup(
            {
                "feature1": table["feature1"],
                "feature2": table["feature2"],
            }
        )

        translator = LinearClassifierTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        predictions = variables.peek_variable("Y")
        scores = variables.peek_variable("Z")

        assert isinstance(scores, ValueVariablesGroup)
        assert len(scores) == 3

        backend = ibis.duckdb.connect()
        # Class 1 has highest score for first row
        assert backend.execute(predictions)[0] == "1"

    def test_linear_classifier_with_intercepts(self):
        """Test LinearClassifierTranslator properly adds intercepts."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (int64[N] Y, float[N,2] Z)
            {
                Y, Z = ai.onnx.ml.LinearClassifier <
                    coefficients: floats = [1.0, -1.0],
                    intercepts: floats = [10.0, 20.0],
                    classlabels_ints: ints = [0, 1]
                > (input)
            }
        """)

        variables = GraphVariables(table, model)

        translator = LinearClassifierTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        scores = variables.peek_variable("Z")

        backend = ibis.duckdb.connect()
        # Verify intercepts are properly added to scores
        assert backend.execute(scores["0"])[0] == 11.0
        assert backend.execute(scores["1"])[0] == 19.0

    def test_linear_classifier_post_transform_logistic(self):
        """Test LinearClassifierTranslator with post_transform='LOGISTIC'."""
        table = ibis.memtable({"input": [0.0, 1.0, -1.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (int64[N] Y, float[N,2] Z)
            {
                Y, Z = ai.onnx.ml.LinearClassifier <
                    coefficients: floats = [1.0, -1.0],
                    intercepts: floats = [0.0, 0.0],
                    classlabels_ints: ints = [0, 1],
                    post_transform: string = "LOGISTIC"
                > (input)
            }
        """)

        variables = GraphVariables(table, model)

        translator = LinearClassifierTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        scores = variables.peek_variable("Z")

        backend = ibis.duckdb.connect()
        # Logistic of 0.0 yields 0.5
        assert abs(backend.execute(scores["0"])[0] - 0.5) < 1e-6

    def test_linear_classifier_post_transform_softmax(self):
        """Test LinearClassifierTranslator with post_transform='SOFTMAX'."""
        table = ibis.memtable(
            {
                "feature1": [1.0, 2.0, 3.0],
                "feature2": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (int64[N] Y, float[N,2] Z)
            {
                Y, Z = ai.onnx.ml.LinearClassifier <
                    coefficients: floats = [1.0, 0.0, 0.0, 1.0],
                    intercepts: floats = [0.0, 0.0],
                    classlabels_ints: ints = [0, 1],
                    post_transform: string = "SOFTMAX"
                > (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = ValueVariablesGroup(
            {
                "feature1": table["feature1"],
                "feature2": table["feature2"],
            }
        )

        translator = LinearClassifierTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        scores = variables.peek_variable("Z")
        assert isinstance(scores, ValueVariablesGroup)
        assert "0" in scores
        assert "1" in scores

    def test_linear_classifier_string_labels(self):
        """Test LinearClassifierTranslator with classlabels_strings instead of classlabels_ints."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (string[N] Y, float[N,2] Z)
            {
                Y, Z = ai.onnx.ml.LinearClassifier <
                    coefficients: floats = [1.0, -1.0],
                    intercepts: floats = [0.0, 0.0],
                    classlabels_strings: strings = ["cat", "dog"]
                > (input)
            }
        """)

        variables = GraphVariables(table, model)

        translator = LinearClassifierTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        predictions = variables.peek_variable("Y")
        scores = variables.peek_variable("Z")

        assert "cat" in scores
        assert "dog" in scores

        backend = ibis.duckdb.connect()
        # "cat" has higher score so it should be predicted
        assert backend.execute(predictions)[0] == "cat"

    def test_linear_classifier_single_input(self):
        """Test LinearClassifierTranslator when input is a single column (not a VariablesGroup)."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (int64[N] Y, float[N,2] Z)
            {
                Y, Z = ai.onnx.ml.LinearClassifier <
                    coefficients: floats = [2.0, -2.0],
                    intercepts: floats = [0.0, 0.0],
                    classlabels_ints: ints = [0, 1]
                > (input)
            }
        """)

        variables = GraphVariables(table, model)

        translator = LinearClassifierTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        predictions = variables.peek_variable("Y")

        backend = ibis.duckdb.connect()
        # Class 0: 1.0*2.0 = 2.0
        # Class 1: 1.0*(-2.0) = -2.0
        # Prediction should be class 0
        assert backend.execute(predictions)[0] == "0"

    def test_linear_classifier_missing_classlabels(self):
        """Test LinearClassifierTranslator raises error when no classlabels defined."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (int64[N] Y, float[N,2] Z)
            {
                Y, Z = ai.onnx.ml.LinearClassifier <
                    coefficients: floats = [1.0, -1.0],
                    intercepts: floats = [0.0, 0.0]
                > (input)
            }
        """)

        variables = GraphVariables(table, model)

        translator = LinearClassifierTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError,
            match="LinearClassifier: classlabels_ints or classlabels_strings must be defined",
        ):
            translator.process()

    def test_linear_classifier_coefficient_mismatch(self):
        """Test LinearClassifierTranslator raises error when coefficients length doesn't match classes × features."""
        table = ibis.memtable(
            {
                "feature1": [1.0, 2.0, 3.0],
                "feature2": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (int64[N] Y, float[N,2] Z)
            {
                Y, Z = ai.onnx.ml.LinearClassifier <
                    coefficients: floats = [1.0, 2.0, 3.0],
                    classlabels_ints: ints = [0, 1]
                > (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = ValueVariablesGroup(
            {
                "feature1": table["feature1"],
                "feature2": table["feature2"],
            }
        )

        translator = LinearClassifierTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError,
            match="Coefficients length must equal number of classes × number of input fields",
        ):
            translator.process()

    def test_linear_classifier_multi_class_mode_not_implemented(self):
        """Test LinearClassifierTranslator raises error when multi_class != 0."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (int64[N] Y, float[N,2] Z)
            {
                Y, Z = ai.onnx.ml.LinearClassifier <
                    coefficients: floats = [1.0, -1.0],
                    classlabels_ints: ints = [0, 1],
                    multi_class: int = 1
                > (input)
            }
        """)

        variables = GraphVariables(table, model)

        translator = LinearClassifierTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            NotImplementedError, match="Multi-class classification is not implemented"
        ):
            translator.process()


class TestLinearRegressorTranslator:
    optimizer = Optimizer(enabled=False)

    def test_linear_regressor_single_target(self):
        """Test LinearRegressorTranslator with single target regression (targets=1)."""
        table = ibis.memtable(
            {
                "feature1": [1.0, 2.0, 3.0],
                "feature2": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] Y)
            {
                Y = ai.onnx.ml.LinearRegressor <
                    coefficients: floats = [0.5, 0.5],
                    intercepts: floats = [10.0],
                    targets: int = 1
                > (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = ValueVariablesGroup(
            {
                "feature1": table["feature1"],
                "feature2": table["feature2"],
            }
        )

        translator = LinearRegressorTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "Y" in variables
        result = variables.peek_variable("Y")

        backend = ibis.duckdb.connect()
        assert backend.execute(result["target_0"])[0] == 12.5

    def test_linear_regressor_multi_target(self):
        """Test LinearRegressorTranslator with multi-target regression (targets=2+)."""
        table = ibis.memtable(
            {
                "feature1": [1.0, 2.0, 3.0],
                "feature2": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N,2] Y)
            {
                Y = ai.onnx.ml.LinearRegressor <
                    coefficients: floats = [1.0, 0.0, 0.0, 1.0],
                    intercepts: floats = [5.0, 10.0],
                    targets: int = 2
                > (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = ValueVariablesGroup(
            {
                "feature1": table["feature1"],
                "feature2": table["feature2"],
            }
        )

        translator = LinearRegressorTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        result = variables.peek_variable("Y")
        assert isinstance(result, ValueVariablesGroup)
        assert "target_0" in result
        assert "target_1" in result

        backend = ibis.duckdb.connect()
        assert backend.execute(result["target_0"])[0] == 6.0
        assert backend.execute(result["target_1"])[0] == 14.0

    def test_linear_regressor_with_intercepts(self):
        """Test LinearRegressorTranslator properly adds intercepts."""
        table = ibis.memtable({"feature": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] Y)
            {
                Y = ai.onnx.ml.LinearRegressor <
                    coefficients: floats = [2.0],
                    intercepts: floats = [100.0]
                > (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = ValueVariablesGroup({"feature": table["feature"]})

        translator = LinearRegressorTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        result = variables.peek_variable("Y")

        backend = ibis.duckdb.connect()
        # Verify intercept is properly added
        assert backend.execute(result["target_0"])[0] == 102.0

    def test_linear_regressor_no_intercepts(self):
        """Test LinearRegressorTranslator without intercepts."""
        table = ibis.memtable({"feature": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] Y)
            {
                Y = ai.onnx.ml.LinearRegressor <
                    coefficients: floats = [3.0]
                > (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = ValueVariablesGroup({"feature": table["feature"]})

        translator = LinearRegressorTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        result = variables.peek_variable("Y")

        backend = ibis.duckdb.connect()
        # Default intercept is 0
        assert backend.execute(result["target_0"])[0] == 3.0

    def test_linear_regressor_single_column_input(self):
        """Test LinearRegressorTranslator with single column input (not VariablesGroup)."""
        table = ibis.memtable({"input": [2.0, 3.0, 4.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] Y)
            {
                Y = ai.onnx.ml.LinearRegressor <
                    coefficients: floats = [5.0],
                    intercepts: floats = [1.0]
                > (input)
            }
        """)

        variables = GraphVariables(table, model)

        translator = LinearRegressorTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        result = variables.peek_variable("Y")

        backend = ibis.duckdb.connect()
        # Prediction: 2.0*5.0 + 1.0 = 11.0
        assert backend.execute(result)[0] == 11.0

    def test_linear_regressor_coefficient_mismatch(self):
        """Test LinearRegressorTranslator raises error when coefficients length mismatch."""
        table = ibis.memtable(
            {
                "feature1": [1.0, 2.0, 3.0],
                "feature2": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] Y)
            {
                Y = ai.onnx.ml.LinearRegressor <
                    coefficients: floats = [1.0, 2.0, 3.0],
                    targets: int = 1
                > (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = ValueVariablesGroup(
            {
                "feature1": table["feature1"],
                "feature2": table["feature2"],
            }
        )

        translator = LinearRegressorTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError,
            match="Coefficients length must equal targets number of input fields",
        ):
            translator.process()

    def test_linear_regressor_intercepts_length_mismatch(self):
        """Test LinearRegressorTranslator raises error when intercepts length mismatch."""
        table = ibis.memtable({"feature": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N,2] Y)
            {
                Y = ai.onnx.ml.LinearRegressor <
                    coefficients: floats = [1.0, 2.0],
                    intercepts: floats = [5.0],
                    targets: int = 2
                > (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = ValueVariablesGroup({"feature": table["feature"]})

        translator = LinearRegressorTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError,
            match="LinearRegressor: intercepts length must match targets or be empty",
        ):
            translator.process()

    def test_linear_regressor_post_transform_not_implemented(self):
        """Test LinearRegressorTranslator raises error when post_transform != 'NONE'."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] Y)
            {
                Y = ai.onnx.ml.LinearRegressor <
                    coefficients: floats = [1.0],
                    post_transform: string = "LOGISTIC"
                > (input)
            }
        """)

        variables = GraphVariables(table, model)

        translator = LinearRegressorTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            NotImplementedError, match="Post transform is not implemented"
        ):
            translator.process()

    def test_linear_regressor_single_input_multiple_targets(self):
        """Test LinearRegressorTranslator raises error with single input and multiple targets."""
        table = ibis.memtable({"input": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N,2] Y)
            {
                Y = ai.onnx.ml.LinearRegressor <
                    coefficients: floats = [1.0, 2.0],
                    intercepts: floats = [0.0, 0.0],
                    targets: int = 2
                > (input)
            }
        """)

        variables = GraphVariables(table, model)

        translator = LinearRegressorTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError,
            match="Single column input expects exactly one target and one coefficient",
        ):
            translator.process()




class TestTreeEnsembleClassifierTranslator:
    optimizer = Optimizer(enabled=False)

    def test_binary_classification_single_tree(self):
        """Test TreeEnsembleClassifier with binary classification and single tree."""
        from onnx import helper, TensorProto
        from orbital.translation.steps.trees.classifier import TreeEnsembleClassifierTranslator

        table = ibis.memtable({"X": [0.3, 0.7, 0.2]})

        # Create a simple binary classification tree:
        # if feature[0] <= 0.5:
        #     return class 0 (weight 1.0)
        # else:
        #     return class 1 (weight 1.0)
        node = helper.make_node(
            op_type="TreeEnsembleClassifier",
            inputs=["X"],
            outputs=["Y", "P"],
            domain="ai.onnx.ml",
            # Tree structure attributes
            nodes_treeids=[0, 0, 0],
            nodes_nodeids=[0, 1, 2],
            nodes_featureids=[0, 0, 0],
            nodes_modes=["BRANCH_LEQ", "LEAF", "LEAF"],
            nodes_values=[0.5, 0.0, 0.0],
            nodes_truenodeids=[1, 0, 0],
            nodes_falsenodeids=[2, 0, 0],
            nodes_missing_value_tracks_true=[0, 0, 0],
            # Class weights
            class_treeids=[0, 0],
            class_nodeids=[1, 2],
            class_ids=[0, 0],
            class_weights=[1.0, -1.0],
            # Class labels
            classlabels_int64s=[0, 1],
            post_transform="NONE",
        )

        # Create a minimal graph and model
        X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [None])
        Y = helper.make_tensor_value_info("Y", TensorProto.INT64, [None])
        P = helper.make_tensor_value_info("P", TensorProto.FLOAT, [None, 2])

        graph = helper.make_graph([node], "test", [X], [Y, P])
        model = helper.make_model(graph)

        variables = GraphVariables(table, model.graph)

        translator = TreeEnsembleClassifierTranslator(
            table, model.graph.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "Y" in variables
        assert "P" in variables

        label_result = variables.peek_variable("Y")
        prob_result = variables.peek_variable("P")

        backend = ibis.duckdb.connect()
        labels = list(backend.execute(label_result))
        assert labels == [1, 0, 1]

        # Check probabilities
        assert isinstance(prob_result, NumericVariablesGroup)
        assert "0" in prob_result
        assert "1" in prob_result

    def test_multiclass_classification_single_tree(self):
        """Test TreeEnsembleClassifier with multi-class classification."""
        from onnx import helper, TensorProto
        from orbital.translation.steps.trees.classifier import TreeEnsembleClassifierTranslator

        table = ibis.memtable({"X": [0.3, 0.7, 0.2]})

        # Create a multi-class tree with 3 classes
        node = helper.make_node(
            op_type="TreeEnsembleClassifier",
            inputs=["X"],
            outputs=["Y", "P"],
            domain="ai.onnx.ml",
            # Simple tree: always returns class weights at leaf 1
            nodes_treeids=[0, 0],
            nodes_nodeids=[0, 1],
            nodes_featureids=[0, 0],
            nodes_modes=["BRANCH_LEQ", "LEAF"],
            nodes_values=[10.0, 0.0],
            nodes_truenodeids=[1, 0],
            nodes_falsenodeids=[1, 0],
            nodes_missing_value_tracks_true=[0, 0],
            # Class weights for 3 classes at leaf node
            class_treeids=[0, 0, 0],
            class_nodeids=[1, 1, 1],
            class_ids=[0, 1, 2],
            class_weights=[0.2, 0.5, 0.3],
            # String class labels
            classlabels_strings=["cat", "dog", "bird"],
            post_transform="NONE",
        )

        X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [None])
        Y = helper.make_tensor_value_info("Y", TensorProto.STRING, [None])
        P = helper.make_tensor_value_info("P", TensorProto.FLOAT, [None, 3])

        graph = helper.make_graph([node], "test", [X], [Y, P])
        model = helper.make_model(graph)

        variables = GraphVariables(table, model.graph)

        translator = TreeEnsembleClassifierTranslator(
            table, model.graph.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "Y" in variables
        assert "P" in variables

        label_result = variables.peek_variable("Y")
        prob_result = variables.peek_variable("P")

        backend = ibis.duckdb.connect()
        labels = list(backend.execute(label_result))
        # All inputs should predict "dog" since it has highest weight (0.5)
        assert labels == ["dog", "dog", "dog"]

        # Check probabilities group structure
        assert isinstance(prob_result, NumericVariablesGroup)
        assert "cat" in prob_result
        assert "dog" in prob_result
        assert "bird" in prob_result

    def test_classifier_invalid_input_type(self):
        """Test TreeEnsembleClassifier raises error for invalid input type."""
        from onnx import helper, TensorProto
        from orbital.translation.steps.trees.classifier import TreeEnsembleClassifierTranslator

        table = ibis.memtable({"X": [0.3, 0.7, 0.2]})

        node = helper.make_node(
            op_type="TreeEnsembleClassifier",
            inputs=["X"],
            outputs=["Y", "P"],
            domain="ai.onnx.ml",
            nodes_treeids=[0],
            nodes_nodeids=[0],
            nodes_featureids=[0],
            nodes_modes=["LEAF"],
            nodes_values=[0.0],
            nodes_truenodeids=[0],
            nodes_falsenodeids=[0],
            nodes_missing_value_tracks_true=[0],
            class_treeids=[0],
            class_nodeids=[0],
            class_ids=[0],
            class_weights=[1.0],
            classlabels_int64s=[0, 1],
        )

        X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [None])
        Y = helper.make_tensor_value_info("Y", TensorProto.INT64, [None])
        P = helper.make_tensor_value_info("P", TensorProto.FLOAT, [None, 2])

        graph = helper.make_graph([node], "test", [X], [Y, P])
        model = helper.make_model(graph)

        variables = GraphVariables(table, model.graph)
        variables["X"] = "invalid_string_input"  # type: ignore[assignment]

        translator = TreeEnsembleClassifierTranslator(
            table, model.graph.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError,
            match="TreeEnsembleClassifier: The first operand must be a column or a column group"
        ):
            translator.process()


class TestTreeEnsembleRegressorTranslator:
    optimizer = Optimizer(enabled=False)

    def test_single_tree_regression(self):
        """Test TreeEnsembleRegressor with a single decision tree."""
        from onnx import helper, TensorProto
        from orbital.translation.steps.trees.regressor import TreeEnsembleRegressorTranslator

        table = ibis.memtable({"X": [0.3, 0.7, 0.2]})

        # Create a simple regression tree:
        # if feature[0] <= 0.5:
        #     return 10.0
        # else:
        #     return 20.0
        node = helper.make_node(
            op_type="TreeEnsembleRegressor",
            inputs=["X"],
            outputs=["Y"],
            domain="ai.onnx.ml",
            # Tree structure
            nodes_treeids=[0, 0, 0],
            nodes_nodeids=[0, 1, 2],
            nodes_featureids=[0, 0, 0],
            nodes_modes=["BRANCH_LEQ", "LEAF", "LEAF"],
            nodes_values=[0.5, 0.0, 0.0],
            nodes_truenodeids=[1, 0, 0],
            nodes_falsenodeids=[2, 0, 0],
            nodes_missing_value_tracks_true=[0, 0, 0],
            # Target weights for regression
            target_treeids=[0, 0],
            target_nodeids=[1, 2],
            target_weights=[10.0, 20.0],
            # Base value (offset)
            base_values=[5.0],
        )

        X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [None])
        Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [None])

        graph = helper.make_graph([node], "test", [X], [Y])
        model = helper.make_model(graph)

        variables = GraphVariables(table, model.graph)

        translator = TreeEnsembleRegressorTranslator(
            table, model.graph.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "Y" in variables
        result = variables.peek_variable("Y")

        backend = ibis.duckdb.connect()
        predictions = list(backend.execute(result))
        assert predictions == [15.0, 25.0, 15.0]

    def test_regression_base_values_applied(self):
        """Test TreeEnsembleRegressor correctly applies base_values."""
        from onnx import helper, TensorProto
        from orbital.translation.steps.trees.regressor import TreeEnsembleRegressorTranslator

        table = ibis.memtable({"X": [1.0, 2.0, 3.0]})

        # Simple tree that always returns same value
        node = helper.make_node(
            op_type="TreeEnsembleRegressor",
            inputs=["X"],
            outputs=["Y"],
            domain="ai.onnx.ml",
            # Single leaf node tree
            nodes_treeids=[0],
            nodes_nodeids=[0],
            nodes_featureids=[0],
            nodes_modes=["LEAF"],
            nodes_values=[0.0],
            nodes_truenodeids=[0],
            nodes_falsenodeids=[0],
            nodes_missing_value_tracks_true=[0],
            # Weight at leaf
            target_treeids=[0],
            target_nodeids=[0],
            target_weights=[7.0],
            # Base value should be added
            base_values=[3.0],
        )

        X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [None])
        Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [None])

        graph = helper.make_graph([node], "test", [X], [Y])
        model = helper.make_model(graph)

        variables = GraphVariables(table, model.graph)

        translator = TreeEnsembleRegressorTranslator(
            table, model.graph.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "Y" in variables
        result = variables.peek_variable("Y")

        backend = ibis.duckdb.connect()
        computed = backend.execute(result)
        # Single-leaf tree returns a constant expression (scalar)
        assert computed == 10.0

    def test_regressor_invalid_input_type(self):
        """Test TreeEnsembleRegressor raises error for invalid input type."""
        from onnx import helper, TensorProto
        from orbital.translation.steps.trees.regressor import TreeEnsembleRegressorTranslator

        table = ibis.memtable({"X": [1.0, 2.0, 3.0]})

        node = helper.make_node(
            op_type="TreeEnsembleRegressor",
            inputs=["X"],
            outputs=["Y"],
            domain="ai.onnx.ml",
            nodes_treeids=[0],
            nodes_nodeids=[0],
            nodes_featureids=[0],
            nodes_modes=["LEAF"],
            nodes_values=[0.0],
            nodes_truenodeids=[0],
            nodes_falsenodeids=[0],
            nodes_missing_value_tracks_true=[0],
            target_treeids=[0],
            target_nodeids=[0],
            target_weights=[1.0],
        )

        X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [None])
        Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [None])

        graph = helper.make_graph([node], "test", [X], [Y])
        model = helper.make_model(graph)

        variables = GraphVariables(table, model.graph)
        variables["X"] = "invalid_string_input"  # type: ignore[assignment]

        translator = TreeEnsembleRegressorTranslator(
            table, model.graph.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError,
            match="TreeEnsembleRegressor: The first operand must be a column or a column group"
        ):
            translator.process()




class TestScalerTranslator:
    optimizer = Optimizer(enabled=False)

    def test_scaler_single_column(self):
        """Test ScalerTranslator with single column scaling."""
        table = ibis.memtable({"input": [2.0, 4.0, 6.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = ai.onnx.ml.Scaler <offset: floats = [1.0], scale: floats = [2.0]> (input)
            }
        """)

        variables = GraphVariables(table, model)
        translator = ScalerTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        # Y = (X - offset) * scale = (X - 1) * 2
        # [2, 4, 6] -> [2, 6, 10]
        backend = ibis.duckdb.connect()
        assert list(backend.execute(result)) == [2.0, 6.0, 10.0]

    def test_scaler_group_columns(self):
        """Test ScalerTranslator with group of columns."""
        table = ibis.memtable(
            {
                "col_a": [2.0, 4.0, 6.0],
                "col_b": [10.0, 20.0, 30.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = ai.onnx.ml.Scaler <offset: floats = [1.0, 5.0], scale: floats = [2.0, 0.5]> (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = ScalerTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        # col_a: Y = (X - 1) * 2 = [2, 6, 10]
        assert list(backend.execute(result["col_a"])) == [2.0, 6.0, 10.0]
        # col_b: Y = (X - 5) * 0.5 = [2.5, 7.5, 12.5]
        assert list(backend.execute(result["col_b"])) == [2.5, 7.5, 12.5]

    def test_scaler_only_offset(self):
        """Test ScalerTranslator with only offset (scale=1.0)."""
        table = ibis.memtable({"input": [2.0, 4.0, 6.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = ai.onnx.ml.Scaler <offset: floats = [3.0], scale: floats = [1.0]> (input)
            }
        """)

        variables = GraphVariables(table, model)
        translator = ScalerTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        # Y = (X - 3) * 1 = X - 3
        backend = ibis.duckdb.connect()
        assert list(backend.execute(result)) == [-1.0, 1.0, 3.0]

    def test_scaler_only_scale(self):
        """Test ScalerTranslator with only scale (offset=0.0)."""
        table = ibis.memtable({"input": [2.0, 4.0, 6.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = ai.onnx.ml.Scaler <offset: floats = [0.0], scale: floats = [3.0]> (input)
            }
        """)

        variables = GraphVariables(table, model)
        translator = ScalerTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        # Y = (X - 0) * 3 = X * 3
        backend = ibis.duckdb.connect()
        assert list(backend.execute(result)) == [6.0, 12.0, 18.0]

    def test_scaler_mismatched_offset_scale_counts(self):
        """Test ScalerTranslator raises error when offset/scale counts don't match."""
        table = ibis.memtable(
            {
                "col_a": [2.0, 4.0, 6.0],
                "col_b": [10.0, 20.0, 30.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input) => (float[N] output) {
                output = ai.onnx.ml.Scaler <offset: floats = [1.0, 2.0, 3.0], scale: floats = [2.0, 3.0]> (input)
            }
        """)

        variables = GraphVariables(ibis.memtable({"input": [1.0]}), model)
        variables["input"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = ScalerTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="offset and scale lists must match"):
            translator.process()


class TestOneHotEncoderTranslator:
    optimizer = Optimizer(enabled=False)

    def test_onehot_string_column(self):
        """Test OneHotEncoderTranslator with string column."""
        table = ibis.memtable({"category": ["cat", "dog", "cat", "bird"]})
        model = onnx.parser.parse_graph("""
            agraph (string[N] category) => (float[N, 3] output) {
                output = ai.onnx.ml.OneHotEncoder <cats_strings: strings = ["cat", "dog", "bird"]> (category)
            }
        """)

        variables = GraphVariables(table, model)
        translator = OneHotEncoderTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        # First row is "cat", so cat=1.0, dog=0.0, bird=0.0
        assert list(backend.execute(result["cat"])) == [1.0, 0.0, 1.0, 0.0]
        assert list(backend.execute(result["dog"])) == [0.0, 1.0, 0.0, 0.0]
        assert list(backend.execute(result["bird"])) == [0.0, 0.0, 0.0, 1.0]

    def test_onehot_output_type(self):
        """Test OneHotEncoderTranslator output is ValueVariablesGroup with correct keys."""
        table = ibis.memtable({"category": ["apple", "banana", "apple"]})
        model = onnx.parser.parse_graph("""
            agraph (string[N] category) => (float[N, 2] output) {
                output = ai.onnx.ml.OneHotEncoder <cats_strings: strings = ["apple", "banana"]> (category)
            }
        """)

        variables = GraphVariables(table, model)
        translator = OneHotEncoderTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)
        assert "apple" in result
        assert "banana" in result
        assert len(result) == 2

    def test_onehot_values(self):
        """Test OneHotEncoderTranslator outputs 0.0 or 1.0 floats."""
        table = ibis.memtable({"category": ["red", "blue", "green", "red"]})
        model = onnx.parser.parse_graph("""
            agraph (string[N] category) => (float[N, 3] output) {
                output = ai.onnx.ml.OneHotEncoder <cats_strings: strings = ["red", "green", "blue"]> (category)
            }
        """)

        variables = GraphVariables(table, model)
        translator = OneHotEncoderTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        # Input: ["red", "blue", "green", "red"]
        # Verify one-hot encoding produces correct 0.0/1.0 values
        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["red"])) == [1.0, 0.0, 0.0, 1.0]
        assert list(backend.execute(result["green"])) == [0.0, 0.0, 1.0, 0.0]
        assert list(backend.execute(result["blue"])) == [0.0, 1.0, 0.0, 0.0]

    def test_onehot_missing_cats_strings(self):
        """Test OneHotEncoderTranslator raises error when cats_strings is missing."""
        table = ibis.memtable({"category": ["cat", "dog"]})
        model = onnx.parser.parse_graph("""
            agraph (string[N] category) => (float[N, 3] output) {
                output = ai.onnx.ml.OneHotEncoder (category)
            }
        """)

        variables = GraphVariables(table, model)
        translator = OneHotEncoderTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="attribute cats_strings not found"):
            translator.process()


class TestLabelEncoderTranslator:
    optimizer = Optimizer(enabled=False)

    def test_labelencoder_string_to_int(self):
        """Test LabelEncoderTranslator encoding string labels to integers."""
        table = ibis.memtable({"label": ["cat", "dog", "cat", "bird", "dog"]})
        model = onnx.parser.parse_graph("""
            agraph (string[N] label) => (int64[N] output) {
                output = ai.onnx.ml.LabelEncoder <keys_strings: strings = ["cat", "dog", "bird"], values_int64s: ints = [0, 1, 2]> (label)
            }
        """)

        variables = GraphVariables(table, model)
        translator = LabelEncoderTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result)) == [0, 1, 0, 2, 1]

    def test_labelencoder_int_to_int(self):
        """Test LabelEncoderTranslator encoding integer labels to different integers."""
        table = ibis.memtable({"label": [10, 20, 10, 30, 20]})
        model = onnx.parser.parse_graph("""
            agraph (int64[N] label) => (int64[N] output) {
                output = ai.onnx.ml.LabelEncoder <keys_int64s: ints = [10, 20, 30], values_int64s: ints = [100, 200, 300]> (label)
            }
        """)

        variables = GraphVariables(table, model)
        translator = LabelEncoderTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result)) == [100, 200, 100, 300, 200]

    def test_labelencoder_default_value(self):
        """Test LabelEncoderTranslator handles default value for unknown labels."""
        table = ibis.memtable({"label": ["cat", "dog", "unknown", "cat"]})
        model = onnx.parser.parse_graph("""
            agraph (string[N] label) => (int64[N] output) {
                output = ai.onnx.ml.LabelEncoder <keys_strings: strings = ["cat", "dog"], values_int64s: ints = [0, 1], default_int64: int = 999> (label)
            }
        """)

        variables = GraphVariables(table, model)
        translator = LabelEncoderTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result)) == [0, 1, 999, 0]

    def test_labelencoder_missing_keys(self):
        """Test LabelEncoderTranslator raises error when keys attribute is missing."""
        table = ibis.memtable({"label": ["cat", "dog"]})
        model = onnx.parser.parse_graph("""
            agraph (string[N] label) => (int64[N] output) {
                output = ai.onnx.ml.LabelEncoder <values_int64s: ints = [0, 1]> (label)
            }
        """)

        variables = GraphVariables(table, model)
        translator = LabelEncoderTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="required mapping attributes not found"):
            translator.process()

class TestWhereTranslator:
    optimizer = Optimizer(enabled=False)

    def test_where_single_columns(self):
        """Test WhereTranslator selecting between two single columns based on condition."""
        table = ibis.memtable(
            {
                "condition": [True, False, True],
                "true_val": [10.0, 20.0, 30.0],
                "false_val": [100.0, 200.0, 300.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (bool[N] condition, float[N] true_val, float[N] false_val) => (float[N] output) {
                output = Where(condition, true_val, false_val)
            }
        """)

        variables = GraphVariables(table, model)
        translator = WhereTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        computed = list(backend.execute(result))
        # When condition is True, take true_val; when False, take false_val
        assert computed == [10.0, 200.0, 30.0]

    def test_where_group_columns(self):
        """Test WhereTranslator selecting between column groups based on condition."""
        table = ibis.memtable(
            {
                "condition": [True, False, True],
                "true_col1": [1.0, 2.0, 3.0],
                "true_col2": [4.0, 5.0, 6.0],
                "false_col1": [10.0, 20.0, 30.0],
                "false_col2": [40.0, 50.0, 60.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (bool[N] condition, float[N] true_expr, float[N] false_expr) => (float[N] output) {
                output = Where(condition, true_expr, false_expr)
            }
        """)

        variables = GraphVariables(
            ibis.memtable(
                {"condition": [True], "true_expr": [1.0], "false_expr": [1.0]}
            ),
            model,
        )
        variables["condition"] = table["condition"]
        variables["true_expr"] = ValueVariablesGroup(
            {"col1": table["true_col1"], "col2": table["true_col2"]}
        )
        variables["false_expr"] = ValueVariablesGroup(
            {"col1": table["false_col1"], "col2": table["false_col2"]}
        )

        translator = WhereTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["c0"])) == [1.0, 20.0, 3.0]
        assert list(backend.execute(result["c1"])) == [4.0, 50.0, 6.0]

    def test_where_broadcast_scalar_true(self):
        """Test WhereTranslator with single true value broadcast to group false."""
        table = ibis.memtable(
            {
                "condition": [True, False, True],
                "true_val": [42.0, 42.0, 42.0],
                "false_col1": [10.0, 20.0, 30.0],
                "false_col2": [100.0, 200.0, 300.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (bool[N] condition, float[N] true_expr, float[N] false_expr) => (float[N] output) {
                output = Where(condition, true_expr, false_expr)
            }
        """)

        variables = GraphVariables(
            ibis.memtable(
                {"condition": [True], "true_expr": [1.0], "false_expr": [1.0]}
            ),
            model,
        )
        variables["condition"] = table["condition"]
        variables["true_expr"] = table["true_val"]
        variables["false_expr"] = ValueVariablesGroup(
            {"col1": table["false_col1"], "col2": table["false_col2"]}
        )

        translator = WhereTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["c0"])) == [42.0, 20.0, 42.0]
        assert list(backend.execute(result["c1"])) == [42.0, 200.0, 42.0]

    def test_where_broadcast_scalar_false(self):
        """Test WhereTranslator with single false value broadcast to group true."""
        table = ibis.memtable(
            {
                "condition": [True, False, True],
                "true_col1": [10.0, 20.0, 30.0],
                "true_col2": [100.0, 200.0, 300.0],
                "false_val": [99.0, 99.0, 99.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (bool[N] condition, float[N] true_expr, float[N] false_expr) => (float[N] output) {
                output = Where(condition, true_expr, false_expr)
            }
        """)

        variables = GraphVariables(
            ibis.memtable(
                {"condition": [True], "true_expr": [1.0], "false_expr": [1.0]}
            ),
            model,
        )
        variables["condition"] = table["condition"]
        variables["true_expr"] = ValueVariablesGroup(
            {"col1": table["true_col1"], "col2": table["true_col2"]}
        )
        variables["false_expr"] = table["false_val"]

        translator = WhereTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["c0"])) == [10.0, 99.0, 30.0]
        assert list(backend.execute(result["c1"])) == [100.0, 99.0, 300.0]

    def test_where_condition_group_error(self):
        """Test WhereTranslator raises error when condition is a group of columns."""
        table = ibis.memtable(
            {
                "cond1": [True, False],
                "cond2": [False, True],
                "true_val": [1.0, 2.0],
                "false_val": [10.0, 20.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (bool[N] condition, float[N] true_expr, float[N] false_expr) => (float[N] output) {
                output = Where(condition, true_expr, false_expr)
            }
        """)

        variables = GraphVariables(
            ibis.memtable(
                {"condition": [True], "true_expr": [1.0], "false_expr": [1.0]}
            ),
            model,
        )
        variables["condition"] = ValueVariablesGroup(
            {"cond1": table["cond1"], "cond2": table["cond2"]}
        )
        variables["true_expr"] = table["true_val"]
        variables["false_expr"] = table["false_val"]

        translator = WhereTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            NotImplementedError,
            match="Where: The condition expression can't be a group of columns",
        ):
            translator.process()

    def test_where_mismatched_group_sizes(self):
        """Test WhereTranslator raises error when true and false groups have different sizes."""
        table = ibis.memtable(
            {
                "condition": [True, False],
                "true_col1": [1.0, 2.0],
                "true_col2": [3.0, 4.0],
                "false_col1": [10.0, 20.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (bool[N] condition, float[N] true_expr, float[N] false_expr) => (float[N] output) {
                output = Where(condition, true_expr, false_expr)
            }
        """)

        variables = GraphVariables(
            ibis.memtable(
                {"condition": [True], "true_expr": [1.0], "false_expr": [1.0]}
            ),
            model,
        )
        variables["condition"] = table["condition"]
        variables["true_expr"] = ValueVariablesGroup(
            {"col1": table["true_col1"], "col2": table["true_col2"]}
        )
        variables["false_expr"] = ValueVariablesGroup({"col1": table["false_col1"]})

        translator = WhereTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError,
            match="Where: The number of values in the true and false expressions must match",
        ):
            translator.process()


class TestZipMapTranslator:
    optimizer = Optimizer(enabled=False)

    def test_zipmap_string_labels_group(self):
        """Test ZipMapTranslator with string labels and group of columns."""
        table = ibis.memtable(
            {
                "prob_class_0": [0.7, 0.2, 0.4],
                "prob_class_1": [0.2, 0.5, 0.1],
                "prob_class_2": [0.1, 0.3, 0.5],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (float[N] output) {
                output = ai.onnx.ml.ZipMap <classlabels_strings: strings = ["negative", "neutral", "positive"]> (data)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = ValueVariablesGroup(
            {
                "class_0": table["prob_class_0"],
                "class_1": table["prob_class_1"],
                "class_2": table["prob_class_2"],
            }
        )

        translator = ZipMapTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        # Check that labels were properly mapped
        assert "negative" in result
        assert "neutral" in result
        assert "positive" in result

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["negative"])) == [0.7, 0.2, 0.4]
        assert list(backend.execute(result["neutral"])) == [0.2, 0.5, 0.1]
        assert list(backend.execute(result["positive"])) == [0.1, 0.3, 0.5]

    def test_zipmap_int64_labels_group(self):
        """Test ZipMapTranslator with int64 labels and group of columns."""
        table = ibis.memtable(
            {
                "feat_0": [1.0, 2.0, 3.0],
                "feat_1": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (float[N] output) {
                output = ai.onnx.ml.ZipMap <classlabels_int64s: ints = [100, 200]> (data)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = ValueVariablesGroup(
            {
                "feat_0": table["feat_0"],
                "feat_1": table["feat_1"],
            }
        )

        translator = ZipMapTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        # int64 labels are converted to strings
        assert "100" in result
        assert "200" in result

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["100"])) == [1.0, 2.0, 3.0]
        assert list(backend.execute(result["200"])) == [4.0, 5.0, 6.0]

    def test_zipmap_single_column(self):
        """Test ZipMapTranslator with single column and single label."""
        table = ibis.memtable({"prob": [0.95, 0.87, 0.92]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (float[N] output) {
                output = ai.onnx.ml.ZipMap <classlabels_strings: strings = ["confidence"]> (data)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = table["prob"]

        translator = ZipMapTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)

        assert "confidence" in result
        assert len(result) == 1

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["confidence"])) == [0.95, 0.87, 0.92]

    def test_zipmap_missing_labels_error(self):
        """Test ZipMapTranslator raises error when no classlabels attribute is found."""
        table = ibis.memtable({"data": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (float[N] output) {
                output = ai.onnx.ml.ZipMap (data)
            }
        """)

        variables = GraphVariables(table, model)

        translator = ZipMapTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError, match="ZipMap: required mapping attributes not found"
        ):
            translator.process()

    def test_zipmap_mismatched_label_count(self):
        """Test ZipMapTranslator raises error when labels count doesn't match columns."""
        table = ibis.memtable(
            {
                "col1": [1.0, 2.0, 3.0],
                "col2": [4.0, 5.0, 6.0],
                "col3": [7.0, 8.0, 9.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (float[N] output) {
                output = ai.onnx.ml.ZipMap <classlabels_strings: strings = ["label1", "label2"]> (data)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = ValueVariablesGroup(
            {
                "col1": table["col1"],
                "col2": table["col2"],
                "col3": table["col3"],
            }
        )

        translator = ZipMapTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError, match="ZipMap: The number of labels and columns must match"
        ):
            translator.process()

class TestConcatTranslator:
    optimizer = Optimizer(enabled=False)

    def test_concat_two_single_columns(self):
        """Test concatenating two single columns."""
        table = ibis.memtable(
            {
                "feature1": [1.0, 2.0, 3.0],
                "feature2": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] feature1, float[N] feature2) => (float[N] output) {
                output = Concat <axis: int = 1> (feature1, feature2)
            }
        """)

        variables = GraphVariables(table, model)
        translator = ConcatTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)
        assert len(result) == 2
        assert "feature1" in result
        assert "feature2" in result

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["feature1"])) == [1.0, 2.0, 3.0]
        assert list(backend.execute(result["feature2"])) == [4.0, 5.0, 6.0]

    def test_concat_column_groups(self):
        """Test concatenating column groups."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
                "col_c": [7.0, 8.0, 9.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] group1, float[N] group2) => (float[N] output) {
                output = Concat <axis: int = -1> (group1, group2)
            }
        """)

        variables = GraphVariables(
            ibis.memtable({"group1": [1.0], "group2": [1.0]}), model
        )
        variables["group1"] = ValueVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )
        variables["group2"] = ValueVariablesGroup(
            {
                "col_c": table["col_c"],
            }
        )

        translator = ConcatTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)
        assert len(result) == 3
        assert "group1.col_a" in result
        assert "group1.col_b" in result
        assert "group2.col_c" in result

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["group1.col_a"])) == [1.0, 2.0, 3.0]
        assert list(backend.execute(result["group1.col_b"])) == [4.0, 5.0, 6.0]
        assert list(backend.execute(result["group2.col_c"])) == [7.0, 8.0, 9.0]

    def test_concat_mix_single_columns_and_groups(self):
        """Test concatenating mix of single columns and groups."""
        table = ibis.memtable(
            {
                "single_col": [1.0, 2.0, 3.0],
                "group_col_a": [4.0, 5.0, 6.0],
                "group_col_b": [7.0, 8.0, 9.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] single, float[N] group) => (float[N] output) {
                output = Concat <axis: int = 1> (single, group)
            }
        """)

        variables = GraphVariables(
            ibis.memtable({"single": [1.0], "group": [1.0]}), model
        )
        variables["single"] = table["single_col"]
        variables["group"] = ValueVariablesGroup(
            {
                "col_a": table["group_col_a"],
                "col_b": table["group_col_b"],
            }
        )

        translator = ConcatTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)
        assert len(result) == 3
        assert "single" in result
        assert "group.col_a" in result
        assert "group.col_b" in result

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["single"])) == [1.0, 2.0, 3.0]
        assert list(backend.execute(result["group.col_a"])) == [4.0, 5.0, 6.0]
        assert list(backend.execute(result["group.col_b"])) == [7.0, 8.0, 9.0]

    def test_concat_unsupported_axis(self):
        """Test error for unsupported axis."""
        table = ibis.memtable(
            {"feature1": [1.0, 2.0, 3.0], "feature2": [4.0, 5.0, 6.0]}
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] feature1, float[N] feature2) => (float[N] output) {
                output = Concat <axis: int = 0> (feature1, feature2)
            }
        """)

        variables = GraphVariables(table, model)
        translator = ConcatTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            NotImplementedError, match="only supports concatenating over columns"
        ):
            translator.process()


class TestFeatureVectorizerTranslator:
    optimizer = Optimizer(enabled=False)

    def test_featurevectorizer_multiple_inputs_with_dimensions(self):
        """Test vectorizing multiple inputs with inputdimensions attribute."""
        table = ibis.memtable(
            {
                "feature1": [1.0, 2.0, 3.0],
                "group_col_a": [4.0, 5.0, 6.0],
                "group_col_b": [7.0, 8.0, 9.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input1, float[N] input2) => (float[N] output) {
                output = ai.onnx.ml.FeatureVectorizer <inputdimensions: ints = [1, 2]> (input1, input2)
            }
        """)

        variables = GraphVariables(
            ibis.memtable({"input1": [1.0], "input2": [1.0]}), model
        )
        variables["input1"] = table["feature1"]
        variables["input2"] = ValueVariablesGroup(
            {
                "col_a": table["group_col_a"],
                "col_b": table["group_col_b"],
            }
        )

        translator = FeatureVectorizerTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)
        assert len(result) == 3
        assert "input1" in result
        assert "input2.col_a" in result
        assert "input2.col_b" in result

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["input1"])) == [1.0, 2.0, 3.0]
        assert list(backend.execute(result["input2.col_a"])) == [4.0, 5.0, 6.0]
        assert list(backend.execute(result["input2.col_b"])) == [7.0, 8.0, 9.0]

    def test_featurevectorizer_single_input_feature(self):
        """Test single input feature."""
        table = ibis.memtable({"input1": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input1) => (float[N] output) {
                output = ai.onnx.ml.FeatureVectorizer <inputdimensions: ints = [1]> (input1)
            }
        """)

        variables = GraphVariables(table, model)
        translator = FeatureVectorizerTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)
        assert len(result) == 1
        assert "input1" in result

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["input1"])) == [1.0, 2.0, 3.0]

    def test_featurevectorizer_mismatched_dimensions(self):
        """Test error when input dimensions don't match actual columns."""
        table = ibis.memtable(
            {
                "feature1": [1.0, 2.0, 3.0],
                "group_col_a": [4.0, 5.0, 6.0],
                "group_col_b": [7.0, 8.0, 9.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] input1, float[N] input2) => (float[N] output) {
                output = ai.onnx.ml.FeatureVectorizer <inputdimensions: ints = [1, 3]> (input1, input2)
            }
        """)

        variables = GraphVariables(
            ibis.memtable({"input1": [1.0], "input2": [1.0]}), model
        )
        variables["input1"] = table["feature1"]
        variables["input2"] = ValueVariablesGroup(
            {
                "col_a": table["group_col_a"],
                "col_b": table["group_col_b"],
            }
        )

        translator = FeatureVectorizerTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(ValueError, match="Number of columns in input input2"):
            translator.process()

    def test_featurevectorizer_wrong_single_column_dimension(self):
        """Test error when single column has wrong dimension."""
        table = ibis.memtable({"input1": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input1) => (float[N] output) {
                output = ai.onnx.ml.FeatureVectorizer <inputdimensions: ints = [2]> (input1)
            }
        """)

        variables = GraphVariables(table, model)
        translator = FeatureVectorizerTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError,
            match="When merging over individual columns, the dimension should be 1",
        ):
            translator.process()

    def test_featurevectorizer_mismatched_input_count(self):
        """Test error when input count doesn't match dimensions count."""
        table = ibis.memtable({"input1": [1.0, 2.0, 3.0], "input2": [4.0, 5.0, 6.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] input1, float[N] input2) => (float[N] output) {
                output = ai.onnx.ml.FeatureVectorizer <inputdimensions: ints = [1]> (input1, input2)
            }
        """)

        variables = GraphVariables(table, model)
        translator = FeatureVectorizerTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError,
            match="Number of input dimensions should be equal to number of inputs",
        ):
            translator.process()


class TestGatherTranslator:
    optimizer = Optimizer(enabled=False)

    def test_gather_single_element_from_group(self):
        """Test extracting single element from group by index."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
                "col_c": [7.0, 8.0, 9.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (float[N] output)
            <int32[1] index = {1}>
            {
                output = Gather <axis: int = 1> (data, index)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
                "col_c": table["col_c"],
            }
        )

        translator = GatherTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result)) == [4.0, 5.0, 6.0]

    def test_gather_element_from_single_column(self):
        """Test extracting element from single column (passthrough)."""
        table = ibis.memtable({"data": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (float[N] output)
            <int32[1] index = {0}>
            {
                output = Gather <axis: int = 1> (data, index)
            }
        """)

        variables = GraphVariables(table, model)

        translator = GatherTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result)) == [1.0, 2.0, 3.0]

    def test_gather_unsupported_axis(self):
        """Test error for unsupported axis."""
        table = ibis.memtable({"data": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (float[N] output)
            <int32[1] index = {0}>
            {
                output = Gather <axis: int = 0> (data, index)
            }
        """)

        variables = GraphVariables(table, model)

        translator = GatherTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            NotImplementedError, match="only selecting columns.*is supported"
        ):
            translator.process()

    def test_gather_index_out_of_bounds(self):
        """Test error for index out of bounds."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (float[N] output)
            <int32[1] index = {5}>
            {
                output = Gather <axis: int = 1> (data, index)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = NumericVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = GatherTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(IndexError, match="index out of bounds"):
            translator.process()

    def test_gather_invalid_index_for_single_column(self):
        """Test error for non-zero index on single column."""
        table = ibis.memtable({"data": [1.0, 2.0, 3.0]})
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (float[N] output)
            <int32[1] index = {1}>
            {
                output = Gather <axis: int = 1> (data, index)
            }
        """)

        variables = GraphVariables(table, model)

        translator = GatherTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            NotImplementedError, match="index 1 not supported for single columns"
        ):
            translator.process()


class TestArrayFeatureExtractorTranslator:
    optimizer = Optimizer(enabled=False)

    def test_arrayfeatureextractor_single_column_from_group(self):
        """Test extracting single column from group."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
                "col_c": [7.0, 8.0, 9.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (float[N] output)
            <int32[1] indices = {1}>
            {
                output = ai.onnx.ml.ArrayFeatureExtractor (data, indices)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = ValueVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
                "col_c": table["col_c"],
            }
        )

        translator = ArrayFeatureExtractorTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        # When extracting with a list of indices (even single), returns ValueVariablesGroup
        assert isinstance(result, ValueVariablesGroup)
        assert len(result) == 1
        assert "col_b" in result

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["col_b"])) == [4.0, 5.0, 6.0]

    def test_arrayfeatureextractor_multiple_columns_from_group(self):
        """Test extracting multiple columns from group."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
                "col_c": [7.0, 8.0, 9.0],
                "col_d": [10.0, 11.0, 12.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (float[N] output)
            <int32[2] indices = {0, 2}>
            {
                output = ai.onnx.ml.ArrayFeatureExtractor (data, indices)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = ValueVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
                "col_c": table["col_c"],
                "col_d": table["col_d"],
            }
        )

        translator = ArrayFeatureExtractorTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")
        assert isinstance(result, ValueVariablesGroup)
        assert len(result) == 2
        assert "col_a" in result
        assert "col_c" in result

        backend = ibis.duckdb.connect()
        assert list(backend.execute(result["col_a"])) == [1.0, 2.0, 3.0]
        assert list(backend.execute(result["col_c"])) == [7.0, 8.0, 9.0]

    def test_arrayfeatureextractor_index_out_of_bounds(self):
        """Test error for index out of bounds."""
        table = ibis.memtable(
            {
                "col_a": [1.0, 2.0, 3.0],
                "col_b": [4.0, 5.0, 6.0],
            }
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] data) => (float[N] output)
            <int32[3] indices = {0, 1, 2}>
            {
                output = ai.onnx.ml.ArrayFeatureExtractor (data, indices)
            }
        """)

        variables = GraphVariables(ibis.memtable({"data": [1.0]}), model)
        variables["data"] = ValueVariablesGroup(
            {
                "col_a": table["col_a"],
                "col_b": table["col_b"],
            }
        )

        translator = ArrayFeatureExtractorTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )

        with pytest.raises(
            ValueError,
            match="Indices requested are more than the available numer of columns",
        ):
            translator.process()

    def test_arrayfeatureextractor_from_list_of_constants(self):
        """Test extracting from a list of constants using column indices."""
        table = ibis.memtable(
            {"indices": [0, 1, 2, 1, 0], "dummy": [1.0, 2.0, 3.0, 4.0, 5.0]}
        )
        model = onnx.parser.parse_graph("""
            agraph (float[N] dummy, int32[N] indices) => (float[N] output) {
                output = ai.onnx.ml.ArrayFeatureExtractor (dummy, indices)
            }
        """)

        variables = GraphVariables(table, model)
        # Override the "dummy" variable with a list of constants (like class labels)
        variables["dummy"] = ["class_a", "class_b", "class_c"]  # type: ignore[assignment]

        translator = ArrayFeatureExtractorTranslator(
            table, model.node[0], variables, self.optimizer, TranslationOptions()
        )
        translator.process()

        assert "output" in variables
        result = variables.peek_variable("output")

        backend = ibis.duckdb.connect()
        # Should map indices to class names
        computed = list(backend.execute(result))
        assert computed == ["class_a", "class_b", "class_c", "class_b", "class_a"]
