"""SQL Expression Transform entity."""

from typing import List

from pyspark.sql import DataFrame
from pyspark.sql.functions import expr

from butterfree.transform.transformations.transform_component import TransformComponent


class SQLExpressionTransform(TransformComponent):
    """Defines an SQL Expression Transformation.

    Attributes:
        expression: SQL expression defined by the user.

    Example:
        It's necessary to declare the custom SQL query, such that it
        corresponds to operations between existing columns in the dataframe.
        Besides, the usage of SparkSQL functions is also allowed. Finally,
        a "complete select statement", such as "select col_a * col_b from
        my_table", is not necessary, just the simple operations are
        required, for instance "col_a / col_b" or "col_a * col_b".

        >>> from butterfree.transform.features import Feature
        >>> from butterfree.transform.transformations import SQLExpressionTransform
        >>> from butterfree.constants import DataType
        >>> from pyspark import SparkContext
        >>> from pyspark.sql import session
        >>> import pyspark.sql.functions as F
        >>> sc = SparkContext.getOrCreate()
        >>> spark = session.SparkSession(sc)
        >>> df = spark.createDataFrame([(1, "2016-04-11 11:31:11", 200, 200),
        ...                             (1, "2016-04-11 11:44:12", 300, 300),
        ...                             (1, "2016-04-11 11:46:24", 400, 400),
        ...                             (1, "2016-04-11 12:03:21", 500, 500)]
        ...                           ).toDF("id", "timestamp", "feature1", "feature2")
        >>> feature = Feature(
        ...    name="feature",
        ...    description="SQL expression transform usage example",
        ...    dtype=DataType.DOUBLE,
        ...    transformation=SQLExpressionTransform(expression="feature1/feature2"),
        ...)
        >>> feature.transform(df).orderBy("timestamp").show()
        +--------+--------+---+-------------------+----------------------+
        |feature1|feature2| id|          timestamp|feature1_over_feature2|
        +--------+--------+---+-------------------+----------------------+
        |     200|     200|  1|2016-04-11 11:31:11|                   1.0|
        |     300|     300|  1|2016-04-11 11:44:12|                   1.0|
        |     400|     400|  1|2016-04-11 11:46:24|                   1.0|
        |     500|     500|  1|2016-04-11 12:03:21|                   1.0|
        +--------+--------+---+-------------------+----------------------+

    """

    def __init__(
        self, expression: str,
    ):
        super().__init__()
        self.expression = expression

    @property
    def output_columns(self) -> List[str]:
        """Columns generated by the transformation."""
        return [self._parent.name]

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

        Args:
            dataframe: input dataframe.

        Returns:
            Transformed dataframe.

        """
        return dataframe.withColumn(self._parent.name, expr(self.expression))
