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

# qis
import qis.utils.struct_ops as sop
import qis.plots.utils as put
from qis.plots.utils import LegendStats, TrendLine, LastLabel
from qis.perfstats.desc_table import compute_desc_table, DescTableType


def plot_time_series(df: Union[pd.Series, pd.DataFrame],
                     linestyles: List[str] = None,
                     linewidth: float = 1.0,
                     x_date_freq: Optional[str] = 'YE',
                     date_format: str = '%d-%b-%y',
                     legend_title: str = None,
                     legend_loc: Optional[Union[str, bool]] = 'upper left',
                     last_label: LastLabel = LastLabel.NONE,
                     sort_by_value_stretch_factor: float = 1.0,
                     trend_line: Optional[TrendLine] = TrendLine.NONE,
                     trend_line_colors: List[str] = None,
                     legend_stats: LegendStats = LegendStats.AVG_LAST,
                     desc_table_type: DescTableType = DescTableType.NONE,
                     legend_labels: List[str] = None,
                     indices_for_shaded_areas: Dict[str, Tuple[int, int]] = None,
                     xlabel: str = None,
                     ylabel: str = None,
                     var_format: Optional[str] = '{:,.2f}',
                     markers: List[str] = False,
                     title: Union[str, bool] = None,
                     fontsize: int = 10,
                     markersize: int = None,
                     colors: List[str] = None,
                     x_limits: Tuple[Optional[float], Optional[float]] = None,
                     y_limits: Tuple[Optional[float], Optional[float]] = None,
                     is_log: bool = False,
                     ax: plt.Subplot = None,
                     **kwargs
                     ) -> Optional[plt.Figure]:

    data1 = df.copy()
    if isinstance(data1, pd.DataFrame):
        pass
    elif isinstance(data1, pd.Series):
        data1 = data1.to_frame()
    else:
        raise TypeError(f"unsuported data type {type(data1)}")
    columns = data1.columns

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

    if colors is None:
        colors = put.get_n_colors(n=len(columns), **kwargs)

    sns.lineplot(data=data1, palette=colors, dashes=False, markers=markers,
                 markersize=markersize,
                 linewidth=linewidth, ax=ax)

    # add tredlines
    if trend_line_colors is None:
        trend_line_colors = colors
    if trend_line in [TrendLine.ZERO_SHADOWS, TrendLine.ABOVE_ZERO_SHADOWS]:
        for column, color in zip(columns, trend_line_colors):
            x0 = data1.index[0]
            x1 = data1.index[-1]
            y0 = 0
            y1 = 0
            x = [x0, x1]
            y = [y0, y1]
            ax.plot(x, y,
                    color='black',
                    linestyle='-',
                    transform=ax.transData,
                    lw=linewidth)

            x_ = (data1.index - data1.index[0]).days
            slope = (y1 - y0) / (x_[-1] - x_[0])
            y_line = [slope * x + y0 for x in x_]
            y = data1[column]
            if trend_line == TrendLine.ZERO_SHADOWS:
                ax.fill_between(data1.index, y, y_line, where=y_line >= y,
                            facecolor=color, interpolate=True, alpha=0.2, lw=linewidth)
            else:
                ax.fill_between(data1.index, y, y_line, where=y_line <= y,
                                facecolor=color, interpolate=True, alpha=0.2, lw=linewidth)

    elif trend_line in [TrendLine.AVERAGE, TrendLine.AVERAGE_SHADOWS]:
        for column, color in zip(columns, trend_line_colors):
            data2 = data1[column].dropna()  # exclude nans from showing the average lines
            if data2.empty:  # skip this columns from ere
                continue
            else:
                average = np.nanmean(data2.to_numpy())

            x0 = data2.index[0]
            x1 = data2.index[-1]
            y0 = average
            y1 = average
            x = [x0, x1]
            y = [y0, y1]

            ax.plot(x, y,
                    color=color,
                    linestyle=':',
                    transform=ax.transData,
                    linewidth=linewidth)

            if trend_line == TrendLine.AVERAGE_SHADOWS:
                x_ = (data2.index - data2.index[0]).days
                slope = (y1 - y0) / (x_[-1] - x_[0])
                y_line = [slope * x + y0 for x in x_]
                y = data2.to_numpy()
                ax.fill_between(data2.index, y, y_line, where=y_line >= y,
                                facecolor=color, interpolate=True, alpha=0.2, lw=linewidth)

    elif trend_line in [TrendLine.TREND_LINE, TrendLine.TREND_LINE_SHADOWS]:
        for column, color in zip(columns, trend_line_colors):
            y = data1[column].dropna()
            if not y.empty:
                x0 = y.first_valid_index()
                x1 = y.index[-1]
                y0 = y[x0]
                y1 = y.iloc[-1]
                x = [x0, x1]
                y = [y0, y1]
                ax.plot(x, y,
                        color=color,
                        linestyle='--',
                        transform=ax.transData,
                        linewidth=linewidth)

                if trend_line == TrendLine.TREND_LINE_SHADOWS:
                    x_ = (data1.index - data1.index[0]).days
                    slope = (y1 - y0) / (x_[-1] - x_[0])
                    y_line = [slope * x + y0 for x in x_]
                    y = data1[column]
                    ax.fill_between(data1.index, y, y_line, where=y_line >= y,
                                    facecolor=color, interpolate=True, alpha=0.2, lw=linewidth)

    # add last labels
    if last_label in [LastLabel.AVERAGE_VALUE, LastLabel.AVERAGE_VALUE_SORTED]:
        average_dict = {}
        for column, color in zip(columns, colors):
            data2 = data1[column].dropna()  # exclude nans from showing the average lines

            if data2.empty:  # skip this columns from ere
                continue
            else:
                average = np.nanmean(data2.to_numpy())

            if var_format is not None:
                average_str = var_format.format(average)
            else:
                average_str = average
            y1 = average
            variable_label = f"{column}, average={average_str}"
            average_dict.update({y1: [variable_label, color]})

        x1 = data1.index[-1]
        ymin, ymax = ax.get_ylim()
        mid = sort_by_value_stretch_factor * (ymax - ymin)
        if last_label == LastLabel.AVERAGE_VALUE_SORTED:
            pivot_dict = sorted(average_dict)
            locs = np.linspace(pivot_dict[0], sort_by_value_stretch_factor*mid, len(pivot_dict), endpoint=True)
        else:
            pivot_dict = average_dict
            locs = [dict for dict in pivot_dict]

        for key, loc in zip(pivot_dict, locs):
            ax.annotate(average_dict[key][0],
                        xy=(x1, key), xytext=(x1, loc),
                        fontsize=fontsize, weight='normal', color=average_dict[key][1],
                        textcoords='data', ha='left', va='bottom',
                        bbox={'boxstyle': 'round,pad=0.5', 'fc': average_dict[key][1], 'alpha': 0.1},
                        arrowprops=dict(arrowstyle='->', connectionstyle='arc3,rad=0'))

    elif last_label in [LastLabel.LAST_VALUE, LastLabel.LAST_VALUE_SORTED]:
        last_dict = {}
        for column, color in zip(columns, colors):
            y = data1[column].dropna()
            if last_label == LastLabel.LAST_VALUE:
                if var_format is not None:
                    if len(y.index) > 0:
                        variable_str = var_format.format(y.iloc[-1])
                    else:
                        variable_str = 'nan'
                else:
                    variable_str = y.iloc[-1]
                variable_label = f"{column}, last = {variable_str}"

                if len(y.index) > 0:
                    y1 = y.iloc[-1]
                else:
                    y1 = np.nan
                last_dict.update({y1: [variable_label, color]})

        # plot dicts sorted by last value
        x1 = data1.index[-1]
        ymin, ymax = ax.get_ylim()
        mid = sort_by_value_stretch_factor * (ymax - ymin)

        if last_label == LastLabel.LAST_VALUE_SORTED:
            pivot_dict = sorted(last_dict)
            locs = np.linspace(pivot_dict[0], sort_by_value_stretch_factor*mid, len(pivot_dict), endpoint=True)
        else:
            pivot_dict = last_dict
            locs = [dict for dict in pivot_dict]

        for key, loc in zip(pivot_dict,locs):
            ax.annotate(last_dict[key][0],
                        xy=(x1, key), xytext=(x1, loc),
                        fontsize=fontsize, weight ='normal', color = last_dict[key][1],
                        textcoords='data', ha='left', va='bottom',
                        bbox={'boxstyle': 'round,pad=0.5', 'fc': last_dict[key][1], 'alpha': 0.1},
                        arrowprops=dict(arrowstyle='->', connectionstyle='arc3,rad=0'))

    if indices_for_shaded_areas is not None:
        if len(data1.columns) == 1:
            warnings.warn(f"cannot do indices_for_shaded_areas = {indices_for_shaded_areas} for dataframe with one columns ")
        else:
            for col, indxs in indices_for_shaded_areas.items():
                y0 = data1.iloc[:, indxs[0]]
                y1 = data1.iloc[:, indxs[1]]
                ax.fill_between(data1.index, y0, y1, where=y0 >= y1,
                                facecolor=col, alpha=0.2,
                                interpolate=True)

    if legend_loc is not None:
        if legend_labels is None:
            if desc_table_type != DescTableType.NONE:  # use get_legend_with_stats_table
                stats_table = compute_desc_table(df=data1,
                                                 desc_table_type=desc_table_type,
                                                 var_format=var_format)
                put.set_legend_with_stats_table(stats_table=stats_table,
                                                ax=ax,
                                                colors=colors,
                                                legend_loc=legend_loc,
                                                fontsize=fontsize,
                                                **kwargs)
            else:  # generate legend_labels
                legend_labels = put.get_legend_lines(data=data1, legend_stats=legend_stats, var_format=var_format)

        if legend_labels is not None:
            put.set_legend(ax=ax,
                           labels=legend_labels,
                           colors=colors,
                           legend_loc=legend_loc,
                           legend_title=legend_title,
                           fontsize=fontsize,
                           **kwargs)

    else:
        ax.legend().set_visible(False)

    if linestyles is not None:
        put.set_linestyles(ax=ax, linestyles=linestyles)

    if x_date_freq is not None and isinstance(data1.index, pd.DatetimeIndex):
        put.set_date_on_axis(data=data1,
                             ax=ax,
                             x_date_freq=x_date_freq,
                             date_format=date_format,
                             fontsize=fontsize,
                             **kwargs)

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

    if x_limits is not None:
        put.set_x_limits(ax=ax, x_limits=x_limits)

    if var_format is not None:
        put.set_ax_ticks_format(ax=ax, fontsize=fontsize, xvar_format=None, yvar_format=var_format, **kwargs)

    if isinstance(data1.index, pd.DatetimeIndex):
        put.set_ax_tick_labels(ax=ax, skip_y_axis=True, fontsize=fontsize, **kwargs)

    put.set_ax_xy_labels(ax=ax, xlabel=xlabel, ylabel=ylabel, fontsize=fontsize, **kwargs)
    put.set_spines(ax=ax, **kwargs)

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

    if is_log:
        ax.set_yscale('log')

    return fig


