import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Any, List, Union

from arthurai.common.exceptions import arthur_excepted
from arthurai.core.viz import style, utils
import warnings

warnings.filterwarnings("ignore")


class DataVisualizer(object):
    def __init__(self, arthur_models):
        self.models = arthur_models

    @arthur_excepted("failed to generate timeline")
    def timeline(self, attribute_name):
        """Generates a visualization of the distribution of an attribute over time.
        For categorical attributes, a stacked area chart over time.
        For continuous attributes, a joyplot showing probability densities over time.
        :raises ArthurUserError: failed due to user error
        :raises ArthurUnexpectedError: failed due to an unexpected error
        """
        for model in self.models:
            if model.get_attribute(attribute_name).categorical:
                self._timeline_categorical(model, attribute_name)
            else:
                self._timeline_continuous(model, attribute_name)

    @arthur_excepted("failed to generate metric series")
    def metric_series(self, metric_names, time_resolution="day"):
        """Generates a line series visualization for selected metrics.

        .. code-block:: python

            # plot both the model's AUC metric and the model's FPR metric by the hour
            arthur_model.viz.metric_series(["auc", "falsePositiveRate"], time_resolution="hour")

        :raises ArthurUserError: failed due to user error
        :raises ArthurUnexpectedError: failed due to an unexpected error
        """
        fig = plt.figure(figsize=(20, 4))
        ax = fig.add_subplot(1, 1, 1)
        predicted_property, ground_truth_property = utils.get_pred_and_gt_attrs(
            self.models[0]
        )
        pal = style.categorical_palette

        # Define list of palettes to be used
        num_metrics = len(metric_names)
        max_metrics = len(style.multi_categorical_palette_names)
        if num_metrics > max_metrics:
            raise UserValueError(
                f"Too many metrics to display. The maximum is {max_metrics}, you wanted to see {num_metrics}"
            )
        color_palette_names = style.multi_categorical_palette_names[0:num_metrics]

        for metric_index, metric_name in enumerate(metric_names):
            color_palette = sns.color_palette(
                palette=color_palette_names[metric_index], n_colors=len(self.models)
            )
            for model_index, model in enumerate(self.models):
                if model.is_batch:
                    time_resolution = "batch_id"
                    query = _metric_series_batch_query(
                        metric_name,
                        ground_truth_property,
                        predicted_property,
                        model.classifier_threshold,
                    )
                else:
                    query = _metric_series_streaming_query(
                        metric_name,
                        time_resolution,
                        ground_truth_property,
                        predicted_property,
                        model.classifier_threshold,
                    )

                response = model.query(query)
                df = pd.DataFrame(response)[::-1]

                label = _label(model, metric_name)
                color = _alpha_color(color_palette[model_index])
                plt.plot(
                    df[time_resolution],
                    df[metric_name],
                    lw=12,
                    color=color,
                    label=label,
                )
                plt.plot(df[time_resolution], df[metric_name], lw=2, color=color)
                plt.xticks([])
                ax.spines["right"].set_visible(False)
                ax.spines["top"].set_visible(False)
                plt.legend(loc="upper left", facecolor="white")

    @arthur_excepted("failed to generate drift series")
    def drift_series(
        self,
        attribute_names: Union[str, List[str]],
        drift_metric: str = "PSI",
        time_resolution: str = "day",
    ):
        """Generates a line series visualization of data drift metrics for selected attributes.

        .. code-block:: python

            # plot the KLDivergence drift of features X1, X2, and X3 by the hour
            arthur_model.viz.drift_series(["X1", "X2", "X3"], drift_metric="KLDivergence", time_resolution="hour")

        :raises ArthurUserError: failed due to user error
        :raises ArthurUnexpectedError: failed due to an unexpected error
        """
        fig = plt.figure(figsize=(20, 4))
        ax = fig.add_subplot(1, 1, 1)
        labels = []

        # if attribute_names is a string, convert to list with one element
        if isinstance(attribute_names, str):
            attribute_names = [attribute_names]

        # Define list of palettes to be used
        color_palettes = _drift_color_palettes(len(attribute_names), len(self.models))

        for model_index, model in enumerate(self.models):
            if model.is_batch:
                time_resolution = "batch_id"
            timestamp_query = _batch_timestamp_query()
            drift_query = {
                "properties": attribute_names,
                "num_bins": 20,
                "rollup": time_resolution,
                "base": {"source": "reference"},
                "target": {"source": "inference"},
                "metric": drift_metric,
            }
            if model.is_batch:
                drift_df = pd.DataFrame(model.query(drift_query, query_type="drift"))
                timestamp_df = pd.DataFrame(model.query(timestamp_query))
                df = (
                    drift_df.rename(columns={"rollup": time_resolution})
                    .join(timestamp_df.set_index(time_resolution), on=time_resolution)
                    .sort_values(by="timestamp")
                    .set_index(time_resolution)
                )
            else:
                df = (
                    pd.DataFrame(model.query(drift_query, query_type="drift"))
                    .sort_values(by="rollup")
                    .set_index("rollup")
                )

            for attribute_name in attribute_names:
                labels.append(_label(model, attribute_name))

            # base_color_palette = style.categorical_palette
            color_palette = color_palettes[model_index]
            # This darkens the color for each successive model version. There are probably better way of doing this
            # but this is just a hack to get it to work for now.
            # for color_tuple in base_color_palette:
            #     color_palette.append((color_tuple[0] / (2 ** model_index), color_tuple[1] / (2 ** model_index), color_tuple[2] / (2 ** model_index), color_tuple[3]))
            df[attribute_names].plot(ax=ax, color=color_palette, lw=8)
            plt.ylabel(drift_metric)
            plt.xticks([])
            plt.xlabel("")
            ax.spines["right"].set_visible(False)
            ax.spines["top"].set_visible(False)
        plt.legend(labels, loc="upper left", facecolor="white")

    def _timeline_categorical(self, model, attribute_name, time_resolution="day"):
        """ """
        print(_label(model, attribute_name))
        if model.is_batch:
            time_resolution = "timestamp"
            query = _timeline_categorical_batch_query(attribute_name, time_resolution)
        else:
            query = _timeline_categorical_streaming_query(
                attribute_name, time_resolution
            )
        response = model.query(query)
        df = (
            pd.DataFrame(response)
            .sort_values(by=time_resolution)
            .pivot(index=time_resolution, columns=attribute_name, values="count")
            .fillna(0.0)
        )

        sns.set(style="white")
        fig = plt.figure(figsize=(20, 4))
        ax = fig.add_subplot(1, 1, 1)
        plt.stackplot(
            df.index.values,
            df.values.T.tolist(),
            labels=df.columns.values,
            colors=style.categorical_palette,
        )
        plt.xticks([])
        plt.ylabel("Count")
        plt.legend(loc="upper right")
        ax.spines["right"].set_visible(False)
        ax.spines["top"].set_visible(False)
        plt.show()

    def _timeline_continuous(self, model, attribute_name, time_resolution="day"):
        """ """
        print(_label(model, attribute_name))
        if model.is_batch:
            time_resolution = "timestamp"
            query = _timeline_continuous_batch_query(attribute_name, time_resolution)
        else:
            query = _timeline_continuous_streaming_query(
                attribute_name, time_resolution
            )
        response = model.query(query)

        dfs = []
        for group in response:
            temp_df = pd.DataFrame(group["distribution"])
            temp_df[time_resolution] = group[time_resolution]
            temp_df["count"] = utils.savgol_filter(temp_df["count"])
            dfs.append(temp_df)
        df = pd.concat(dfs).sort_values([time_resolution, "lower"])

        global_y_min = df["lower"].min()
        global_y_max = df["lower"].max()
        global_x_min = df["count"].min()
        global_x_max = df["count"].max()

        sns.set(style="white", rc={"axes.facecolor": (0, 0, 0, 0)})
        # Initialize the FacetGrid object
        pal = style.continuous_palette
        g = sns.FacetGrid(
            df,
            col=time_resolution,
            hue=time_resolution,
            aspect=0.2,
            height=8,
            palette=pal,
            subplot_kws={
                "xlim": (global_x_max, global_x_min),
                "ylim": (global_y_min, global_y_max),
            },
        )

        # Draw the densities in a few steps
        g.map(
            plt.plot, "count", "lower", clip_on=False, color="white", lw=4.5, alpha=0.6
        )
        g.map(plt.fill_between, "count", "lower", clip_on=False, lw=2)

        # Set the subplots to overlap
        g.fig.subplots_adjust(wspace=-0.55, left=0.125)
        g.set_titles("")
        g.set(yticks=[])
        g.set(xticks=[])
        g.set(ylabel="")
        g.set(xlabel="")
        g.despine(bottom=True, left=True)


