"""
boxplot
"""
# packages
import warnings
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from typing import Union, List, Dict, Optional, Tuple
from enum import Enum

# qis
import qis.utils.df_cut as dfc
import qis.utils.df_melt as dfm
import qis.plots.utils as put


def plot_box(df: Union[pd.Series, pd.DataFrame],
             x: str = None,
             y: str = None,
             xlabel: Optional[Union[str, bool]] = True,
             ylabel: Optional[Union[str, bool]] = True,
             hue: Optional[str] = None,
             hue_order: List[str] = None,
             xvar_format: str = '{:.2f}',
             yvar_format: str = '{:.2f}',
             colors: Union[List[str], List[Tuple[float, float, float]]] = None,
             title: str = None,
             meanline: bool = False,
             showfliers: bool = False,  # outliers
             showmeans: bool = False,
             showmedians: bool = False,
             add_y_mean_labels: bool = False,
             add_xy_mean_labels: bool = False,
             add_xy_median_labels: bool = False,
             add_y_median_labels: bool = False,
             x_rotation: Optional[int] = 0,
             legend_loc: Optional[str] = 'upper right',
             fontsize: int = 10,
             linewidth: float = 1.0,
             y_lines: List[Dict] = None,
             labels: List[str] = None,
             original_index: Union[pd.DatetimeIndex, pd.Index] = None,
             x_date_freq: Union[str, None] = 'YE',
             date_format: str = '%b-%y',
             num_obs_for_ci: np.ndarray = None,
             add_zero_line: bool = False,
             add_mean_std_bound: bool = False,
             add_autocorr_std_bound: bool = False,
             continuous_x_col: str = None,
             y_limits: Tuple[Optional[float], Optional[float]] = None,
             whis: Optional[float] = 1.5,  # sns default
             add_hue_to_legend_title: bool = True,
             ax: plt.Subplot = None,
             **kwargs
             ) -> Optional[plt.Figure]:
    """
    plot boxplot of df[[x, y]]
    original_index use for melted df in term of original index
    """
    if x is None and y is None:
        if len(df.columns) == 2:
            x = df.columns[0]
            y = df.columns[1]
        else:
            raise ValueError(f"x and y must be defined for more than one columns")

    if ax is None:
        fig, ax = plt.subplots()
    else:
        fig = None

    if colors is not None:
        palette = colors
    else:
        if hue is None:
            palette = put.get_data_group_colors(df=df, x=x, y=y)
        else:
            palette = put.get_n_colors(n=len(df[x].unique()))

    sns.boxplot(x=x, y=y, data=df,
                hue=hue or x,
                hue_order=hue_order,
                palette=palette,
                linewidth=linewidth,
                meanline=meanline,
                showmeans=showmeans,
                showfliers=showfliers,
                medianprops={'visible': showmedians, 'color': 'black'},
                meanprops={'visible': showmeans or meanline, 'color': 'black'},
                whis=whis,
                ax=ax)

    if num_obs_for_ci is not None:
        se = 1.96 / np.sqrt(num_obs_for_ci)
        mean = np.zeros_like(se)
        xx = np.arange(0, len(num_obs_for_ci))
        ax.plot(xx, -se+mean, color='red', lw=linewidth, linestyle='--')
        ax.plot(xx, mean, color='red', lw=linewidth, linestyle='-')
        ax.plot(xx, se + mean, color='red', lw=linewidth, linestyle='--')

    elif add_mean_std_bound:
        mean = np.nanmean(df[y].to_numpy())
        stdev = np.nanstd(df[y].to_numpy())
        se = 1.96 * stdev / np.sqrt(len(df.index))
        y_lines = [dict(x=-se+mean, color='blue', linestyle='--'),
                   dict(x=mean, color='blue', linestyle='-'),
                   dict(x=se+mean, color='blue', linestyle='--')]
    elif add_autocorr_std_bound:
        se = 1.96 * 2.0 / np.sqrt(len(df.index))
        y_lines = [dict(x=-se, color='blue', linestyle='--'),
                   dict(x=0, color='blue', linestyle='-'),
                   dict(x=se, color='blue', linestyle='--')]

    elif add_zero_line:
        y_lines = [dict(x=0, color='black', linestyle='-')]

    if y_lines is not None:
        for y_line in y_lines:
            ax.axhline(y=y_line['x'], color=y_line['color'], linestyle=y_line['linestyle'], lw = 1)

    # x ticks
    if original_index is not None:
        dates = pd.Series(data=original_index, index=original_index)
        re_indexed_data, datalables = put.map_dates_index_to_str(data=dates,
                                                                 x_date_freq=x_date_freq,
                                                                 date_format=date_format)

        current_ticks = ax.get_xticks()
        xticks = list(np.linspace(current_ticks[0], current_ticks[-1], len(datalables)))
        put.set_ax_tick_labels(ax=ax, x_labels=datalables, x_rotation=x_rotation, xticks=xticks, fontsize=fontsize)

    else:

        # ax.set_xticklabels(ax.get_xticklabels(), fontsize=fontsize, rotation=x_rotation, minor=False)
        put.set_ax_tick_params(ax=ax, fontsize=fontsize)
        put.set_ax_tick_labels(ax=ax, x_rotation=x_rotation, fontsize=fontsize)

    if len(ax.get_yticks()) > 0:
        ax.set_yticks(ax.get_yticks())
        ax.set_yticklabels([yvar_format.format(x) for x in ax.get_yticks()],
                           fontsize=fontsize, minor=False)

    if add_y_mean_labels or add_xy_mean_labels:
        if hue is not None:
            df_by_x = df.sort_values(hue).groupby(hue, sort=False)
            x_col = x
        else:
            df_by_x = df.sort_values(x).groupby(x, sort=False)
            if continuous_x_col is None:
                raise ValueError(f"continuous_x_col must be given")
            x_col = continuous_x_col

        for x_idx, (key, g_data) in enumerate(df_by_x):
            y_val = g_data[y].mean()
            x_val = g_data[x_col].mean()
            if add_y_mean_labels:
                ax.text(x_idx, y_val, f"{yvar_format.format(y_val)}\n",
                        ha='center', va='center',
                        color='black', fontsize=fontsize)
            else:
                ax.text(x_idx, y_val, f"avg:\ny={yvar_format.format(y_val)}\nx={xvar_format.format(x_val)}",
                        ha='center', va='center',
                        color='black', fontsize=fontsize)

    elif add_y_median_labels or add_xy_median_labels:
        if continuous_x_col is None:
            raise ValueError(f"continuous_x_col must be given")
        df_by_x = df.groupby(x, sort=False)

        for x_idx, (key, g_data) in enumerate(df_by_x):
            y_val = g_data[y].median()
            x_val = g_data[continuous_x_col].median()
            if add_y_median_labels:
                ax.text(x_idx, y_val, f"{yvar_format.format(y_val)}\n\n",
                        ha='center', va='center',
                        color='black', fontsize=fontsize)
            else:
                ax.text(x_idx, y_val, f"med:\ny={yvar_format.format(y_val)}\nx={xvar_format.format(x_val)}",
                        ha='center', va='center',
                        color='black', fontsize=fontsize)

    # labels
    if xlabel is not None:
        if isinstance(xlabel, bool):
            if xlabel is True:
                xlabel = f"x = {x}"
            else:
                xlabel = None

    if ylabel is not None:
        if isinstance(ylabel, bool):
            if ylabel is True:
                ylabel = f"y = {y}"
            else:
                ylabel = None

    put.set_ax_xy_labels(ax=ax, xlabel=xlabel, ylabel=ylabel, fontsize=fontsize)

    if y_limits is not None:
        put.set_y_limits(ax=ax, y_limits=y_limits)

    # need to add line for legend with hues
    lines = None
    if labels is not None:
        lines = []
        for label, color in zip(labels, palette):
            lines.append((label, {'color': color}))

    if legend_loc is not None:
        put.set_legend(ax=ax,
                       legend_title=hue if add_hue_to_legend_title else None,
                       legend_loc=legend_loc,
                       fontsize=fontsize,
                       lines=lines,
                       **kwargs)

    if title is not None:
        put.set_title(ax=ax, title=title, fontsize=fontsize)
    put.set_spines(ax=ax)

    return fig


