# Copyright (c) 2023 ING Analytics Wholesale Banking
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
# the Software, and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
from __future__ import annotations

import warnings

import numpy as np
import pandas as pd

from popmon.base import Module
from popmon.utils import parallel


class ApplyFunc(Module):
    """This module applies functions to specified feature and metrics.

    Extra parameters (kwargs) can be passed to the apply function.
    """

    _input_keys = ("apply_to_key", "assign_to_key", "apply_funcs_key")
    _output_keys = ("store_key",)

    def __init__(
        self,
        apply_to_key,
        store_key: str = "",
        assign_to_key: str = "",
        apply_funcs_key: str = "",
        features=None,
        apply_funcs=None,
        metrics=None,
        msg: str = "",
    ) -> None:
        """Initialize an instance of ApplyFunc.

        :param str apply_to_key: key of the input data to apply funcs to.
        :param str assign_to_key: key of the input data to assign function applied-output to. (optional)
        :param str store_key: key of the output data to store in the datastore (optional)
        :param str apply_funcs_key: key of to-be-applied functions in data to store (optional)
        :param list features: list of features to pick up from input data and apply funcs to (optional)
        :param list metrics: list of metrics to apply funcs to (optional)
        :param str msg: message to print out at start of transform function. (optional)
        :param list apply_funcs: functions to apply (list of dicts):

          - 'func': function to apply
          - 'suffix' (string, optional): suffix added to each metric. default is function name.
          - 'prefix' (string, optional): prefix added to each metric.
          - 'features' (list, optional): features the function is applied to. Overwrites features above
          - 'metrics' (list, optional): metrics the function is applied to. Overwrites metrics above
          - 'entire' (boolean, optional): apply function to the entire feature's dataframe of metrics?
          - 'args' (tuple, optional): args for 'func'
          - 'kwargs' (dict, optional): kwargs for 'func'

        """
        super().__init__()
        self.apply_to_key = apply_to_key
        self.assign_to_key = assign_to_key or apply_to_key
        self.apply_funcs_key = apply_funcs_key
        self.store_key = store_key or self.assign_to_key

        self.features = features or []
        self.metrics = metrics or []
        self.msg = msg
        self.apply_funcs = []
        # import applied functions
        apply_funcs = apply_funcs or []
        for af in apply_funcs:
            self.add_apply_func(**af)

    def get_description(self):
        if len(self.apply_funcs) > 0:
            return " and ".join([x["func"].__name__ for x in self.apply_funcs])
        elif self.apply_funcs_key:
            return f"functions from arg '{self.apply_funcs_key}'"
        else:
            raise NotImplementedError

    def add_apply_func(
        self,
        func,
        suffix=None,
        prefix=None,
        metrics=None,
        features=None,
        entire=None,
        *args,
        **kwargs,
    ):
        """Add function to be applied to dataframe.

        Can call this function after module instantiation to add new functions.

        :param func: function to apply
        :param suffix: (string, optional) suffix added to each metric. default is function name.
        :param prefix: (string, optional) prefix added to each metric.
        :param features: (list, optional) features the function is applied to. Overwrites features above
        :param metrics: (list, optional) metrics the function is applied to. Overwrites metrics above
        :param entire: (boolean, optional) apply function to the entire feature's dataframe of metrics?
        :param args: (tuple, optional) args for 'func'
        :param kwargs: (dict, optional) kwargs for 'func'
        """
        # check inputs
        if features is None:
            features = []
        if metrics is None:
            metrics = []
        if not callable(func):
            raise TypeError("functions in ApplyFunc must be callable objects")
        if suffix is not None and not isinstance(suffix, str):
            raise TypeError("prefix, and suffix in ApplyFunc must be strings or None.")
        if prefix is not None and not isinstance(prefix, str):
            raise TypeError("prefix, and suffix in ApplyFunc must be strings or None.")
        if not isinstance(metrics, list) or not isinstance(features, list):
            raise TypeError("metrics and features must be lists of strings.")

        # add function
        self.apply_funcs.append(
            {
                "features": features,
                "metrics": metrics,
                "func": func,
                "entire": entire,
                "suffix": suffix,
                "prefix": prefix,
                "args": args,
                "kwargs": kwargs,
            }
        )

    def transform(
        self,
        apply_to_data: dict,
        assign_to_data: dict | None = None,
        apply_funcs: list | None = None,
    ):
        """
        Apply functions to specified feature and metrics

        Each feature/metric combination is treated as a pandas series

        :param datastore: input datastore
        :return: updated datastore
        :rtype: dict
        """
        assert isinstance(apply_to_data, dict)
        if assign_to_data is None:
            assign_to_data = {}

        if apply_funcs is not None:
            self.apply_funcs += apply_funcs

        if self.msg:
            self.logger.info(self.msg)

        features = self.get_features(list(apply_to_data.keys()))

        same_key = self.assign_to_key == self.apply_to_key

        args = [
            {
                "feature": feature,
                "metrics": self.metrics,
                "apply_to_df": self.get_datastore_object(
                    apply_to_data, feature, dtype=pd.DataFrame
                ),
                "assign_to_df": (
                    None
                    if same_key
                    else self.get_datastore_object(
                        assign_to_data,
                        feature,
                        dtype=pd.DataFrame,
                        default=pd.DataFrame(),
                    )
                ),
                "apply_funcs": self.apply_funcs,
                "same_key": same_key,
            }
            for feature in features
        ]
        result = parallel(apply_func_array, args, mode="kwargs")
        new_metrics = dict(result)
        return new_metrics