def _metric_series_streaming_query(
    metric_name,
    time_resolution,
    ground_truth_property,
    predicted_property,
    classifier_threshold,
):
    """Generates query for fetching metrics for streaming model."""
    query = {
        "select": [
            {
                "function": metric_name,
                "alias": metric_name,
                "parameters": {
                    "ground_truth_property": ground_truth_property,
                    "predicted_property": predicted_property,
                    "threshold": classifier_threshold,
                },
            },
            {
                "function": "roundTimestamp",
                "alias": time_resolution,
                "parameters": {
                    "property": "inference_timestamp",
                    "time_interval": time_resolution,
                },
            },
        ],
        "group_by": [{"alias": time_resolution}],
        "order_by": [{"alias": time_resolution, "direction": "desc"}],
    }
    return query


def _metric_series_batch_query(
    metric_name, ground_truth_property, predicted_property, classifier_threshold
):
    """Generates query for fetching metrics for batch model."""
    query = {
        "select": [
            {
                "function": metric_name,
                "alias": metric_name,
                "parameters": {
                    "ground_truth_property": ground_truth_property,
                    "predicted_property": predicted_property,
                    "threshold": classifier_threshold,
                },
            },
            {
                "function": "max",
                "parameters": {"property": "inference_timestamp"},
                "alias": "timestamp",
            },
            {"property": "batch_id"},
        ],
        "group_by": [{"property": "batch_id"}],
        "order_by": [{"alias": "timestamp", "direction": "desc"}],
    }
    return query