def df_boxplot_by_index(df: Union[pd.Series, pd.DataFrame],
                        index_var_name: str = 'date',
                        ylabel: str = 'returns',
                        xlabel: str = None,
                        show_ylabel: bool = True,
                        title: str = None,
                        colors: Optional[List[str]] = None,
                        ax: plt.Subplot = None,
                        **kwargs
                        ) -> Optional[plt.Figure]:
    """
    plot data frame by aggregating by columns
    xvar is the index
    appropriate when y data is uniform conditional n regimes
    """
    # make sure it is given
    df.index.name = index_var_name
    box_data = dfm.melt_df_by_columns(df=df, x_index_var_name=index_var_name, y_var_name=ylabel)
    if colors is None:
        colors = put.compute_heatmap_colors(a=np.nanmean(df.to_numpy(), axis=1))

    fig = plot_box(df=box_data,
                   x=index_var_name,
                   y=ylabel,
                   continuous_x_col=ylabel,
                   ylabel=ylabel if show_ylabel else None,
                   original_index=df.index,
                   colors=colors,
                   xlabel=xlabel or False,
                   title=title,
                   ax=ax,
                   **kwargs)
    return fig


def df_boxplot_by_columns(df: Union[pd.Series, pd.DataFrame],
                          hue_var_name: str = 'instruments',
                          y_var_name: str = 'weights',
                          ylabel: str = 'weights',
                          show_ylabel: bool = True,
                          title: str = None,
                          colors: Optional[List[str]] = None,
                          ax: plt.Subplot = None,
                          **kwargs
                          ) -> Optional[plt.Figure]:
    """
    plot boxplot using data by columns
    x-axis will be colum names
    """
    box_data = dfm.melt_df_by_columns(df=df, x_index_var_name=None, hue_var_name=hue_var_name, y_var_name=y_var_name)

    if colors is None:
        colors = put.compute_heatmap_colors(a=np.nanmean(df.to_numpy(), axis=0), **kwargs)

    fig = plot_box(df=box_data,
                   x=hue_var_name,
                   y=y_var_name,
                   continuous_x_col=y_var_name,
                   ylabel=ylabel if show_ylabel else None,
                   xlabel=hue_var_name,
                   original_index=None,
                   labels=df.columns.to_list(),
                   colors=colors,
                   title=title,
                   ax=ax,
                   **kwargs)

    return fig


