import logging
import os
import time
import uuid
from ast import literal_eval

import mlflow
import numpy as np
import pandas as pd
import pycountry_convert as pc

from utils.aws_connectors import s3_to_pandas, pandas_to_s3
from utils.base_attributes_generation import get_fan_attribute_values
from utils.data_model_utils import write_csv_to_efs
from .algo import Algo
from .enrichment import MissingInputFieldException

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

PROFILE = os.getenv("PROFILE", "devel")

MLFLOW_URL = f"https://mlflow.fansifter.cloud{':444' if PROFILE == 'devel-local-toomas' else ''}"

class EnrichMachineLearningClusters(Algo):
    """
    Makes call to ml-engine and does machine learning clustering. Works better the more qualitative features there are.
    Also does some model results interpretation.
    """

    # NB! Notice that when adding columns, add them to either numeric_cols or categorical_cols list as well
    source_attribute_ids = [
        97,  # 'enrAge'
        # 83,  # 'purchaseQuantity' # This is already present in the face of enrPurchaseQuantity
        # 84,  # 'purchaseMonetary' # This is already present in the face of enrPurchaseMonetary
        42,  # 'userRating'
        43,  # 'userMailchimpRating'
        98,  # 'enrUniqueGender'
        # Time columns don't return actionable segments. See discussion in #data-science
        # 51,  # 'enrHourOfDay'
        # 52,  # 'enrDayOfMonth'
        # 53,  # 'enrDayOfWeek'
        54,  # 'enrCountry'
        55,  # 'enrLocality'
        61,  # 'enrTransactionRecency'
        74,  # 'enrPurchaseMonetary'
        75,  # 'enrPurchaseQuantity'
        76,  # 'enrTransactionFrequency'
        94,  # 'enrMaxMonetary'
    ]  # Works if any of those are present
    result_attribute_ids = [70]  # enrMLCluster
    segment_attribute_name = "enrMLCluster"
    segment_name = "ML: {}"
    data_gathering_rule = "strict"

    new_variable_importance = {
        # dictionary to hold new feature importance values so we can order features based on business logic
        "74": 0.9,  # enrPurchaseMonetary
        "75": 0.85,  # enrPurchaseMonetary
        "94": 0.8,  # enrMaxMonetary
        "37": 0.75,  # merchItemPrice
        "42": 0.7,  # userRating
        "43": 0.65,  # userMailchimpRating
        "76": 0.6,  # enrTransactionFrequency
        "61": 0.55,  # enrTransactionRecency
        "98": 0.5,  # enrUniqueGender,
        "genderfemale": 0.5,  # female gender,
        "gendermale": 0.5,  # male gender,
        "genderunknown": 0.5,  # unknown gender,
        "97": 0.45,  # enrAge
        "54": 0.4,  # enrCountry
        "continentAF": 0.39,  # continentAF
        "continentNA": 0.39,  # continentNA
        "continentOC": 0.39,  # continentOC
        "continentAN": 0.39,  # continentAN
        "continentAS": 0.39,  # continentAS
        "continentEU": 0.39,  # continentEU
        "continentSA": 0.39,  # continentSA
    }

    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)

        # ['preprocessing', 'postprocessing'], 'processing' (ECS or Batch) is executed and monitored by Airflow itself
        self.phase = kwargs["phase"]

        ## FOR POSTPROCESSING PHASE
        # retrieving preprocessing data back for postprocessing call
        self.grouped_df_csv_path = kwargs.get("grouped_df_csv_path")
        self.enrichment_collection_id = kwargs.get("enrichment_collection_id")

        # retrieving processing results for postprocessing call
        self.mlflow_run_id = kwargs.get("mlflow_run_id")

    def do_everything(self):
        try:
            """ Running the enrichment """
            if self.phase == "preprocessing":
                is_prepared = self._prepare(populate_tables=True)
                if not is_prepared:
                    return None
                preprocessing_results = self._run_preprocessing()
                return preprocessing_results
            elif self.phase == "postprocessing":
                is_prepared = self._prepare(populate_tables=False)
                if not is_prepared:
                    return None
                self._run_postprocessing()
                csv_name = self._finish()
                result = {
                    "csv": csv_name,
                    "collection_id": self.enrichment_collection_id,
                    "schema": self.schema,
                    **self._get_segment_spec(),
                }
                return result
            else:
                raise ValueError(f"Unknown phase: {self.phase}")

        except MissingInputFieldException as e:
            """ We don't really use this exception so we'd not pollute alerts monitoring with log errors, and we 
            handle this by returning instead. But leaving this be in case we'll figure out a more elegant solution here.
            """
            self._update_collection_status("failed")
            return {
                "error": str(e),
            }
        except Exception as e:
            component_logger.exception(e)
            self._update_collection_status("failed")
            return {
                "error": str(e),
            }

    def get_keyword_mapping(self):
        keyword_mapping_dict = {
            # userAge
            "97": {
                "formulas": [
                    {"formula": "between", "threshold": "1 and 24", "keyword": "Age up to 25 yrs"},
                    {"formula": "between", "threshold": "25 and 44", "keyword": "Age 25-44 yrs"},
                    {"formula": "between", "threshold": "45 and 65", "keyword": "Age 45-65 yrs"},
                    {"formula": "between", "threshold": "65 and 125", "keyword": "Age 65+ yrs"},
                ]
            },
            # userRating
            "42": {
                "formulas": [
                    {"formula": "<", "threshold": "1.5", "keyword": "Lowest Rating Fans"},
                    {"formula": ">", "threshold": "3.5", "keyword": "Highest Rating Fans"},
                ]
            },
            # purchaseQuantity
            "83": {
                "formulas": [
                    {"formula": "==", "threshold": "1", "keyword": "Single Item Buyers"},
                    {"formula": "between", "threshold": "2 and 3", "keyword": "Multiple Item Buyers"},
                    {"formula": ">", "threshold": "3", "keyword": "Bulk Item Buyers"},
                ]
            },
            # purchaseMonetary
            "84": {
                "formulas": [
                    {"formula": "<", "threshold": "quartile_1", "keyword": "Budget-Friendly Fans"},
                    {
                        "formula": "between",
                        "threshold": "quartile_1 and quartile_3",
                        "keyword": "Fans With Average Cost Baskets",
                    },
                    {"formula": ">", "threshold": "quartile_3", "keyword": "Big Spenders"},
                ]
            },
            # userMailchimpRating
            "43": {
                "formulas": [
                    {"formula": "<", "threshold": "1.5", "keyword": "Mailchimp Resubscribers/Softbouncers"},
                    {"formula": "between", "threshold": "1.5 and 2.5", "keyword": "Mailchimp New Contacts"},
                    {"formula": "between", "threshold": "2.5 and 3.5", "keyword": "Mailchimp Low Engagement Fans"},
                    {"formula": "between", "threshold": "3.5 and 4.2", "keyword": "Mailchimp Moderate Engagement Fans"},
                    {"formula": ">", "threshold": "4.2", "keyword": "Mailchimp High Engagement Fans"},
                ]
            },
            # enrUniqueGender
            "98": {
                "formulas": [
                    {"formula": "between", "threshold": "0.65 and 1", "keyword": "Mostly Women"},
                    {"formula": "between", "threshold": "0 and 0.35", "keyword": "Mostly Men"},
                    {"formula": "between", "threshold": "0.35 and 0.65", "keyword": "Mixed with Men and Women"},
                    {"formula": "between", "threshold": "-1 and -0.35", "keyword": "Gender mostly unknown"},
                ]
            },
            # gender female
            "genderfemale": {
                "formulas": [{"formula": "between", "threshold": "0.65 and 1", "keyword": "Mostly Women"}]
            },
            # gender male
            "gendermale": {"formulas": [{"formula": "between", "threshold": "0.65 and 1", "keyword": "Mostly Men"}]},
            # enrHourOfDay
            "51": {
                "formulas": [
                    {"formula": "between", "threshold": "6 and 8", "keyword": "Early Morning Buyers"},
                    {"formula": "between", "threshold": "0 and 5", "keyword": "Night Buyers"},
                    {"formula": "between", "threshold": "9 and 12", "keyword": "Morning/Noon Buyers"},
                    {"formula": "between", "threshold": "13 and 18", "keyword": "Afternoon Buyers"},
                    {"formula": "between", "threshold": "19 and 21", "keyword": "Evening Buyers"},
                    {"formula": "between", "threshold": "22 and 24", "keyword": "Late Evening Buyers"},
                ]
            },
            # enrDayOfWeek
            "53": {
                "formulas": [
                    {"formula": "between", "threshold": "1 and 2", "keyword": "Beginning of Week Buyers"},
                    {"formula": "between", "threshold": "3 and 4", "keyword": "Midweek Buyers"},
                    {"formula": "==", "threshold": "5", "keyword": "Friday Buyers"},
                    {"formula": "between", "threshold": "6 and 7", "keyword": "Weekend Buyers"},
                ]
            },
            # enrDayOfMonth
            "52": {
                "formulas": [
                    {"formula": "between", "threshold": "1 and 7", "keyword": "First Week of Month Buyers"},
                    {"formula": "between", "threshold": "8 and 15", "keyword": "Second Week of Month Buyers"},
                    {"formula": "between", "threshold": "16 and 23", "keyword": "Third Week of Month Buyers"},
                    {"formula": "between", "threshold": "24 and 31", "keyword": "Last Week of Month Buyers"},
                ]
            },
            # enrMaxMonetary; use same mapping logic as for regular monetary field
            "94": {
                "formulas": [
                    {"formula": "<", "threshold": "quartile_1", "keyword": "Low Cost Item Buyers"},
                    {"formula": ">", "threshold": "quartile_3", "keyword": "High Cost Item Buyers"},
                ]
            },
            # enrPurchaseMonetary
            "74": {
                "formulas": [
                    {"formula": "<", "threshold": "quartile_1", "keyword": "Budget-Friendly Fans"},
                    {
                        "formula": "between",
                        "threshold": "quartile_1 and quartile_3",
                        "keyword": "Fans with Average Cost Baskets",
                    },
                    {"formula": ">", "threshold": "quartile_3", "keyword": "Big Spenders"},
                ]
            },
            # enrPurchaseQuantity
            "75": {
                "formulas": [
                    {"formula": "==", "threshold": "1", "keyword": "Single Item Buyers"},
                    {"formula": "between", "threshold": "2 and 3", "keyword": "Multiple Item Buyers"},
                    {"formula": ">", "threshold": "3", "keyword": "Bulk Item Buyers"},
                ]
            },
            # enrTransactionFrequency
            "76": {
                "formulas": [
                    {"formula": "<", "threshold": "1.7", "keyword": "Single Transaction Fans"},
                    {"formula": "between", "threshold": "1.7 and 3", "keyword": "Fans with Multiple Transactions"},
                ]
            },
            # merchItemPrice
            "37": {
                "formulas": [
                    {"formula": "<", "threshold": "quartile_1", "keyword": "Inexpensive Item Buyers"},
                    {"formula": ">", "threshold": "quartile_3", "keyword": "Expensive Item Buyers"},
                ]
            },
            # enrTransactionRecency
            "61": {
                "formulas": [
                    {"formula": "<", "threshold": "quartile_1", "keyword": "Early Bird Buyers"},
                    {"formula": ">", "threshold": "quartile_3", "keyword": "Last Minute Buyers"},
                ]
            },
            # enrCountry
            "54": {
                "formulas": [
                    {"formula": "==", "threshold": "1", "keyword": "Mostly Fans from North America"},
                    {"formula": "==", "threshold": "2", "keyword": "Mostly Fans from Great Britain"},
                    {"formula": "==", "threshold": "3", "keyword": "Mostly Fans from Nordic Countries"},
                    {"formula": "==", "threshold": "4", "keyword": "Mostly Fans from DACH Countries"},
                    {"formula": "==", "threshold": "5", "keyword": "Mostly Fans from Benelux Countries"},
                    {"formula": "==", "threshold": "6", "keyword": "Mostly Fans from Spain and Portugal"},
                    {"formula": "==", "threshold": "7", "keyword": "Mostly Fans from Australia and NZ"},
                    {"formula": "==", "threshold": "8", "keyword": "Mostly Fans from France and Italy"},
                ]
            },
            # continent Africa(AF)
            "continentAF": {"formulas": [{"formula": "between", "threshold": "0.65 and 1", "keyword": "From Africa"}]},
            # continent North America(NA)
            "continentNA": {
                "formulas": [{"formula": "between", "threshold": "0.65 and 1", "keyword": "From North America"}]
            },
            # continent Oceania(OC)
            "continentOC": {"formulas": [{"formula": "between", "threshold": "0.65 and 1", "keyword": "from Oceania"}]},
            # continent Antartica(AN)
            "continentAN": {
                "formulas": [{"formula": "between", "threshold": "0.65 and 1", "keyword": "from Antartica"}]
            },
            # continent Asia(AS)
            "continentAS": {"formulas": [{"formula": "between", "threshold": "0.65 and 1", "keyword": "from Asia"}]},
            # continent Europe(EU)
            "continentEU": {"formulas": [{"formula": "between", "threshold": "0.65 and 1", "keyword": "from Europe"}]},
            # continent South America(SA)
            "continentSA": {
                "formulas": [{"formula": "between", "threshold": "0.65 and 1", "keyword": "from south America"}]
            },
        }
        return keyword_mapping_dict

    def map_segments_to_keywords(self, df_aggregated, quantiles):
        """ Maps dataframe to the keywords map based on this logic
        - https://fansifter.atlassian.net/wiki/spaces/FM/pages/233930753/Interpretation+to+ML+modeling+outputs

        The aim of this is to get some interpretable labels/naming to ML clustering results """

        keyword_dict = self.get_keyword_mapping()

        # Initialize dict with default values we'll be showing when we don't get interpretations
        segment_keywords_results = {x: f"{x} - Uninterpreted segment" for x in df_aggregated["enrMLCluster_"].unique()}

        # Iterate over each segment
        for idx, row in df_aggregated.iterrows():

            segment_keywords = []

            # Iterate over each feature
            segment = row["enrMLCluster_"]
            for column in [x for x in df_aggregated.columns if x != "enrMLCluster_"]:

                # In case our aggregated dataframe median for the cluster is nan (can happen), skip this iteration
                # Evaluating to np.nan or None did not work, that's why we compare to "nan" string
                if str(row[column]) == "nan":
                    continue

                # With some necessary dataframe aggregation and flattening operations we get unnecessary underscores
                original_col = column.split("_")[0]

                # don't apply rule for monetary value if it's 0
                if (original_col in ("74", "94")) & (str(row[column]) == "0.0"):
                    continue

                # Iterate over each formula, unless we have don't have a mapping defined
                try:
                    formulas = keyword_dict[original_col]["formulas"]
                except Exception as e:
                    component_logger.error(
                        f"{self.schema} - No ML segment keyword mapping defined for feature {original_col}: {e}"
                    )
                    continue

                for formula in formulas:
                    threshold = formula["threshold"]
                    operator = formula["formula"]
                    keyword = formula["keyword"]

                    # handle quartiles and generate threshold values
                    if threshold == "quartile_1":
                        threshold = quantiles.loc[0.33, original_col]  # locate first quartile for feature
                    elif threshold == "quartile_3":
                        threshold = quantiles.loc[0.75, original_col]  # locate third quartile for feature

                    try:
                        # TODO! refactor and remove duplication
                        # handle the operation between quartiles
                        if operator == "between" and threshold == "quartile_1 and quartile_3":
                            threshold_q1 = quantiles.loc[0.33, original_col]  # locate first quartile for feature
                            threshold_q3 = quantiles.loc[0.75, original_col]  # locate third quartile for feature
                            if eval(f"{row[column]} >= {threshold_q1}") and eval(f"{row[column]} <= {threshold_q3}"):
                                segment_keywords.append(keyword)
                        # handle the SQL-like "between A and B" operation
                        elif operator == "between":
                            left = threshold.split(" and ")[0]
                            right = threshold.split(" and ")[1]
                            if eval(f"{row[column]} >= {float(left)}") and eval(f"{row[column]} <= {float(right)}"):
                                segment_keywords.append(keyword)
                        # otherwise use the operator we're provided with
                        else:
                            if eval(f"{row[column]} {operator} {float(threshold)}"):
                                segment_keywords.append(keyword)
                    except Exception as e:
                        component_logger.exception(f"Error applying keywords to segments: {e}")
                        """ We don't care when we get issues, we simply don't map segment and thus exclude later.
                        For example, sometimes we get None for aggregated median (happens in cases we are not handling 
                        TOTAL rows in input data, and when extreme (total) values get their own cluster, and when the
                        median will also thus be "nan" (this will produce error in "eval"). We handle this in FSM-574.
                        """
                        continue

            # TODO! This is ugly, should fix it in above logic instead removing here afterwards
            if len(segment_keywords) > 1 and "Fans with Average Cost Baskets" in segment_keywords:
                segment_keywords.remove("Fans with Average Cost Baskets")

            # We have in some cases same mapping for different input labels. Keep only unique keywords and keep order.
            segment_keywords = list(dict.fromkeys(segment_keywords))

            # In case we don't get any keywords, we use the default values
            if len(segment_keywords) > 0:
                segment_keywords_results[segment] = ", ".join(segment_keywords)
        return segment_keywords_results

    def _run_preprocessing(self):
        def country_to_iso_to_cont(country_name):
            try:
                iso_country = pc.country_name_to_country_alpha2(country_name)
                iso_continent = pc.country_alpha2_to_continent_code(iso_country)
                return iso_continent
            except:
                return "unknown"

        source_df = self.available_data["source"]

        # Drop dupes that we might get with reruns etc
        source_df = source_df.drop_duplicates()

        """ Dropping columns that don't contain any values """
        source_df = source_df.dropna(axis=1, how="all")

        """ Only include numeric columns that we want to sum. Don't include "hour of week" and alike in numeric columns.
        Since we previously also drop rows that don't have any data, we create intersection of the available cols 
        and of what cols we have left in the source df.
        """
        numeric_cols = list({"83", "74", "75"} & set(source_df.columns.to_list()))
        categorical_cols = list(
            {"97", "42", "43", "98", "51", "52", "53", "54", "55", "61", "76", "94"} & set(source_df.columns.to_list())
        )

        # Sort df so we'd ensure consecutive runs give same output. Important when we take "first" value later.
        source_df = source_df.sort_values(by=["fan_id", "row_id"], ascending=True).reset_index(drop=True)

        # Turn our columns to numeric or aggregation fails
        if len(source_df) > 1000 and hasattr(source_df, "parallel_apply"):
            source_df[numeric_cols] = source_df[numeric_cols].parallel_apply(pd.to_numeric, axis=0, args=["coerce"])
        else:
            source_df[numeric_cols] = source_df[numeric_cols].apply(lambda x: pd.to_numeric(x, errors="coerce"))

        # Aggregate numeric columns and take their sum. We aggregate so we'd do fan not row_id based clustering.
        df_numeric = source_df.groupby(["fan_id"]).agg(
            {x: ["sum"] for x in numeric_cols} if len(numeric_cols) > 0 else []
        )
        df_numeric = df_numeric.reset_index(level=0)
        df_numeric.columns = [col[0] for col in df_numeric.columns]

        """ 
        We need a strategy to handle aggregating categorical fields. Ideally, we should take the median non-null value.
        We fillna so we could take the max(). We could use the first(), which would be massively faster, but then we'd 
        also need to order results. With ordering (e.g. fan_id+row_id order), we can't guarantee each run will have 
        same order, in case results come from different collections.
        """
        source_df[numeric_cols].fillna(0, inplace=True)
        source_df[categorical_cols].fillna("", inplace=True)

        """ We can have situations when either numeric cols are missing, or cat cols are missing. """
        if not source_df[categorical_cols].empty:
            # max() on string columns is slow AF because it's doing alphanumeric sorting for all values
            # df_text = source_df.groupby("fan_id").max().reset_index()[categorical_cols + ['fan_id']]
            df_text = source_df.groupby("fan_id").first().reset_index()[categorical_cols + ["fan_id"]]

            # Merge texts with the numeric cols aggregations if possible
            if not source_df[numeric_cols].empty:
                grouped_df = df_numeric.merge(df_text, how="inner", on="fan_id")
            else:
                grouped_df = df_text
            del df_text  # conserve some memory
        else:
            # In case we don't get any categorical columns
            grouped_df = df_numeric

        del df_numeric  # conserve some memory
        del source_df

        # apply dummy encoding to selected features
        # apply dummy encoding for gender
        if "98" in categorical_cols:
            grouped_df = pd.get_dummies(grouped_df, prefix="gender", prefix_sep="", columns=["98"])

        if "54" in categorical_cols:
            grouped_df["country_continent"] = grouped_df["54"].apply(country_to_iso_to_cont)
            grouped_df = pd.get_dummies(grouped_df, prefix="continent", prefix_sep="", columns=["country_continent"])

        bucket = "fansifter-model-data"
        unique_id = uuid.uuid4()
        timestr = time.strftime("%Y%m%d-%H%M%S")
        file_key = f"{self.schema}/{unique_id}-{timestr}"
        pandas_to_s3(bucket, grouped_df, file_key)

        collection_ids = [int(cid) for cid in self.collection_ids["source"]]

        grouped_df_csv_path = write_csv_to_efs(
            df=grouped_df,
            schema=self.schema,
            table="fan_attribute",
            header=True,
            attributes=[self.enrichment_name, "preprocessing"],
        )

        return {
            "collection_ids": collection_ids,
            "enrichment_collection_id": self.enrichment_collection_id,
            "grouped_df_csv_path": grouped_df_csv_path,
            "bucket": bucket,
            "file_key": file_key,
            "schema": self.schema,
            "dataset_size": len(grouped_df),
        }

    def _processing(self): # noqa
        # Processing is now performed on the Airflow side:
        # Either calling stand-by ECS ML Engine container for small files
        # Or submitting an AWS Batch job for larger files
        pass

    def _run_postprocessing(self):
        grouped_df = pd.read_csv(self.grouped_df_csv_path, sep="\t")

        # Get model data from Mlflow
        client = mlflow.tracking.MlflowClient(tracking_uri=MLFLOW_URL)
        run_data = client.get_run(self.mlflow_run_id)

        artifact_uri = run_data.info.artifact_uri
        cluster_output_file = run_data.data.params["cluster_output_file"]

        cluster_output_location = artifact_uri + "/" + cluster_output_file
        cluster_output_location = cluster_output_location.replace("s3://fansifter-model-data/", "")

        # download S3 csv to dataframe
        bucket = "fansifter-model-data"
        cluster_output = s3_to_pandas(bucket, cluster_output_location, header="infer", file_format="csv")[
            ["fan_id", "cluster"]
        ]

        cluster_output = cluster_output.astype({"fan_id": int})

        # Take the aggregated dataframe as the source df to do further operations
        source_df = grouped_df.merge(cluster_output, how="left", on="fan_id").rename(
            columns={"cluster": f"enrMLCluster"}
        )

        # Get most important features
        params = run_data.data.params
        signifigance_threshold = 0.05
        importance = literal_eval(params.get("feature_importance"))
        features = literal_eval(params.get("original_features"))
        importance_df = pd.DataFrame({"features": features, "importance": importance})
        important_features = importance_df.loc[importance_df["importance"] >= signifigance_threshold]
        important_features["importance"] = important_features["features"].map(self.new_variable_importance)
        important_features.sort_values(by=["importance"], inplace=True, ascending=False)
        important_features = important_features["features"].tolist()
        component_logger.info(f"{self.schema} - Important features for ML interpretation: {important_features}")

        # Replace gender strings with numeric values otherwise aggregation fails on string(female, male)
        if "98" in source_df.columns:
            source_df["98"] = source_df["98"].replace(["female", "male", "unknown"], ["1", "0", "-1"])

        # replace enrCountry strings with numeric values otherwise aggregations fails
        # mapping rules described here: https://fansifter.atlassian.net/wiki/spaces/FM/pages/233930753/Interpretation+to+ML+modeling+outputs
        enrCountry_dic = {
            "United States": "1",
            "Canada": "1",
            "United Kingdom": "2",
            "Ireland": "2",
            "Northern Ireland": "2",
            "Scotland": "2",
            "Wales": "2",
            "Sweden": "3",
            "Finland": "3",
            "Denmark": "3",
            "Norway": "3",
            "Iceland": "3",
            "Germany": "4",
            "Austria": "4",
            "Switzerland": "4",
            "Netherlands": "5",
            "Belgium": "5",
            "Luxembourg": "5",
            "Spain": "6",
            "Portugal": "6",
            "Australia": "7",
            "New Zealand": "7",
            "France": "8",
            "Italy": "8",
        }
        if "54" in source_df.columns:
            source_df = source_df.replace({"54": enrCountry_dic})

        # Turn our columns to numeric
        if len(source_df) > 1000 and hasattr(source_df, "parallel_apply"):
            source_df[important_features] = source_df[important_features].parallel_apply(
                pd.to_numeric, axis=0, args=["coerce"]
            )
        else:
            source_df[important_features] = source_df[important_features].apply(
                lambda x: pd.to_numeric(x, errors="coerce")
            )

        # calculate cluster sizes
        cluster_sizes_ratios = source_df.groupby(["enrMLCluster"])["fan_id"].count().reset_index(name="cluster_size")

        # Artificially add Purchase Monetary to important features so we'd get keyword mapping when it is present
        if "74" in source_df.columns and "74" not in important_features:
            important_features.append("74")

        # TODO! Do this more gracefully
        # Replace 0 values for numerical monetary related features to make aggregation/quantile functions ignore these values
        if "74" in source_df.columns:
            source_df[["74"]] = source_df[["74"]].replace([0], np.nan)
            # calculate feature 74 size for non missing fans for each cluster
            feature_74_size = (
                source_df[~source_df["74"].isna()]
                .groupby(["enrMLCluster"])["fan_id"]
                .count()
                .reset_index(name="feature_74")
            )
            cluster_sizes_ratios = pd.merge(cluster_sizes_ratios, feature_74_size, how="left", on="enrMLCluster")
            # calculate ratio to be used in keyword mapping for avoiding keyword names when actual fan count with feature 74 is low
            cluster_sizes_ratios["ratio_74"] = cluster_sizes_ratios["feature_74"] / cluster_sizes_ratios["cluster_size"]
            cluster_sizes_ratios.drop("feature_74", axis=1, inplace=True)
            cluster_sizes_ratios.fillna(0, inplace=True)
        if "94" in source_df.columns:
            source_df[["94"]] = source_df[["94"]].replace([0], np.nan)

        # If no important features are present for some reason (e.g., file is too small) - terminate the enrichment
        if len(important_features) == 0:
            self.result_df = pd.DataFrame({"fan_id": [], "enrMLCluster": []})
            return None

        # Get median values for every feature, grouped by ML cluster, and then flatten this structure
        df_aggregated = source_df.groupby(["enrMLCluster"]).agg({x: ["median"] for x in important_features})
        df_aggregated = df_aggregated.reset_index(level=0)
        df_aggregated.columns = ["_".join(col) for col in df_aggregated.columns]

        if "54_median" in df_aggregated.columns:
            # TODO refactor next chunk of code to be more efficient
            # get most frequent value for country per cluster
            df_aggregated_country = (
                source_df[~source_df["54"].isna()]
                .groupby("enrMLCluster")["54"]
                .apply(lambda x: x.value_counts().head(1))
                .reset_index()
                .rename(columns={"level_1": "most_frequent_country"})
            )
            df_aggregated_country.set_index("enrMLCluster", inplace=True)
            # get total fan count per cluster
            df_aggregated_totals = (
                source_df.groupby("enrMLCluster")["fan_id"]
                .count()
                .reset_index()
                .rename(columns={"fan_id": "total_fans"})
            )
            df_aggregated_totals.set_index("enrMLCluster", inplace=True)
            df_aggregated_country = pd.concat([df_aggregated_country, df_aggregated_totals], axis=1)
            # calculate proportion of most frequent country vs total records
            df_aggregated_country["proportion"] = df_aggregated_country["54"] / df_aggregated_country["total_fans"]
            # remove countries below threshold
            df_aggregated_country["most_frequent_country"] = np.where(
                df_aggregated_country["proportion"] > 0.5, df_aggregated_country["most_frequent_country"], np.nan
            )
            df_aggregated_country.drop(["54", "proportion", "total_fans"], axis=1, inplace=True)

            # create dictionary
            dict_54_replacement = df_aggregated_country.T.to_dict("records")[0]
            df_aggregated.rename(columns={"54_median": "54_mode"}, inplace=True)
            # replace values in aggregated ML dataframe
            df_aggregated["54_mode"] = df_aggregated["enrMLCluster_"].apply(lambda x: dict_54_replacement.get(x))

        # for column 74 where ratio is under 0.65 replace monetary value with 0 to avoid keyword mapping
        if "74_median" in df_aggregated.columns:
            for i, row in cluster_sizes_ratios.iterrows():
                cluster = row["enrMLCluster"]
                if row["ratio_74"] < 0.65:
                    df_aggregated.loc[df_aggregated.enrMLCluster_ == cluster, "74_median"] = 0

        # Create quartiles over full dataset
        quantiles = source_df.quantile([0.33, 0.75])

        segment_keywords_results = self.map_segments_to_keywords(df_aggregated, quantiles)

        # Switch existing labels to new ones
        source_df = source_df.replace({"enrMLCluster": segment_keywords_results})

        # As per FSM-540, remove "Uninterpreted segment". TODO! It's a bit ugly removal, handle this logic elsewhere
        # TODO! This bit here is because "nan" fails for the next bit. We get a ton of "nan"s, but ideally shouldn't
        source_df[["enrMLCluster"]] = source_df[["enrMLCluster"]].fillna("Uninterpreted segment")
        source_df = source_df[~source_df.enrMLCluster.str.contains("Uninterpreted segment")]

        if os.path.exists(self.grouped_df_csv_path):
            os.remove(self.grouped_df_csv_path)

        self.result_df = source_df[["fan_id", "enrMLCluster"]]

    # 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
        )

        # Due to some logic weirdness, we used to get duplicate rows. We might not need this any more.
        fan_attributes_df = fan_attributes_df.drop_duplicates()

        self.fan_attribute = fan_attributes_df