"""Aggregated Transform entity."""
from collections import namedtuple
from typing import List, Tuple

from pyspark.sql import DataFrame
from pyspark.sql.functions import col, expr, when

from butterfree.transform.transformations.transform_component import TransformComponent
from butterfree.transform.utils.function import Function


class AggregatedTransform(TransformComponent):
    """Specifies an aggregation.

    This transformation needs to be used within an AggregatedFeatureSet. Unlike the
    other transformations, this class won't have a transform method implemented.

    The idea behing aggregating is that, in spark, we should execute all aggregation
    functions after a single groupby. So an AggregateFeatureSet will have many Features
    with AggregatedTransform. If each one of them needs to apply a groupby.agg(), then
    we must join all the results in the end, making this computation extremely slow.

    Now, the AggregateFeatureSet will collect all Features' AggregatedTransform
    definitions and run, at once, a groupby.agg(*aggregations).

    This class helps defining on a feature, which aggregation function will be applied
    to build a new aggregated column. Allowed aggregations are registered under the
     allowed_aggregations property.

    Attributes:
        functions: namedtuple with aggregation function and data type.
        filter_expression: sql boolean expression to be used inside agg function.
            The filter expression can be used to aggregate some column only with
            records that obey certain condition. Has the same behaviour of the
            following SQL expression: `agg(case when filter_expression then col end)`

    Example:
        >>> from butterfree.transform.transformations import AggregatedTransform
        >>> from butterfree.transform.features import Feature
        >>> from butterfree.constants import DataType
        >>> from butterfree.transform.utils import Function
        >>> import pyspark.sql.functions as F
        >>> feature = Feature(
        ...     name="feature",
        ...     description="aggregated transform",
        ...     transformation=AggregatedTransform(
        ...         functions=[
        ...                    Function(F.avg, DataType.DOUBLE),
        ...                    Function(F.stddev_pop, DataType.DOUBLE)],
        ...     ),
        ...     from_column="somenumber",
        ...)
        >>> feature.get_output_columns()
        ['feature__avg', 'feature__stddev_pop']
        >>> feature.transform(anydf)
        NotImplementedError: ...
    """

    def __init__(self, functions: List[Function], filter_expression: str = None):
        super(AggregatedTransform, self).__init__()
        self.functions = functions
        self.filter_expression = filter_expression

    @property
    def aggregations(self) -> List[Tuple]:
        """Aggregated spark columns."""
        column_name = self._parent.from_column or self._parent.name

        # if transform has a filter expression apply inside agg function
        # if not, use just the target column name inside agg function
        expression = (
            when(expr(self.filter_expression), col(column_name))
            if self.filter_expression
            else column_name
        )
        Function = namedtuple("Function", ["function", "data_type"])

        return [
            Function(f.func(expression), f.data_type.spark,) for f in self.functions
        ]

    def _get_output_name(self, function: object) -> str:
        if not hasattr(function, "__name__"):
            feature_name = self._parent.name
            raise AttributeError(
                f"""Anonymous functions are not supported on AggregatedTransform.
                    Check feature {feature_name} transforms.
                """
            )

        base_name = "__".join([self._parent.name, str(function.__name__).lower()])
        return base_name

    @property
    def output_columns(self) -> List[str]:
        """Columns names generated by the transformation."""
        return [self._get_output_name(f.func) for f in self.functions]

    def transform(self, dataframe: DataFrame) -> DataFrame:
        """(NotImplemented) Performs a transformation to the feature pipeline.

        For the AggregatedTransform, the transformation won't be applied without
        using an AggregatedFeatureSet.

        Args:
            dataframe: input dataframe.

        Raises:
            NotImplementedError.
        """
        raise NotImplementedError(
            "AggregatedTransform won't be used outside an AggregatedFeatureSet, "
            "meaning the responsibility of aggregating and apply the transformation is "
            "now over the FeatureSet component. This should optimize the ETL process."
        )