def df_dict_boxplot_by_columns(dfs: Dict[str, Union[pd.Series, pd.DataFrame]],
                               hue_var_name: str = 'instruments',
                               y_var_name: str = 'weights',
                               ylabel: str = 'weights',
                               hue: str = 'Portfolio',
                               show_ylabel: bool = True,
                               title: str = None,
                               colors: Optional[List[str]] = None,
                               ax: plt.Subplot = None,
                               **kwargs
                               ) -> Optional[plt.Figure]:
    """
    dict keys are added as hue
    plot boxplot using data by columns
    x-axis will be colum names
    """
    box_datas = []
    for key, df in dfs.items():
        box_data = dfm.melt_df_by_columns(df=df, x_index_var_name=None, hue_var_name=hue_var_name, y_var_name=y_var_name)
        box_data[hue] = key
        box_datas.append(box_data)
    box_datas = pd.concat(box_datas, axis=0)
    if colors is None:
        colors = put.get_n_colors(n=len(dfs.keys()), **kwargs)

    fig = plot_box(df=box_datas,
                   x=hue_var_name,
                   y=y_var_name,
                   hue=hue,
                   labels=list(dfs.keys()),
                   continuous_x_col=y_var_name,
                   ylabel=ylabel if show_ylabel else None,
                   xlabel=hue_var_name,
                   original_index=None,
                   colors=colors,
                   title=title,
                   ax=ax,
                   **kwargs)

    return fig