def plot_time_series_2ax(df1: Union[pd.Series, pd.DataFrame],
                         df2: Union[pd.Series, pd.DataFrame],
                         legend_loc: Optional[str] = 'upper left',
                         legend_stats: LegendStats = LegendStats.NONE,
                         legend_stats2: LegendStats = LegendStats.NONE,
                         title: Optional[str] = None,
                         var_format: str = '{:,.0f}',
                         var_format_yax2: str = '{:,.0f}',
                         ylabel1: str = None,
                         ylabel2: str = None,
                         x_date_freq: Union[str, None] = 'YE',
                         legend_labels: List[str] = None,
                         linestyles: List[str] = None,
                         linestyles_ax2: List[str] = None,
                         y_limits: Tuple[Optional[float], Optional[float]] = None,
                         y_limits_ax2: Tuple[Optional[float], Optional[float]] = None,
                         trend_line1: put.TrendLine = put.TrendLine.NONE,
                         trend_line2: put.TrendLine = put.TrendLine.NONE,
                         yvar_major_ticks1: np.ndarray = None,
                         yvar_major_ticks2: np.ndarray = None,
                         colors: List[str] = None,
                         fontsize: int = 10,
                         x_rotation: int = 90,
                         is_logs: Tuple[bool, bool] = (False, False),
                         ax: plt.Subplot = None,
                         **kwargs
                         ) -> plt.Figure:
    if ax is None:
        fig, ax = plt.subplots()
    else:
        fig = None

    if isinstance(df1, pd.Series):
        df1 = df1.to_frame()
    if isinstance(df2, pd.Series):
        df2 = df2.to_frame()

    ncols1 = len(df1.columns)
    ncols2 = len(df2.columns)
    if colors is None:
        colors = put.get_n_colors(n=ncols1 + ncols2, **kwargs)

    ax_twin = ax.twinx()
    plot_time_series(df=df1,
                     legend_loc=None,
                     colors=colors[:ncols1],
                     var_format=None,
                     linestyles=linestyles,
                     trend_line=trend_line1,
                     x_date_freq=x_date_freq,
                     is_log=is_logs[0],
                     y_limits=y_limits,
                     ylabel=ylabel1,
                     fontsize=fontsize,
                     ax=ax,
                     **kwargs)

    plot_time_series(df=df2,
                     legend_loc=None,
                     colors=colors[ncols1:],
                     var_format=None,
                     linestyles=linestyles_ax2,
                     trend_line=trend_line2,
                     x_date_freq=x_date_freq,
                     is_log=is_logs[1],
                     y_limits=y_limits_ax2,
                     ylabel=ylabel2,
                     fontsize=fontsize,
                     ax=ax_twin,
                     **kwargs)

    ax.tick_params(axis='x', which='both', bottom=False)

    put.set_ax_ticks_format(ax=ax, fontsize=fontsize, xvar_format=None, yvar_format=var_format, set_ticks=False,
                            yvar_major_ticks=yvar_major_ticks1, x_rotation=x_rotation, **kwargs)
    put.set_ax_ticks_format(ax=ax_twin, fontsize=fontsize, xvar_format=None, yvar_format=var_format_yax2, set_ticks=False,
                            yvar_major_ticks=yvar_major_ticks2, x_rotation=x_rotation, **kwargs)

    if legend_loc is not None:
        if legend_labels is None:
            df1.columns = [f"{x} (left)" for x in df1.columns]
            df2.columns = [f"{x} (right)" for x in df2.columns]
            legend_labels1 = put.get_legend_lines(data=df1,
                                                  legend_stats=legend_stats,
                                                  var_format=var_format)
            legend_labels2 = put.get_legend_lines(data=df2,
                                                  legend_stats=legend_stats2,
                                                  var_format=var_format_yax2)
            legend_labels = sop.to_flat_list(legend_labels1 + legend_labels2)
        put.set_legend(ax=ax,
                       labels=legend_labels,
                       colors=colors,
                       fontsize=fontsize,
                       legend_loc=legend_loc,
                       **kwargs)

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

    return fig