def apply_func_array(
    feature, metrics, apply_to_df, assign_to_df, apply_funcs, same_key
):
    """Apply list of functions to dataframe

    Split off for parallelization reasons

    :param str feature: feature currently looping over
    :param list metrics: list of selected metrics to apply functions to
    :param apply_to_df: pandas data frame that function in arr is applied to
    :param assign_to_df: pandas data frame the output of function is assigned to
    :param apply_funcs: list of functions to apply to
    :param same_key: if True, merge apply_to_df and assign_to_df before returning assign_to_df
    :return: union of feature and assign_to_df
    """
    if not isinstance(apply_to_df, pd.DataFrame):
        raise TypeError(
            f'apply_to_df of feature "{feature}" is not a pandas dataframe.'
        )

    if same_key or (len(assign_to_df.index) == 0 and len(assign_to_df.columns) == 0):
        assign_to_df = pd.DataFrame(index=apply_to_df.index)

    for arr in apply_funcs:
        obj = apply_func(feature, metrics, apply_to_df, arr)
        if len(obj) == 0:
            # no metrics were found in apply_to_df
            continue
        for new_metric, o in obj.items():
            if isinstance(o, pd.Series):
                if len(assign_to_df.index) == len(o) and all(
                    assign_to_df.index == o.index
                ):
                    assign_to_df[new_metric] = o
                else:
                    warnings.warn(
                        f"{feature}:{new_metric}: df_out and object have inconsistent lengths."
                    )
            else:
                # o is number or object, assign to every element of new column
                assign_to_df[new_metric] = [o] * len(assign_to_df.index)
    if same_key:
        assign_to_df = pd.concat([apply_to_df, assign_to_df], axis=1)
    return feature, assign_to_df