def df_boxplot_by_hue_var(df: Union[pd.Series, pd.DataFrame],
                          hue_var_name: str = 'date',
                          x_index_var_name: str = None,
                          y_var_name: str = 'returns',
                          is_heatmap_colors: bool = True,
                          hue_order: List[str] = None,
                          title: str = None,
                          colors: List[str] = None,
                          ax: plt.Subplot = None,
                          **kwargs
                          ) -> Optional[plt.Figure]:
    """
    plot time series of data frame by melting columns
    hue_var_name is hue index
    """
    box_data = dfm.melt_df_by_columns(df=df,
                                      x_index_var_name=x_index_var_name,
                                      y_var_name=y_var_name,
                                      hue_var_name=hue_var_name,
                                      hue_order=hue_order)
    if colors is None:
        if is_heatmap_colors:
            a = np.nanmedian(df.to_numpy(), axis=0)
            colors = put.compute_heatmap_colors(a=np.where(np.isfinite(a), a, 0.0))
        else:
            n = len(df.columns) if isinstance(df, pd.DataFrame) else 1
            colors = put.get_n_colors(n=n, **kwargs)

    fig = plot_box(df=box_data,
                   x=x_index_var_name,
                   y=y_var_name,
                   hue=hue_var_name,
                   colors=colors,
                   title=title,
                   ax=ax,
                   **kwargs)

    return fig


def df_boxplot_by_classification_var(df: pd.DataFrame,
                                     x: str,  # classification variables
                                     y: str,
                                     num_buckets: int = 6,
                                     x_hue_name: str = 'x_regimes',
                                     hue: Optional[str] = None,
                                     hue_order: List[str] = None,
                                     title: str = None,
                                     is_add_xlabel: bool = True,
                                     xvar_format: str = '{:.2f}',
                                     yvar_format: str = '{:.2f}',
                                     showfliers: bool = False,
                                     showmeans: bool = False,
                                     meanline: bool = False,
                                     medianline: bool = True,
                                     add_xy_mean_labels: bool = False,
                                     add_xy_med_labels: bool = False,
                                     is_value_labels: bool = True,
                                     colors: Union[List[str], List[Tuple[float, float, float]]] = None,
                                     ax: plt.Subplot = None,
                                     **kwargs
                                     ) -> None:
    """
    use x as classification var and plod box plot relative to quatiles of x
    """
    df = df.dropna()
    scatter_data, _ = dfc.add_quantile_classification(df=df, x_column=x,
                                                      hue_name=x_hue_name,
                                                      num_buckets=num_buckets,
                                                      bins=None,
                                                      xvar_format=xvar_format,
                                                      is_value_labels=is_value_labels)
    if hue_order is not None:
        labels = hue_order
    else:
        labels = None

    plot_box(df=scatter_data,
             x=x_hue_name,
             y=y,
             hue=hue,
             continuous_x_col=x,
             xlabel=x_hue_name if is_add_xlabel else None,
             hue_order=hue_order,
             title=title,
             xvar_format=xvar_format,
             yvar_format=yvar_format,
             showfliers=showfliers,
             showmeans=showmeans,
             meanline=meanline,
             medianline=medianline,
             add_xy_mean_labels=add_xy_mean_labels,
             add_xy_median_labels=add_xy_med_labels,
             labels=labels,
             colors=colors,
             ax=ax,
             **kwargs)


