# --------------------------------------------------------------------------------------
# Part of the interTwin Project: https://www.intertwin.eu/
#
# Created by: Matteo Bunino
#
# Credit:
# - Matteo Bunino <matteo.bunino@cern.ch> - CERN
# --------------------------------------------------------------------------------------

"""Here we show how to implement component interfaces in a simple way.

New components are defined as classes inheriting from the BaseComponent class
or an already defined component.
"""

from typing import Any, List, Optional, Tuple

from itwinai.components import DataGetter, DataSplitter, Saver, Trainer, monitor_exec


class MyDataGetter(DataGetter):
    def __init__(self, data_size: int, name: Optional[str] = None) -> None:
        super().__init__(name)
        self.data_size = data_size
        self.save_parameters(data_size=data_size)

    @monitor_exec
    def execute(self) -> List[int]:
        """Return a list dataset.

        Returns:
            List[int]: dataset
        """
        return list(range(self.data_size))


class MyDatasetSplitter(DataSplitter):
    @monitor_exec
    def execute(self, dataset: List[int]) -> Tuple[List[int], List[int], List[int]]:
        """Splits a list dataset into train, validation and test datasets.

        Args:
            dataset (List[int]): input list dataset.

        Returns:
            Tuple[List[int], List[int], List[int]]: train, validation, and
            test datasets respectively.
        """
        train_n = int(len(dataset) * self.train_proportion)
        valid_n = int(len(dataset) * self.validation_proportion)
        train_set = dataset[:train_n]
        vaild_set = dataset[train_n : train_n + valid_n]
        test_set = dataset[train_n + valid_n :]
        return train_set, vaild_set, test_set


class MyTrainer(Trainer):
    def __init__(self, lr: float = 1e-3, name: Optional[str] = None) -> None:
        super().__init__(name)
        self.save_parameters(name=name, lr=lr)

    @monitor_exec
    def execute(
        self, train_set: List[int], vaild_set: List[int], test_set: List[int]
    ) -> Tuple[List[int], List[int], List[int], str]:
        """Dummy ML trainer mocking a ML training algorithm.

        Args:
            train_set (List[int]): training dataset.
            vaild_set (List[int]): validation dataset.
            test_set (List[int]): test dataset.

        Returns:
            Tuple[List[int], List[int], List[int], str]: train, validation,
            test datasets, and trained model.
        """
        return train_set, vaild_set, test_set, "my_trained_model"


class MySaver(Saver):
    @monitor_exec
    def execute(self, artifact: Any) -> Any:
        """Saves an artifact to disk.

        Args:
            artifact (Any): artifact to save (e.g., dataset, model).

        Returns:
            Any: input artifact.
        """
        return artifact
