"""
plot aggregate report on pandas data and plot time series of columns
each time series plot = data[column]
 # = # columns
"""
# packages
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from typing import Tuple, Optional, Dict, List
from enum import Enum

# qis
import qis as qis
import qis.utils.dates as da
import qis.utils.df_agg as dfa
import qis.utils.df_groups as dfg
import qis.utils.df_ops as dfo
from qis.perfstats.desc_table import DescTableType
from qis.plots import histogram as pdf
from qis.plots.reports.utils import get_summary_table_fig, set_x_date_freq, ReportType
from qis.plots.derived.data_timeseries import plot_data_timeseries


def econ_data_report(data: pd.DataFrame,
                     complimentary_datas: Dict[str, pd.DataFrame] = None,
                     group_data: pd.Series = None,
                     descriptive_data: pd.DataFrame = None,
                     inst_names: pd.Series = None,
                     feature_name: str = None,
                     report_type: ReportType = ReportType.SingleTimeSeries,
                     is_add_summary_table: bool = True,
                     file_name: str = 'econ_data',
                     var_format: Optional[str] = None,
                     is_price_data: bool = False,
                     desc_table_type: DescTableType = DescTableType.WITH_KURTOSIS,
                     x_min_max_quantiles: Tuple[Optional[float], Optional[float]] = (0.005, 0.995),
                     background: str = 'darkgrid',
                     **kwargs
                     ) -> List[plt.Figure]:
    """
    report for economic data
    """
    local_kwargs = {'fontsize': 6,
                    'linewidth': 0.5,
                    'weight': 'normal',
                    'framealpha': 0.75}
    qis.update_kwargs(local_kwargs, kwargs)

    figsize_table = (8.0, 2.9)  # table for 1 column
    figsize = (14, 10)

    if var_format is None:
        if is_price_data:
            var_format = '{:,.2f}'
        else:
            var_format = '{:.2%}'

    # get summary table
    price_start_dates = {}
    price_end_dates = {}
    insample_prices = {}
    for inst in data.columns:
        insample_price = dfo.drop_first_nan_data(df=data[inst]).dropna()

        if insample_price.empty:
            bbg_start = np.nan
            bbg_finish = np.nan
        else:
            bbg_start = insample_price.index[0]
            bbg_finish = insample_price.index[-1]

        price_start_dates.update({inst: bbg_start})
        price_end_dates.update({inst: bbg_finish})
        insample_prices.update({inst: insample_price})
    price_start_dates = pd.Series(price_start_dates).rename('StartDate')
    price_end_dates = pd.Series(price_end_dates).rename('EndDate')
    price_start_dates = pd.concat([price_start_dates, price_end_dates], axis=1)

    figs = []
    if is_add_summary_table:
        fig = get_summary_table_fig(data_start_dates=price_start_dates,
                                    descriptive_data=descriptive_data,
                                    figsize=figsize_table,
                                    **kwargs)
        figs.append(fig)

    group_means = None
    if group_data is not None:
        group_means = dfg.agg_df_by_groups(df=data,
                                           group_data=group_data,
                                           agg_func=dfa.nanmean,
                                           total_column=None)
    figure_num = 0
    for inst, insample_price in insample_prices.items():
        if np.all(insample_price.isnull()):  # skip for all nans
            continue
        figure_num += 1
        figure_caption = f"Figure {figure_num}. {inst}"
        if inst_names is not None:
            figure_caption += f", {inst_names[inst]}"

        if descriptive_data is not None:
            for k, v in descriptive_data.loc[inst].to_dict().items():
                figure_caption += f", {k}: {v}"

        local_kwargs = set_x_date_freq(data=insample_price, kwargs=kwargs)

        # change names
        if feature_name is None:
            inst_name = f"{inst} "
        else:
            insample_price = insample_price.rename(feature_name)
            inst_name = ""

        if complimentary_datas is not None:
            for name, data in complimentary_datas.items():
                if inst in data.columns:
                    comp_name = f"{inst_name}{name}"
                    insample_price = pd.concat([insample_price, data[inst].rename(comp_name)], axis=1)

        if group_means is not None:
            inst_group = group_data[inst]
            insample_price = pd.concat([insample_price, group_means[inst_group].rename(f"{inst_group} ac avg")], axis=1)

        if report_type == ReportType.SingleTimeSeries:
            with sns.axes_style(background):
                fig, ax = plt.subplots(1, 1, figsize=figsize, tight_layout=True)
                plot_data_timeseries(data=insample_price,
                                     ax=ax,
                                     title=f"In-sample period = ",
                                     is_price_data=is_price_data,
                                     var_format=var_format,
                                     **local_kwargs)

        elif report_type == ReportType.SingleTimeSeriesWithPDF:
            fig, axs = plt.subplots(1, 2, figsize=figsize, tight_layout=True)
            plot_data_timeseries(data=insample_price,
                                 ax=axs[0],
                                 title=f"In-sample period = ",
                                 is_price_data=is_price_data,
                                 var_format=var_format,
                                 **local_kwargs)

            pdf.plot_histogram(df=insample_price,
                               xvar_format=var_format,
                               x_min_max_quantiles=x_min_max_quantiles,
                               add_last_value=False,
                               desc_table_type=desc_table_type,
                               title=f"Distribution: {da.get_time_period(df=data).to_str()}",
                               bbox_to_anchor=None,
                               ax=axs[1],
                               **kwargs)

        else:
            raise TypeError(f"report_type={report_type} is not implemented")

        qis.set_suptitle(fig, title=f"{file_name}_{figure_num}_{inst}")
        figs.append(fig)

    return figs


class LocalTests(Enum):
    UPDATE_DATA = 1
    RUN_REPORT = 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.
    """

    import matplotlib.pyplot as plt
    if local_test == LocalTests.UPDATE_DATA:
        from bbg_fetch import fetch_field_timeseries_per_tickers
        vix_tickers = ['VIX1D Index', 'VIX9D Index', 'VIX Index', 'VIX3M Index', 'VIX6M Index', 'VIX1Y Index']
        vols = 0.01 * fetch_field_timeseries_per_tickers(tickers=vix_tickers, field='PX_LAST', CshAdjNormal=True)
        print(vols)
        qis.save_df_to_csv(df=vols, file_name='vix_indices')

    elif local_test == LocalTests.RUN_REPORT:
        df = qis.load_df_from_csv(file_name='vix_indices')
        figs = econ_data_report(data=df)
        qis.save_figs_to_pdf(figs, file_name='vix_indices')

    plt.show()


if __name__ == '__main__':

    run_local_test(local_test=LocalTests.RUN_REPORT)
