import copy
import logging
from datetime import timedelta, datetime

import dateparser
import numpy as np
import pandas as pd

from utils.aws_connectors import run_query
from utils.base_attributes_generation import get_fan_attribute_values
from .algo import Algo

component_logger = logging.getLogger().getChild("enrichment.enrich_rfm")


class EnrichRFM(Algo):
    """
    Calculates Request, Frequency, Monetary analysis. This approach mimics almost exactly what's in the SQL
    sql/analytics_rfm_analysis.sql. Python inspiration from https://towardsdatascience.com/recency-frequency-monetary-model-with-python-and-how-sephora-uses-it-to-optimize-their-google-d6a0707c5f17

    General Logic:
        - will run the final Segments assignment only if all 3 attributes are available 73, 74, 75
    """

    source_attribute_ids = [
        # enrPurchaseDate
        73,
        # enrPurchaseMonetary
        74,
        # enrPurchaseQuantity
        75,
    ]
    result_attribute_ids = [
        50,  # enrUserRFM
        61,  # enrTransactionRecency
        76,  # enrTransactionFrequency
        110,  # enrTransactionMonetary
        77,  # enrR
        78,  # enrF
        79,  # enrM
    ]
    segment_attribute_name = "enrUserRFM"
    segment_name = "RFM: {}"
    data_gathering_rule = "strict"
    # TODO! create extra method that extends self._get_collection_type() to update rule to loose to support SETS

    # Needed for dynamically inserting attributes into collection_attribute table
    result_attribute_dict = {
        "enrUserRFM": 50,
        "enrTransactionRecency": 61,
        "enrRFMScore": 62,
        "enrTransactionFrequency": 76,
        "enrR": 77,
        "enrF": 78,
        "enrM": 79,
        "enrTransactionMonetary": 110,
    }
    # TODO! can be obtained dynamically with extra SQL query

    def __init__(self, collection_id: int, schema: str, user_id: str, default_value: str = "unknown", **kwargs):
        super().__init__(collection_id, schema, user_id, default_value)

    # @timeit
    def _run(self):
        """ Main enrichment method
            Will NOT produce RFM_Segment_Concat, enrRFMScore, enrUserRFM if any of R F M are not present
        """
        source_df = self.available_data["source"]

        # Convert date string to date object first, so we could add timedelta later; numeric cols to numeric;
        source_df[["73", "75", "74"]] = source_df.apply(
            lambda x: pd.Series(
                [
                    self._get_timestamp(x["73"]),
                    pd.to_numeric(x["75"], errors="coerce"),
                    pd.to_numeric(x["74"], errors="coerce"),
                ]
            ),
            axis=1,
        )
        """ Interpolate all empty values for Monetary to 0 """
        source_df["74"] = source_df[["74"]].apply(lambda x: x[0] if not pd.isnull(x[0]) else 0, axis=1)

        """ Interpolate all empty values for Quantity to 0 """
        source_df["75"] = source_df[["75"]].apply(lambda x: x[0] if not pd.isnull(x[0]) else 0, axis=1)

        # Aggregate numeric columns per fan and take their sum
        df_numeric = source_df.groupby(["fan_id"]).agg({x: ["sum"] for x in ["74", "75"]})
        df_numeric = df_numeric.reset_index(level=0)
        df_numeric.columns = [col[0] for col in df_numeric.columns]

        # Aggregate date col and take max per fan
        df_date = source_df.groupby(["fan_id"]).agg({x: ["max"] for x in ["73"]})
        df_date = df_date.reset_index(level=0)
        df_date.columns = [col[0] for col in df_date.columns]

        # Merge prepared and aggregated data
        source_df = df_numeric.merge(df_date, how="inner", on="fan_id")

        source_df = self._generate_r_f_m_values(source_df)

        """ Only generate RFM if all necessary columns (R, F and M) are available """
        if len(set(source_df.columns).intersection({"R", "F", "M"})) == 3:
            """ Getting rfm segments """
            segments_df = self._get_rfm_segments()
            segments_df.rename(
                columns={"recency": "R", "frequency": "F", "monetary": "M", "segment_name": "enrUserRFM"}, inplace=True
            )

            """ Calculate RFM_Score """
            source_df["enrRFMScore"] = source_df[["R", "F", "M"]].sum(axis=1)

            """ Create a new variable, with some classic RFM segmenting descriptors
                Dropping rows where any of R, F or M are nulls"""
            # source_df['enrUserRFM'] = source_df.apply(self._rfm_level, axis=1)
            data_for_rfm = copy.deepcopy(source_df.dropna(axis=0, how="any", subset=["R", "F", "M"]))
            data_for_rfm = data_for_rfm.merge(segments_df, how="left", on=["R", "F", "M"])

            source_df = source_df.merge(data_for_rfm[["fan_id", "enrUserRFM"]], how="left", on="fan_id")
            source_df = source_df[~source_df["enrUserRFM"].isna()]

        else:
            """ Making sure that we don't run segments generation if the above logic did not run"""
            self.run_generate_segments = False

        source_df.rename(
            columns={
                "Recency": "enrTransactionRecency",
                "Frequency": "enrTransactionFrequency",
                "MonetaryValue": "enrTransactionMonetary",
                "R": "enrR",
                "F": "enrF",
                "M": "enrM",
            },
            inplace=True,
        )

        self.result_df = source_df

    def _generate_r_f_m_values(self, source_df):
        """Generating separate columns for R, F and M
        Indicating if main RFM logic needs to run or not (if all fields are available)"""
        rfm_dict = {
            "R": lambda x: self._calculate_r(x),
            "F": lambda x: self._calculate_f(x),
            "M": lambda x: self._calculate_m(x),
        }

        """ Dropping columns that doesn't contain any values """
        check = source_df[["73", "75", "74"]].dropna(axis=1, how="all")

        available_columns = list(check.rename(columns={"73": "R", "75": "F", "74": "M"}).columns)

        result_df = pd.DataFrame({"fan_id": []})
        for column in available_columns:
            data = rfm_dict[column](
                copy.deepcopy(source_df)
            )  # TODO! what is happening here, when do we use "check" vs "source_df"
            # TODO! Can be handled without copy.deepcopy() to reduce memory usage
            result_df = result_df.merge(data, on="fan_id", how="outer")

        return result_df

    @staticmethod
    def _calculate_r(source_df):  # Triggers for 73
        snapshot_date = source_df["73"].max() + timedelta(days=1)
        source_df = source_df.groupby(["fan_id"]).agg({"73": lambda x: (snapshot_date - x.max()).days})
        source_df = source_df[["73"]].rename(columns={"73": "Recency"})

        # Create labels for Recency
        r_labels = range(4, 0, -1)
        # Assign these labels to 4 equal percentile groups
        r_groups = pd.qcut(source_df["Recency"].rank(method="first"), q=4, labels=r_labels)
        # Create new column
        source_df = source_df.assign(R=r_groups.values)

        source_df.reset_index(inplace=True)
        return source_df[["fan_id", "R", "Recency"]]

    def _calculate_f(self, source_df):  # Triggers for 75
        source_df["Frequency"] = source_df.groupby("fan_id")["75"].transform("sum")
        source_df = source_df[["fan_id", "Frequency"]]
        source_df.drop_duplicates(inplace=True)

        # # Create labels for Frequency
        # f_labels = range(1, 5)
        # # Assign these labels to 4 equal percentile groups
        # f_groups = pd.qcut(source_df['Frequency'].rank(method='first'), q=4, labels=f_labels)
        # # Create new column
        # source_df = source_df.assign(F=f_groups.values)

        source_df["F"] = source_df.apply(lambda x: pd.Series(self._get_f_value(x["Frequency"])), axis=1)

        source_df.reset_index(inplace=True)
        return source_df[["fan_id", "F", "Frequency"]]

    @staticmethod
    def _get_f_value(x):
        """ Specifying the conditions to assign F values"""
        if x == 1:
            return 1
        elif int(x) == 0:
            # Had 0 "tracks played" in one pilot dataset, which we can also label as quantity
            return 0
        elif x in [2, 3]:
            return 2
        elif x >= 4:
            return 3
        else:
            raise RuntimeError(f"enrRFM - Unexpected F value assigned - {x}")

    @staticmethod
    def _calculate_m(source_df):  # Triggers for 74
        source_df = source_df.groupby(["fan_id"]).agg({"74": "sum"})
        source_df.rename(columns={"74": "MonetaryValue"}, inplace=True)

        # Create labels for Monetary
        m_labels = range(1, 5)
        # Assign these labels to 4 equal percentile groups
        m_groups = pd.qcut(source_df["MonetaryValue"].rank(method="first"), q=4, labels=m_labels)
        # Create new column
        source_df = source_df.assign(M=m_groups.values)

        source_df.reset_index(inplace=True)
        return source_df[["fan_id", "M", "MonetaryValue"]]

    @staticmethod
    def _get_date_object(date_string):
        """ Get date object out of string. We use dateparser to conveniently make date object out
        of almost any string representation of date - this is pretty fool-proof and can work with stuff like
        "20 November 2020" as well as "22/11/2016"
        """
        date_object = np.nan
        try:
            date_object = dateparser.parse(str(date_string), settings={"PREFER_DATES_FROM": "past"})
        except Exception as e:
            pass
        return date_object

    @staticmethod
    def _get_timestamp(date_string):
        try:
            return datetime.strptime(date_string, "%Y-%m-%d")
        except:
            return np.nan

    @staticmethod
    def _get_rfm_segments() -> pd.DataFrame:
        """ Obtaining the predefined RFM segments"""
        query = """
            SELECT *
            FROM commons.rfm_segments
        """
        segments_df = run_query(query, return_type="df")
        return segments_df

    @staticmethod
    def _populate_date_field(df):
        """ Replacing the null values of the date_column with max value per fan_id"""
        if "73" in df.columns:
            agg_df = df[["fan_id", "73"]].groupby(by="fan_id").max()
            agg_df = agg_df.reset_index()
            agg_df.rename(columns={"73": "aggregated_date"}, inplace=True)

            df = df.merge(agg_df, on="fan_id", how="left")
            df["73"] = df[["73", "aggregated_date"]].apply(lambda x: x[0] if not pd.isnull(x[0]) else x[1], axis=1)
            df.drop(columns="aggregated_date", inplace=True)
        else:
            component_logger.info("Skipping _populate_date_field - attribute is missing in source_df")
            pass

        return df

    # We have this here because we don't need to merge with original row (and cannot in case of multiple sources)
    def _generate_fan_attributes(self):
        """ Getting a list of attributes that are present in this enrichment"""
        condition = self.system_fields["a_id"].isin(self.result_attribute_ids)
        attributes_df = self.system_fields[condition][["a_system_name", "a_id"]]
        attributes_df.rename(columns={"a_system_name": "name", "a_id": "attribute_id"}, inplace=True)

        fan_attributes_df = get_fan_attribute_values(
            source_df=self.result_df, attributes_df=attributes_df, collection_id=self.enrichment_collection_id
        )
        self.fan_attribute = fan_attributes_df