def df_dict_boxplot_by_classification_var(data_dict: Dict[Tuple[str, str], pd.DataFrame],
                                          num_buckets: int = 6,
                                          x_hue_name: str = 'x_regimes',
                                          var_hue_name: str = 'vars',
                                          y_var_name: str = 'y_var',
                                          title: str = None,
                                          is_add_xlabel: bool = True,
                                          xvar_format: str = '{:.2f}',
                                          yvar_format: str = '{:.2f}',
                                          showfliers: bool = False,
                                          showmeans: bool = False,
                                          meanline: bool = False,
                                          medianline: bool = True,
                                          add_xy_mean_labels: bool = False,
                                          add_xy_med_labels: bool = False,
                                          colors: Union[List[str], List[Tuple[float, float, float]]] = None,
                                          is_value_labels: bool = False,
                                          is_add_last_value: bool = False,
                                          ax: plt.Subplot = None,
                                          **kwargs
                                          ) -> None:
    """
    get dict of dfs and add regime vars, key = (x_column, y_column)
    use x as classification var and plod box plot relative to quatiles of x
    is_value_labels will be same across different df and x
    """
    scatter_datas = None
    hue_order = []
    for key, df in data_dict.items():
        scatter_data, _ = dfc.add_quantile_classification(df=df, x_column=key[0],
                                                          hue_name=x_hue_name,
                                                          num_buckets=num_buckets,
                                                          bins=None,
                                                          xvar_format=xvar_format,
                                                          is_value_labels=is_value_labels,
                                                          **kwargs)
        scatter_data[var_hue_name] = key[1]  # y is id
        scatter_data = scatter_data.rename({key[1]: y_var_name}, axis=1)
        if scatter_datas is None:
            scatter_datas = scatter_data
        else:
            scatter_datas = pd.concat([scatter_datas, scatter_data], axis=0)
        hue_order.append(key[1])
    plot_box(df=scatter_datas,
             x=x_hue_name,
             y=y_var_name,
             hue=var_hue_name,
             xlabel=x_hue_name if is_add_xlabel else None,
             hue_order=hue_order,
             title=title,
             xvar_format=xvar_format,
             yvar_format=yvar_format,
             showfliers=showfliers,
             showmeans=showmeans,
             meanline=meanline,
             medianline=medianline,
             add_xy_mean_labels=add_xy_mean_labels,
             add_xy_median_labels=add_xy_med_labels,
             labels=hue_order,
             colors=colors,
             ax=ax,
             **kwargs)
    if is_add_last_value:
        scatter_datas = scatter_datas.sort_index().ffill()
        last = scatter_datas.iloc[-1, :]
        for idx, key in enumerate(data_dict.keys()):
            ax.annotate('Last',
                        xy=(last[x_hue_name], last[key[1]]), xytext=(1, 1),
                        textcoords='offset points', ha='left', va='bottom',
                        fontsize=12)
            ax.scatter(x=last[x_hue_name], y=last[key[1]], s=20)


class LocalTests(Enum):
    RETURNS_BOXPLOT = 1
    DF_BOXPLOT = 2
    DF_BOXPLOT_INDEX = 3
    DF_WEIGHTS = 4
    DF_DICT = 5


def run_local_test(local_test: LocalTests):
    """Run local tests for development and debugging purposes.

    These are integration tests that download real data and generate reports.
    Use for quick verification during development.
    """

    from qis.test_data import load_etf_data
    prices = load_etf_data()
    returns = prices.asfreq('QE', method='ffill').pct_change()

    if local_test == LocalTests.RETURNS_BOXPLOT:
        fig, ax = plt.subplots(1, 1, figsize=(8, 6))
        global_kwargs = dict(fontsize=8, linewidth=0.5, weight='normal', markersize=1)

        index_name = 'date'
        value_name = 'returns'

        box_data = dfm.melt_df_by_columns(df=returns, x_index_var_name=index_name, y_var_name=value_name)
        print(box_data)
        colors = put.compute_heatmap_colors(a=np.nanmean(returns.to_numpy(), axis=1))

        plot_box(df=box_data,
                 x=index_name,
                 y=value_name,
                 original_index=returns.index,
                 colors=colors,
                 xlabel=False,
                 ax=ax,
                 **global_kwargs)

    elif local_test == LocalTests.DF_BOXPLOT:
        # returns by the quantiles of the first variable
        var = returns.columns[0]
        df_boxplot_by_classification_var(df=returns[var].to_frame(), x=var, y=var)

    elif local_test == LocalTests.DF_BOXPLOT_INDEX:
        df_boxplot_by_index(df=returns)

    elif local_test == LocalTests.DF_WEIGHTS:
        df_boxplot_by_columns(df=prices,
                              hue_var_name='instruments',
                              y_var_name='weights',
                              ylabel='weights',
                              legend_loc=None,
                              showmedians=True,
                              add_y_median_labels=True)

    elif local_test == LocalTests.DF_DICT:
        dfs = {'alts': prices, 'bal': 0.5*prices}
        with sns.axes_style("darkgrid"):
            fig, ax = plt.subplots(1, 1, figsize=(8, 6))
            df_dict_boxplot_by_columns(dfs=dfs,
                                       hue_var_name='instruments',
                                       y_var_name='weights',
                                       ylabel='weights',
                                       legend_loc='upper center',
                                       showmedians=True,
                                       add_y_median_labels=True,
                                       ncols=2,
                                       ax=ax)

    plt.show()


if __name__ == '__main__':

    run_local_test(local_test=LocalTests.DF_DICT)