def _timeline_continuous_streaming_query(attribute_name, time_resolution):
    """Generates a query for continuous attribute in a streaming model."""
    return {
        "select": [
            {
                "function": "distribution",
                "alias": "distribution",
                "parameters": {"property": attribute_name, "num_bins": 50},
            },
            {
                "function": "roundTimestamp",
                "alias": time_resolution,
                "parameters": {
                    "property": "inference_timestamp",
                    "time_interval": time_resolution,
                },
            },
        ],
        "group_by": [{"alias": time_resolution}],
        "order_by": [{"alias": time_resolution, "direction": "desc"}],
    }


def _timeline_continuous_batch_query(attribute_name, time_resolution):
    """Generates a query for continuous attribute in a batch model."""
    return {
        "select": [
            {
                "function": "distribution",
                "alias": "distribution",
                "parameters": {"property": attribute_name, "num_bins": 50},
            },
            {
                "function": "max",
                "alias": "timestamp",
                "parameters": {"property": "inference_timestamp"},
            },
            {"property": "batch_id"},
        ],
        "group_by": [{"property": "batch_id"}],
        "order_by": [{"alias": "timestamp", "direction": "desc"}],
    }


def _timeline_categorical_streaming_query(attribute_name, time_resolution):
    """Generates a query for categorical attribute in a streaming model."""
    return {
        "select": [
            {"property": attribute_name},
            {"function": "count"},
            {
                "function": "roundTimestamp",
                "alias": time_resolution,
                "parameters": {
                    "property": "inference_timestamp",
                    "time_interval": time_resolution,
                },
            },
        ],
        "group_by": [{"property": attribute_name}, {"alias": time_resolution}],
        "order_by": [
            {"property": attribute_name},
            {"alias": time_resolution, "direction": "desc"},
        ],
    }


def _timeline_categorical_batch_query(attribute_name, time_resolution):
    """Generates a query for categorical attribute in a batch model."""
    return {
        "select": [
            {"property": attribute_name},
            {"function": "count"},
            {"property": "batch_id"},
            {
                "function": "max",
                "parameters": {"property": "inference_timestamp"},
                "alias": "timestamp",
            },
        ],
        "group_by": [{"property": attribute_name}, {"property": "batch_id"}],
        "order_by": [
            {"property": attribute_name},
            {"alias": "timestamp", "direction": "desc"},
        ],
    }


def _batch_timestamp_query():
    return {
        "select": [
            {
                "function": "max",
                "alias": "timestamp",
                "parameters": {"property": "inference_timestamp"},
            },
            {"property": "batch_id"},
        ],
        "group_by": [{"property": "batch_id"}],
        "order_by": [{"alias": "timestamp", "direction": "desc"}],
    }


def _label(model, label):
    model_name = model.display_name
    if len(model_name) <= 20:
        short_model_name = model_name
    else:
        short_model_name = model_name[0:17] + "..."
    return (
        "(" + str(model.version_sequence_num) + ")__" + short_model_name + "__" + label
    )


def _alpha_color(color):
    return (color[0], color[1], color[2], style.alpha)


def _drift_color_palettes(num_attributes, num_models):
    max_attributes = len(style.multi_categorical_palette_names)
    if num_attributes > max_attributes:
        raise UserValueError(
            f"Too many attributes to display. The maximum is {max_attributes}, you wanted to see {num_attributes}"
        )

    color_palettes = []
    color_palette_names = style.multi_categorical_palette_names[0:num_attributes]
    for i in range(num_models):
        color_palettes.append([])
    for color_palette_name in color_palette_names:
        base_palette = sns.color_palette(
            palette=color_palette_name, n_colors=num_models
        )
        for i, base_palette_color in enumerate(base_palette):
            color_palettes[i].append(_alpha_color(base_palette_color))

    return color_palettes
