import pytest

from itwinai.components import Trainer, Adapter
from itwinai.pipeline import Pipeline
from itwinai.tests import (
    FakeGetterExec, FakeSplitterExec, FakeTrainerExec, FakeSaverExec
)
from itwinai.serialization import SerializationError


def test_serializable():
    """Test serialization of components."""
    comp = FakeGetterExec(data_uri='http://...')
    dict_serializ = comp.to_dict()
    assert isinstance(dict_serializ, dict)
    assert comp.name == "FakeGetterExec"
    assert dict_serializ == dict(
        class_path="itwinai.tests.dummy_components.FakeGetterExec",
        init_args=dict(data_uri='http://...', name=None)
    )

    # List
    comp = FakeGetterExec(data_uri=[1, 2, 3])
    dict_serializ = comp.to_dict()
    assert isinstance(dict_serializ, dict)
    assert comp.name == "FakeGetterExec"
    assert dict_serializ == dict(
        class_path="itwinai.tests.dummy_components.FakeGetterExec",
        init_args=dict(data_uri=[1, 2, 3], name=None)
    )

    # Tuple
    comp = FakeGetterExec(data_uri=(1, 2, 3))
    dict_serializ = comp.to_dict()
    assert isinstance(dict_serializ, dict)
    assert comp.name == "FakeGetterExec"
    assert dict_serializ == dict(
        class_path="itwinai.tests.dummy_components.FakeGetterExec",
        init_args=dict(data_uri=[1, 2, 3], name=None)
    )

    # Set
    comp = FakeGetterExec(data_uri={1, 2, 3})
    dict_serializ = comp.to_dict()
    assert isinstance(dict_serializ, dict)
    assert comp.name == "FakeGetterExec"
    assert dict_serializ == dict(
        class_path="itwinai.tests.dummy_components.FakeGetterExec",
        init_args=dict(data_uri=[1, 2, 3], name=None)
    )

    # Dict
    comp = FakeGetterExec(data_uri=dict(foo=12, bar="123", hl=3.14))
    dict_serializ = comp.to_dict()
    assert isinstance(dict_serializ, dict)
    assert comp.name == "FakeGetterExec"
    assert dict_serializ == dict(
        class_path="itwinai.tests.dummy_components.FakeGetterExec",
        init_args=dict(data_uri=dict(foo=12, bar="123", hl=3.14), name=None)
    )

    # Non serializable obj
    class NonSerializable:
        ...

    comp = FakeGetterExec(data_uri=NonSerializable())
    with pytest.raises(SerializationError) as exc_info:
        dict_serializ = comp.to_dict()
    assert ("is not a Python built-in type and it does not extend"
            in str(exc_info.value))

    # Local component class
    class MyTrainer(Trainer):
        def execute(self):
            ...

    comp = MyTrainer()
    with pytest.raises(SerializationError) as exc_info:
        dict_serializ = comp.to_dict()
    assert ("is defined locally, which is not supported for serialization."
            in str(exc_info.value))


def test_adapter():
    """Test Adapter component."""
    prefix = Adapter.INPUT_PREFIX
    adapter = Adapter(
        policy=[f"{prefix}{3-i}" for i in range(4)]
    )
    result = adapter.execute(0, 1, 2, 3)
    assert result == (3, 2, 1, 0)

    result = adapter.execute(*tuple(range(100)))
    assert result == (3, 2, 1, 0)

    adapter = Adapter(
        policy=[f"{prefix}0" for i in range(4)]
    )
    result = adapter.execute(0, 1, 2, 3)
    assert result == (0, 0, 0, 0)

    adapter = Adapter(
        policy=[f"{prefix}{i % 2}" for i in range(4)]
    )
    result = adapter.execute(0, 1, 2, 3)
    assert result == (0, 1, 0, 1)

    adapter = Adapter(
        policy=[f"{prefix}2", "hello", "world", 3.14]
    )
    result = adapter.execute(0, 1, 2, 3)
    assert result == (2, "hello", "world", 3.14)

    adapter = Adapter(
        policy=[1, 3, 5, 7, 11]
    )
    result = adapter.execute(0, 1, 2, 3)
    assert result == (1, 3, 5, 7, 11)

    adapter = Adapter(
        policy=[f"{prefix}{9-i}" for i in range(10)]
    )
    with pytest.raises(IndexError) as exc_info:
        result = adapter.execute(0, 1)
    assert str(exc_info.value) == (
        "The args received as input by 'Adapter' are not consistent with "
        "the given adapter policy because input args are too few! Input "
        "args are 2 but the policy foresees at least 10 items."
    )

    adapter = Adapter(
        policy=[]
    )
    result = adapter.execute(*tuple(range(100)))
    assert result == ()


@pytest.mark.integration
def test_adapter_integration_pipeline():
    """Test integration of Adapter component in the pipeline,
    connecting other components.
    """
    pipeline = Pipeline([
        FakeGetterExec(data_uri='http://...'),
        FakeSplitterExec(train_prop=.7),
        FakeTrainerExec(lr=1e-3, batch_size=32),
        Adapter(policy=[f"{Adapter.INPUT_PREFIX}-1"]),
        FakeSaverExec(save_path="my_model.pth")
    ])
    saved_model = pipeline.execute()
    assert saved_model == FakeTrainerExec.model
