"""All abstract classes later pipeline elements inherit from."""

import abc
import copy
import inspect
from collections.abc import Iterable
from typing import Any, NamedTuple, Self
from uuid import uuid4

from joblib import Parallel, delayed
from loguru import logger
from rdkit import Chem
from rdkit.Chem import Mol as RDKitMol
from sklearn.utils import Tags, TargetTags
from typing_extensions import override

from molpipeline.utils.multi_proc import check_available_cores


class InvalidInstance(NamedTuple):
    """Object which is returned when an instance cannot be processed.

    Attributes
    ----------
    element_id: str
        Id of the element which could not be processed.
    message: str
        Message why the element could not be processed.
    element_name: str | None
        Optional name of the element which could not be processed.
        The name of the pipeline element is often more descriptive than the id.

    """

    element_id: str
    message: str
    element_name: str | None = None

    def __repr__(self) -> str:
        """Return string representation of InvalidInstance.

        Returns
        -------
        str
            String representation of InvalidInstance.

        """
        return (
            f"InvalidInstance({self.element_name or self.element_id}, {self.message})"
        )


OptionalMol = RDKitMol | InvalidInstance


class RemovedInstance:  # pylint: disable=too-few-public-methods
    """Object which is returned by a ErrorFilter if an Invalid instance was removed."""

    def __init__(self, filter_element_id: str, message: str | None = None) -> None:
        """Initialize RemovedInstance.

        Parameters
        ----------
        filter_element_id: str
            FilterElement which removed the molecule.
        message: str | None, optional
            Information why the molecule was removed.

        """
        self.filter_element_id = filter_element_id
        self.message = message

    def __repr__(self) -> str:
        """Return string representation of RemovedInstance.

        Returns
        -------
        str
            String representation of RemovedInstance.

        """
        return f"RemovedInstance({self.filter_element_id}, {self.message})"