class LocalTests(Enum):
    PRICES = 1
    PRICES_2AX = 2


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().dropna()

    if local_test == LocalTests.PRICES:
        fig, axs = plt.subplots(2, 1, figsize=(8, 6), tight_layout=True)
        global_kwargs = {'fontsize': 8,
                         'linewidth': 0.5,
                         'weight': 'normal',
                         'markersize': 1}
        plot_time_series(df=prices,
                         legend_stats=put.LegendStats.AVG_LAST,
                         last_label=LastLabel.AVERAGE_VALUE_SORTED,
                         trend_line=TrendLine.AVERAGE_SHADOWS,
                         ax=axs[0],
                         **global_kwargs)
        plot_time_series(df=prices,
                         legend_stats=put.LegendStats.AVG_LAST,
                         last_label=LastLabel.LAST_VALUE,
                         trend_line=TrendLine.AVERAGE_SHADOWS,
                         ax=axs[1],
                         **global_kwargs)

    elif local_test == LocalTests.PRICES_2AX:
        fig, ax = plt.subplots(1, 1, figsize=(4, 3), tight_layout=True)
        global_kwargs = {'fontsize': 6,
                         'linewidth': 0.5,
                         'weight': 'normal',
                         'markersize': 1}

        plot_time_series_2ax(df1=prices.iloc[:, -1],
                             df2=prices.iloc[:, :-1],
                             legend_stats=put.LegendStats.AVG_LAST,
                             var_format_yax2='{:.0f}',
                             ax=ax,
                             **global_kwargs)

    plt.show()


if __name__ == '__main__':

    run_local_test(local_test=LocalTests.PRICES)
