from featuretools.primitives import TransformPrimitive
from featuretools.tests.testing_utils import make_ecommerce_entityset
from woodwork.column_schema import ColumnSchema
from woodwork.logical_types import Datetime, NaturalLanguage
import featuretools as ft
import numpy as np
import re

"""
Explanation: Advanced Custom Primitives Guide
End of explanation
"""


class StringCount(TransformPrimitive):
    '''Count the number of times the string value occurs.'''
    name = 'string_count'
    input_types = [ColumnSchema(logical_type=NaturalLanguage)]
    return_type = ColumnSchema(semantic_tags={'numeric'})

    def __init__(self, string=None):
        self.string = string

    def get_function(self):
        def string_count(column):
            assert self.string is not None, "string to count needs to be defined"
            # this is a naive implementation used for clarity
            counts = [text.lower().count(self.string) for text in column]
            return counts

        return string_count

"""
Explanation: Primitives with Additional Arguments
Some features require more advanced calculations than others. Advanced features usually entail additional arguments to help output the desired value. With custom primitives, you can use primitive arguments to help you create advanced features.
String Count Example
In this example, you will learn how to make custom primitives that take in additional arguments. You will create a primitive to count the number of times a specific string value occurs inside a text.
First, derive a new transform primitive class using TransformPrimitive as a base. The primitive will take in a text column as the input and return a numeric column as the output, so set the input type to a Woodwork ColumnSchema with logical type NaturalLanguage and the return type to a Woodwork ColumnSchema with the semantic tag 'numeric'. The specific string value is the additional argument, so define it as a keyword argument inside __init__. Then, override get_function to return a primitive function that will calculate the feature.
Featuretools' primitives use Woodwork's ColumnSchema to control the input and return types of columns for the primitive. For more information about using the Woodwork typing system in Featuretools, see the Woodwork Typing in Featuretools guide.
End of explanation
"""


es = make_ecommerce_entityset()

feature_matrix, features = ft.dfs(
    entityset=es,
    target_dataframe_name="sessions",
    agg_primitives=["sum", "mean", "std"],
    trans_primitives=[StringCount(string="the")],
)

feature_matrix[[
    'STD(log.STRING_COUNT(comments, string=the))',
    'SUM(log.STRING_COUNT(comments, string=the))',
    'MEAN(log.STRING_COUNT(comments, string=the))',
]]

"""
Explanation: Now you have a primitive that is reusable for different string values. For example, you can create features based on the number of times the word "the" appears in a text. Create an instance of the primitive where the string value is "the" and pass the primitive into DFS to generate the features. The feature name will automatically reflect the string value of the primitive.
End of explanation
"""


class CaseCount(TransformPrimitive):
    '''Return the count of upper case and lower case letters of a text.'''
    name = 'case_count'
    input_types = [ColumnSchema(logical_type=NaturalLanguage)]
    return_type = ColumnSchema(semantic_tags={'numeric'})
    number_output_features = 2

    def get_function(self):
        def case_count(array):
            # this is a naive implementation used for clarity
            upper = np.array([len(re.findall('[A-Z]', i)) for i in array])
            lower = np.array([len(re.findall('[a-z]', i)) for i in array])
            return upper, lower

        return case_count

"""
Explanation: Features with Multiple Outputs
Some calculations output more than a single value. With custom primitives, you can make the most of these calculations by creating a feature for each output value.
Case Count Example
In this example, you will learn how to make custom primitives that output multiple features. You will create a primitive that outputs the count of upper case and lower case letters of a text.
First, derive a new transform primitive class using TransformPrimitive as a base. The primitive will take in a text column as the input and return two numeric columns as the output, so set the input type to a Woodwork ColumnSchema with logical type NaturalLanguage and the return type to a Woodwork ColumnSchema with semantic tag 'numeric'. Since this primitive returns two columns, also set number_output_features to two. Then, override get_function to return a primitive function that will calculate the feature and return a list of columns.
End of explanation
"""


feature_matrix, features = ft.dfs(
    entityset=es,
    target_dataframe_name="sessions",
    agg_primitives=[],
    trans_primitives=[CaseCount],
)

feature_matrix[[
    'customers.CASE_COUNT(favorite_quote)[0]',
    'customers.CASE_COUNT(favorite_quote)[1]',
]]

"""
Explanation: Now you have a primitive that outputs two columns. One column contains the count for the upper case letters. The other column contains the count for the lower case letters. Pass the primitive into DFS to generate features. By default, the feature name will reflect the index of the output.
End of explanation
"""


class HourlySineAndCosine(TransformPrimitive):
    '''Returns the sine and cosine of the hour.'''
    name = 'hourly_sine_and_cosine'
    input_types = [ColumnSchema(logical_type=Datetime, semantic_tags={'time_index'})]
    return_type = ColumnSchema(semantic_tags={'numeric'})

    number_output_features = 2

    def get_function(self):
        def hourly_sine_and_cosine(column):
            sine = np.sin(column.dt.hour)
            cosine = np.cos(column.dt.hour)
            return sine, cosine

        return hourly_sine_and_cosine

    def generate_names(self, base_feature_names):
        name = self.generate_name(base_feature_names)
        return f'{name}[sine]', f'{name}[cosine]'

"""
Explanation: Custom Naming for Multiple Outputs
When you create a primitive that outputs multiple features, you can also define custom naming for each of those features.
Hourly Sine and Cosine Example
In this example, you will learn how to apply custom naming for multiple outputs. You will create a primitive that outputs the sine and cosine of the hour.
First, derive a new transform primitive class using TransformPrimitive as a base. The primitive will take in the time index as the input and return two numeric columns as the output. Set the input type to a Woodwork ColumnSchema with a logical type of Datetime and the semantic tag 'time_index'. Next, set the return type to a Woodwork ColumnSchema with semantic tag 'numeric' and set number_output_features to two. Then, override get_function to return a primitive function that will calculate the feature and return a list of columns. Also, override generate_names to return a list of the feature names that you define.
End of explanation
"""


feature_matrix, features = ft.dfs(
    entityset=es,
    target_dataframe_name="log",
    agg_primitives=[],
    trans_primitives=[HourlySineAndCosine],
)

feature_matrix.head()[[
    'HOURLY_SINE_AND_COSINE(datetime)[sine]',
    'HOURLY_SINE_AND_COSINE(datetime)[cosine]',
]]

"""
Explanation: Now you have a primitive that outputs two columns. One column contains the sine of the hour. The other column contains the cosine of the hour. Pass the primitive into DFS to generate features. The feature name will reflect the custom naming you defined.
End of explanation
"""