class ABCPipelineElement(abc.ABC):
    """Ancestor of all PipelineElements."""

    name: str
    _requires_fitting: bool = False
    uuid: str

    def __init__(
        self,
        name: str | None = None,
        n_jobs: int = 1,
        uuid: str | None = None,
    ) -> None:
        """Initialize ABCPipelineElement.

        Parameters
        ----------
        name: str | None, optional
            Name of PipelineElement
        n_jobs: int, default=1
            Number of cores used for processing.
        uuid: str | None, optional
            Unique identifier of the PipelineElement.

        """
        if name is None:
            name = self.__class__.__name__
        self.name = name
        self.n_jobs = n_jobs
        if uuid is None:
            self.uuid = str(uuid4())
        else:
            self.uuid = uuid

    def __repr__(self) -> str:
        """Return string representation of object.

        Returns
        -------
        str
            String representation of object.

        """
        parm_list = []
        for key, value in self._get_non_default_parameters().items():
            parm_list.append(f"{key}={value}")
        parm_str = ", ".join(parm_list)
        return f"{self.__class__.__name__}({parm_str})"

    def __sklearn_tags__(self) -> Tags:
        """Return Tags for the Element.

        Returns
        -------
        Tags
            Tags for the Element.

        """
        return Tags(
            estimator_type=None,
            target_tags=TargetTags(required=False),
            transformer_tags=None,
            regressor_tags=None,
            classifier_tags=None,
        )

    def _get_non_default_parameters(self) -> dict[str, Any]:
        """Return all parameters which are not default.

        Returns
        -------
        dict[str, Any]
            Dictionary of non default parameters.

        """
        signature = inspect.signature(self.__class__.__init__)
        self_params = self.get_params()
        non_default_params = {}
        for parm, value_default in signature.parameters.items():
            if parm == "uuid":
                continue
            if value_default is inspect.Parameter.empty:
                non_default_params[parm] = self_params[parm]
                continue
            if parm in self_params and value_default.default != self_params[parm]:
                non_default_params[parm] = self_params[parm]
        return non_default_params

    def get_params(self, deep: bool = True) -> dict[str, Any]:
        """Return the parameters of the object.

        Parameters
        ----------
        deep: bool, default=True
            If True get a deep copy of the parameters.

        Returns
        -------
        dict[str, Any]
            Parameters of the object.

        """
        if deep:
            return {
                "name": copy.copy(self.name),
                "n_jobs": copy.copy(self.n_jobs),
                "uuid": copy.copy(self.uuid),
            }

        return {
            "name": self.name,
            "n_jobs": self.n_jobs,
            "uuid": self.uuid,
        }

    def set_params(self, **parameters: Any) -> Self:
        """Change the parameters of the object.

        Parameters
        ----------
        parameters: Any
            Parameters to be set.

        Returns
        -------
        Self
            Self with updated parameters.

        Raises
        ------
        ValueError
            If the parameter is not a valid parameter of the object.

        """
        for att_name, att_value in parameters.items():
            if not hasattr(self, att_name):
                raise ValueError(
                    f"Cannot set attribute {att_name} on {self.__class__.__name__}",
                )
            setattr(self, att_name, att_value)
        return self

    @property
    def n_jobs(self) -> int:
        """Get the number of cores."""
        return self._n_jobs

    @n_jobs.setter
    def n_jobs(self, n_jobs: int) -> None:
        """Set the number of cores.

        Parameters
        ----------
        n_jobs: int
            Number of cores used for processing.

        """
        self._n_jobs = check_available_cores(n_jobs)

    @property
    def requires_fitting(self) -> bool:
        """Return whether the object requires fitting or not."""
        return self._requires_fitting

    def fit(self, values: Any, labels: Any = None) -> Self:
        """Fit object to input_values.

        Most objects might not need fitting, but it is implemented for
        consistency for all PipelineElements.

        Parameters
        ----------
        values: Any
            List of molecule representations.
        labels: Any
            Optional label for fitting.

        Returns
        -------
        Self
            Fitted object.

        """
        _ = self.fit_transform(values, labels)
        return self

    def fit_to_result(self, values: Any) -> Self:  # noqa: ARG002
        """Fit object to result of transformed values.

        Fit object to the result of the transform function. This is useful for catching
        invalid or removed molecules.

        Parameters
        ----------
        values: Any
            List of molecule representations.

        Returns
        -------
        Self
            Fitted object.

        """
        return self

    @abc.abstractmethod
    def fit_transform(
        self,
        values: Any,
        labels: Any = None,
    ) -> Any:
        """Apply fit function and subsequently transform the input.

        Parameters
        ----------
        values: Any
            Apply transformation specified in transform_single to the value_list.
        labels: Any
            Optional label for fitting.

        Returns
        -------
        Any
            List of instances in new representation.

        """

    @abc.abstractmethod
    def transform(self, values: Any) -> Any:
        """Transform input_values according to object rules.

        Parameters
        ----------
        values: Any
            Iterable of instances to be transformed.
            Could be list of molecule representations (SMILES, Molecules, PhysChem
            vectors etc.).

        Returns
        -------
        Any
            Transformed input_values.

        """

    @abc.abstractmethod
    def transform_single(self, value: Any) -> Any:
        """Transform a single value.

        Parameters
        ----------
        value: Any
            Value to be transformed.

        Returns
        -------
        Any
            Transformed value.

        """


