""" Basically small util to visualize the graphs """
from core.utils.Aggregators import Approach, Aggregators, PlotTypes
from core.db_style.basic_functions import PostgresDBUnit
from core.pandas_style.basic_functions import PandasUtil

import plotly
from typing import List, Tuple
import plotly.graph_objs as go

from core.utils.formatters import date_from_string, get_months, increment_month
from core.utils.validation import is_time_frame_valid


# TODO: add functionality for aggregated stats except balance (not only for periods)
def collect_date_intervals(start_date: str, end_date: str) -> dict:
    periods = {}
    start = date_from_string(start_date)
    end = date_from_string(end_date)
    # collect pairs year/month from start till end
    inner_interval = start
    while not inner_interval > end:
        key = "{}-{}".format(inner_interval.year, inner_interval.month)
        periods[key] = inner_interval
        inner_interval = increment_month(inner_interval.isoformat())
    return periods


class OnlineVisualizer:

    def __init__(self, approach, file_path=None):
        """ connects to a data source """
        if approach == Approach.DB:
            self.data_source = PostgresDBUnit()
        elif approach == Approach.DF:
            assert file_path is not None
            self.data_source = PandasUtil(file_path=file_path)

    def _prepare_cost_structure(self, start_date:str = None, end_date:str = None,
                                exclude_category: List[str] = None) -> dict:
        """
        Helper function to unite single and several subplots
        :param start_date: str start date
        :param end_date: str end date
        :param exclude_category: list of categories not to show on plot
        :return: dict {category: SUM cost}
        """
        cats = self.data_source.get_categories()
        cats = [x if x !=  "" else "Unknown" for x in cats]
        # TODO: fix non-vanishing categories
        if exclude_category is not None:
            for cat in cats:
                if cat in exclude_category:
                    cats.remove(cat)
        subsets = {}
        for cat in cats:
            subsets[cat] = self.data_source.generic_metrics_aggregator(Aggregators.SUM,
                                                                       start_date, end_date, cat)[0]

        return subsets

    def simple_pie_chart(self, start_date:str = None, end_date:str = None,
                         exclude_category: List[str] = None) -> go.Pie:
        """
        Possible combinations:
        1. total spending per category for time
        2. avg spending per category per time
        """
        subsets = self._prepare_cost_structure(start_date, end_date, exclude_category)
        return go.Pie(labels=list(subsets.keys()),
                      values=list(subsets.values()),
                      hoverinfo='label',
                      textinfo='value',
                      textfont=dict(size=14),
                      name='Cost structure')

    def simple_bar_chart(self, start_date: str = None, end_date: str = None,
                         category: str = None, aggregator: Aggregators = Aggregators.SUM,
                         formatter: str = "%Y-%m-%d") -> go.Bar:
        """
        Sums cost of a category for given period per month
        :param start_date: str start date of a main date frame
        :param end_date:  str end date of a main date frame
        :param category: str cost category name
        :param aggregator: Aggregator to build a bar chart for
        :param formatter: str formatter to parse datetime from
        :return: None
        """
        if is_time_frame_valid(start_date, end_date, formatter):
            periods = collect_date_intervals(start_date, end_date)
            # compute a sums for collected month:
            res = {}
            # TODO: check if the category is being validated
            for k, v in periods.items():
                res[k] = self.data_source.generic_metrics_aggregator(aggregator,
                                                                     v.isoformat(),
                                                                     v.isoformat(),
                                                                     category)[0]

        else:
            print("Chosen time frame {}/{} is nt valid".format(start_date, end_date))
            res = {}
        return go.Bar(
            x=list(res.keys()),
            y=list(res.values()),
            hoverinfo='label',
            name='{} cost'.format(category))

    def simple_balance_chart(self, start_date: str = None, end_date: str = None) -> go.Bar:
        """
        Draws balance over time
        :param start_date: str start date
        :param end_date: str end date
        :return: None
        """
        res = {}
        periods = collect_date_intervals(start_date, end_date)
        for k, v in periods.items():
            res[k] = self.data_source.balance(v.isoformat(), v.isoformat())

        return go.Bar(
            x=list(res.keys()),
            y=list(res.values()),
            hoverinfo='value',
            name="simple_balance")

    def accumulated_balance_chart(self, start_date: str = None, end_date: str = None) -> go.Scatter:
        """
        Accumulated balance over time
        :param start_date: str stat
        :param end_date: srt end
        :return: None
        """
        tmp = {}
        periods = collect_date_intervals(start_date, end_date)
        for k, v in periods.items():
            tmp[k] = self.data_source.balance(v.isoformat(), v.isoformat())

        res = tmp
        indexes = list(tmp.keys())
        for i in range(1, len(indexes)):
            res[indexes[i]] = res[indexes[i-1]] + tmp[indexes[i]]

        return go.Scatter(
            x=list(res.keys()),
            y=list(res.values()),
            hoverinfo='value',
            name="accumulated_balance")

    def show_pie_subplots(self, start_date: str = None, end_date: str = None,
                          exclude_category: List[str] = None, name: str = "Cost Structure") -> None:
        """
        Several charts for monthly intervals within given date range
        Hardcoded for 6 month
        :return: None, fig being rendered
        """
        # validate time frame is 6 month
        month = get_months(start_date, end_date)
        if month > 6:
            print("This option works for 6 month period only")
            raise Exception
        else:
            # collect all data for all sub-periods
            periods = collect_date_intervals(start_date, end_date)
            cost_structure = {}
            for k, v in periods.items():
                cost_structure[k] = self._prepare_cost_structure(v.isoformat(),
                                                                 v.isoformat(),
                                                                 exclude_category)
            # create a separate pie for each data set
            data = []
            annotations = []
            for i, v in enumerate(cost_structure.values()):
                data.append({
                    "values": list(v.values()),
                    "labels": list(v.keys()),
                    "type": "pie",
                    "domain": {"x": [0  + i * .15, .15 + i * .15],
                               "y": [0, -1]},
                    "hoverinfo": "label+value+name",
                    "name": periods[list(cost_structure.keys())[i]],
                    "textinfo": periods[list(cost_structure.keys())[i]]

                })
                annotations.append({
                    "font": {"size":14},
                    "text": "{}".format(periods[list(cost_structure.keys())[i]]),
                    "x": .055 + i * .15,
                    "y": .8,
                    "showarrow": False
                })
            layout = {'title': '{} for {}/{}'.format(name, start_date, end_date),
                      'showlegend': True,
                      "annotations": annotations}

            fig = {"data": data, "layout": layout}
            plotly.offline.plot(fig)

    @staticmethod
    def show_combined_plot(aggregators: List, start_date: str, end_date: str, name: str) -> None:
        """
        Plots several figures as one
        :return: None
        """
        layout = go.Layout(
            title="{} over {}/{}".format(name, start_date, end_date),
        )

        fig = go.Figure(data=aggregators, layout=layout)
        plotly.offline.plot(fig)