def apply_func(feature, selected_metrics, df, arr):
    """Apply function to dataframe

    :param str feature: feature currently looping over
    :param list selected_metrics: list of selected metrics to apply to
    :param df: pandas data frame that function in arr is applied to
    :param dict arr: dictionary containing the function to be applied to pandas dataframe.
    :return: dictionary with outputs of applied-to metric pd.Series
    """
    # basic checks of feature
    if (
        "features" in arr
        and len(arr["features"]) > 0
        and feature not in arr["features"]
    ):
        return {}

    # get func input
    keys = list(arr.keys())

    assert "func" in keys, "function input is insufficient."
    func = arr["func"]

    if "prefix" not in keys or arr["prefix"] is None:
        arr["prefix"] = ""
    if len(arr["prefix"]) > 0 and not arr["prefix"].endswith("_"):
        arr["prefix"] = arr["prefix"] + "_"
    if "suffix" not in keys or arr["suffix"] is None:
        arr["suffix"] = func.__name__ if len(arr["prefix"]) == 0 else ""
    if len(arr["suffix"]) > 0 and not arr["suffix"].startswith("_"):
        arr["suffix"] = "_" + arr["suffix"]
    suffix = arr["suffix"]
    prefix = arr["prefix"]

    args = ()
    kwargs = {}
    if "kwargs" in keys:
        kwargs = arr["kwargs"]
    if "args" in keys:
        args = arr["args"]

    # apply func
    if len(selected_metrics) > 0 or ("metrics" in keys and len(arr["metrics"]) > 0):
        metrics = (
            arr["metrics"]
            if ("metrics" in keys and len(arr["metrics"]) > 0)
            else selected_metrics
        )
        metrics = [m for m in metrics if m in df.columns]
        if len(metrics) == 0:
            return {}
        df = df[metrics] if len(metrics) >= 2 else df[metrics[0]]

    if (
        "entire" in arr
        and arr["entire"] is not None
        and arr["entire"] is not False
        and arr["entire"] != 0
    ):
        obj = func(df, *args, **kwargs)
    else:
        obj = df.apply(func, args=args, **kwargs)

    # convert object to dict format
    if not isinstance(
        obj, (pd.Series, pd.DataFrame, list, tuple, np.ndarray)
    ) and isinstance(df, pd.Series):
        obj = {df.name: obj}
    elif not isinstance(
        obj, (pd.Series, pd.DataFrame, list, tuple, np.ndarray)
    ) and isinstance(df, pd.DataFrame):
        obj = {"_".join(df.columns): obj}
    elif (
        isinstance(obj, (list, tuple, np.ndarray))
        and isinstance(df, pd.DataFrame)
        and len(df.columns) == len(obj)
    ):
        obj = dict(zip(df.columns, obj))
    elif (
        isinstance(obj, (list, tuple, np.ndarray))
        and isinstance(df, pd.Series)
        and len(df.index) == len(obj)
    ):
        obj = {df.name: pd.Series(data=obj, index=df.index)}
    elif (
        isinstance(obj, (list, tuple, np.ndarray))
        and isinstance(df, pd.DataFrame)
        and len(df.index) == len(obj)
    ):
        obj = {"_".join(df.columns): pd.Series(data=obj, index=df.index)}
    elif (
        isinstance(obj, pd.Series)
        and isinstance(df, pd.Series)
        and len(obj) == len(df)
        and all(obj.index == df.index)
    ):
        obj = {df.name: obj}
    elif (
        isinstance(obj, pd.Series)
        and isinstance(df, pd.DataFrame)
        and len(obj) == len(df)
        and all(obj.index == df.index)
    ):
        obj = {"_".join(df.columns): obj}
    # e.g. output of normalized_hist_mean_cov: a dataframe with one column, actually a series
    elif (
        isinstance(obj, pd.DataFrame)
        and len(obj.columns) == 1
        and (
            len(obj.index) != len(df.index)
            or (len(obj.index) == len(df.index) and (obj.index != df.index).any())
        )
    ):
        obj = obj[obj.columns[0]].to_dict()
    elif isinstance(obj, pd.Series):
        # e.g. output of np.mean of np.std: results in one number per column when applied to a dataframe
        obj = obj.to_dict()
    elif isinstance(obj, pd.DataFrame):
        # when apply to one column returns another column
        obj = {c: obj[c] for c in obj.columns}
    assert isinstance(obj, dict)

    # add prefix and suffix to dict keys
    keys = list(obj.keys())
    for k in keys:
        obj[prefix + k + suffix] = obj.pop(k)

    return obj