class TransformingPipelineElement(ABCPipelineElement):
    """Ancestor of all PipelineElements."""

    _input_type: str
    _output_type: str
    name: str

    def __init__(
        self,
        name: str | None = None,
        n_jobs: int = 1,
        uuid: str | None = None,
    ) -> None:
        """Initialize ABCPipelineElement.

        Parameters
        ----------
        name: str | None, optional
            Name of PipelineElement
        n_jobs: int, default=1
            Number of cores used for processing.
        uuid: str | None, optional
            Unique identifier of the PipelineElement.

        """
        super().__init__(name=name, n_jobs=n_jobs, uuid=uuid)
        self._is_fitted = False

    @property
    def input_type(self) -> str:
        """Return the input type."""
        return self._input_type

    @property
    def is_fitted(self) -> bool:
        """Return whether the object is fitted or not."""
        return self._is_fitted

    @property
    def output_type(self) -> str:
        """Return the output type."""
        return self._output_type

    @property
    def parameters(self) -> dict[str, Any]:
        """Return the parameters of the object."""
        return self.get_params()

    @parameters.setter
    def parameters(self, parameters: Any) -> None:
        """Set the parameters of the object.

        Parameters
        ----------
        parameters: Any
            Object parameters as a dictionary.

        """
        self.set_params(**parameters)

    def fit_to_result(self, values: Any) -> Self:
        """Fit object to result of transformed values.

        Fit object to the result of the transform function. This is useful
        catching invalid or removed molecules.

        Parameters
        ----------
        values: Any
            List of molecule representations.

        Returns
        -------
        Self
            Fitted object.

        """
        self._is_fitted = True
        return super().fit_to_result(values)

    @override
    def fit_transform(
        self,
        values: Any,
        labels: Any = None,
    ) -> Any:
        """Apply fit function and subsequently transform the input.

        Parameters
        ----------
        values: Any
            Apply transformation specified in transform_single to the value_list.
        labels: Any
            Optional label for fitting.

        Returns
        -------
        Any
            List of molecules in new representation.

        """
        self._is_fitted = True
        if self.requires_fitting:
            pre_value_list = self.pretransform(values)
            self.fit_to_result(pre_value_list)
            output_list = self.finalize_list(pre_value_list)
            if hasattr(self, "assemble_output"):
                return self.assemble_output(output_list)
        return self.transform(values)

    def transform_single(self, value: Any) -> Any:
        """Transform a single molecule to the new representation.

        RemovedMolecule objects are passed without change.

        Parameters
        ----------
        value: Any
            Current representation of the molecule.

        Returns
        -------
        Any
            New representation of the molecule.

        """
        if isinstance(value, InvalidInstance):
            return value
        pre_value = self.pretransform_single(value)
        if isinstance(pre_value, InvalidInstance):
            return pre_value
        return self.finalize_single(pre_value)

    def assemble_output(self, value_list: Iterable[Any]) -> Any:
        """Aggregate rows, which in most cases is just return the list.

        Some representations might be better representd as a single object.
        For example a list of vectors can be transformed to a matrix.

        Parameters
        ----------
        value_list: Iterable[Any]
            Iterable of transformed rows.

        Returns
        -------
        Any
            Aggregated output. This can also be the original input.

        """
        return list(value_list)

    @abc.abstractmethod
    def pretransform_single(self, value: Any) -> Any:
        """Transform the instance, but skips parameters learned during fitting.

        This is the first step for the full transformation. It is followed by the
        finalize_single method and assemble output which collects all single
        transformations. These functions are split as they need to be accessed
        separately from outside the object.

        Parameters
        ----------
        value: Any
            Value to be pretransformed.

        Returns
        -------
        Any
            Pretransformed value.

        """

    def finalize_single(self, value: Any) -> Any:
        """Apply parameters learned during fitting to a single instance.

        Parameters
        ----------
        value: Any
            Value obtained from pretransform_single.

        Returns
        -------
        Any
            Finalized value.

        """
        return value

    def pretransform(self, value_list: Iterable[Any]) -> list[Any]:
        """Transform without using fitting specific parameters.

        Parameters
        ----------
        value_list: Iterable[Any]
            Iterable of instances to be pretransformed.

        Returns
        -------
        list[Any]
            Transformed input_values.

        """
        parallel = Parallel(n_jobs=self.n_jobs)
        return parallel(
            delayed(self.pretransform_single)(value) for value in value_list
        )

    def finalize_list(self, value_list: Iterable[Any]) -> list[Any]:
        """Transform list of values according to parameters learned during fitting.

        Parameters
        ----------
        value_list:  Iterable[Any]
            List of values to be transformed.

        Returns
        -------
        list[Any]
            List of transformed values.

        """
        parallel = Parallel(n_jobs=self.n_jobs)
        return parallel(delayed(self.finalize_single)(value) for value in value_list)

    def transform(self, values: Any) -> Any:
        """Transform input_values according to object rules.

        Parameters
        ----------
        values: Any
            Iterable of instances to be transformed.
            Could be list of molecule representations (SMILES, Molecules, PhysChem
            vectors etc.).

        Returns
        -------
        Any
            Transformed input_values.

        """
        output_rows = self.pretransform(values)
        output_rows = self.finalize_list(output_rows)
        return self.assemble_output(output_rows)


class MolToMolPipelineElement(TransformingPipelineElement, abc.ABC):
    """Abstract PipelineElement where input and outputs are molecules."""

    _input_type = "RDKitMol"
    _output_type = "RDKitMol"

    def transform(self, values: list[OptionalMol]) -> list[OptionalMol]:
        """Transform list of molecules to list of molecules.

        Parameters
        ----------
        values: list[OptionalMol]
            List of molecules to be transformed.

        Returns
        -------
        list[OptionalMol]
            List of molecules.
            InvalidInstances, if the transformation was not successful.

        """
        mol_list: list[OptionalMol] = super().transform(values)  # Stupid mypy...
        return mol_list

    def transform_single(self, value: OptionalMol) -> OptionalMol:
        """Apply pretransform_single and finalize_single in one step.

        Parameters
        ----------
        value: OptionalMol
            Molecule to be transformed.

        Returns
        -------
        OptionalMol
            Transformed molecule if transformation was successful, else InvalidInstance.

        """
        try:
            return super().transform_single(value)
        except Exception as exc:
            if isinstance(value, Chem.Mol):
                logger.error(
                    f"Failed to process: {Chem.MolToSmiles(value)} | "
                    f"({self.name}) ({self.uuid})",
                )
            else:
                logger.error(
                    f"Failed to process: {value} ({type(value)}) | "
                    f"({self.name}) ({self.uuid})",
                )
            raise exc

    @abc.abstractmethod
    def pretransform_single(self, value: RDKitMol) -> OptionalMol:
        """Transform the molecule to another molecule object.

        Do not apply parameters learned during fitting.

        Parameters
        ----------
        value: RDKitMol
            Molecule to be transformed.

        Returns
        -------
        OptionalMol
            Transformed molecule if transformation was successful, else InvalidInstance.

        """


class AnyToMolPipelineElement(TransformingPipelineElement, abc.ABC):
    """Abstract PipelineElement which creates molecules from different inputs."""

    _output_type = "RDKitMol"

    def transform(self, values: Any) -> list[OptionalMol]:
        """Transform list of instances to list of molecules.

        Parameters
        ----------
        values: Any
            Instances to be transformed to a list of molecules.

        Returns
        -------
        list[OptionalMol]
            List of molecules.
            InvalidInstance if the representation was invalid.

        """
        mol_list: list[OptionalMol] = super().transform(values)  # Stupid mypy...
        return mol_list

    @abc.abstractmethod
    def pretransform_single(self, value: Any) -> OptionalMol:
        """Transform a single instance to a molecule.

        The transformations only includes steps which do not depend on parameters
        learned during fitting.

        Parameters
        ----------
        value: Any
            Representation to be transformed to a molecule.

        Returns
        -------
        OptionalMol
            Obtained molecule if valid representation, else InvalidInstance.

        """


class MolToAnyPipelineElement(TransformingPipelineElement, abc.ABC):
    """Abstract PipelineElement which creates molecules from different inputs."""

    _input_type = "RDKitMol"

    @abc.abstractmethod
    def pretransform_single(self, value: RDKitMol) -> Any:
        """Transform the molecule, but skip parameters learned during fitting.

        Parameters
        ----------
        value: RDKitMOl
            Molecule to be transformed.

        Returns
        -------
        Any
            Transformed molecule.

        """
