#
# Copyright (c) Daniel Sheffield 2023
#
# All rights reserved
#
# THIS SOFTWARE IS PROVIDED AS IS WITHOUT WARRANTY
from io import StringIO
from datetime import date, datetime
from queue import Queue
from bottle import (
    FormsDict,
    HTTPError,
    template,
)
import matplotlib.pyplot as plt
import matplotlib
import numpy as np
from pandas import DataFrame
import seaborn as sns
from psycopg import Connection
from psycopg.connection import TupleRow

from . import ALL_UNITS, BOOLEAN, PARAMS
from ..data.QueryManager import (
    display_mapper,
    QueryManager,
)
from ..data.filter import (
    get_filter,
    get_query_param,
)
from .form import(
    get_form,
)

matplotlib.use('agg')

plot_style = {
    "lines.color": "#ffffff",
    "patch.edgecolor": "#ffffff",
    "text.color": "#ffffff",
    "axes.facecolor": "#7f7f7f",
    "axes.edgecolor": "#ffffff",
    "axes.labelcolor": "#ffffff",
    "xtick.color": "#ffffff",
    "ytick.color": "#ffffff",
    "grid.color": "#ffffff",
    "figure.facecolor": "#7f7f7f",
    "figure.edgecolor": "#7f7f7f",
    "savefig.facecolor": "#7f7f7f",
    "savefig.edgecolor": "#7f7f7f",
}

def get_data(query_manager: QueryManager, unit=None, **kwargs) -> DataFrame:
    d = DataFrame(query_manager.get_historic_prices_data(unit, **kwargs))
    if d.empty:
        return d
    d['ts_month'] = d['ts_raw'].apply(lambda x: date(x.date().year, x.date().month,1))
    d[['price','quantity']] = d[['price','quantity']].apply(
        lambda y: y.apply(lambda x: x and float(x)),
    )
    return d

def abort(code, text):
    raise HTTPError(code, text)

def trend(queue: Queue, conn: Connection[TupleRow], path: str, query: FormsDict):
    for item in trend_internal(conn, path, query):
        queue.put(item, block=True)
    queue.put(None)

def trend_internal(conn: Connection[TupleRow], path: str, query: FormsDict):
    progress = {
        'stage': None,
        'percent': None,
    }
    action = path.split('/')[-1]
    organic = BOOLEAN.get(query.organic, None)
    _filter = get_filter(query, allow=PARAMS)
    yield template("trend", start=True)
    try:
        with conn.cursor() as cur:
            query_manager = QueryManager(cur, display_mapper)
            fields = {
                k: get_query_param(*_filter[k])
                for k in sorted(_filter) if k not in ('organic', 'unit') and _filter[k]
            }
            unit = fields['unit'] = query.unit or 'kg'
            fields['organic'] = BOOLEAN.get(query.organic, None)
            if unit and unit not in ALL_UNITS:
                abort(400, f"Unsupported unit: {unit}")

            progress.update({ "stage": "Querying database", "percent": "10"})
            yield template("done") + template("progress", **progress)
            data = get_data(query_manager, **fields)
            
            if data.empty:
                abort(404, f"No data.")
            
            progress.update({ "stage": "Preparing data", "percent": "30"})
            yield template("done") + template("progress", **progress)
            
            in_chart = data['$/unit'].apply(lambda x: (x or False) and True)
            data = data[in_chart]
            
            pivot = data.pivot_table(index=['ts_raw',], columns=['product',], values=['$/unit'], aggfunc='mean')
            pivot.columns = pivot.columns.droplevel()
            sns.set_theme(style='darkgrid', palette='pastel', context="talk")
            plt.style.use("dark_background")
            plt.rcParams.update(plot_style)
            plt.rcParams.update({"grid.linewidth":0.2, "grid.alpha":0.5})
            plt.figure(figsize=[16, 9], layout="tight")
            xlabel='Time'
            ylabel=f'$ / {unit}'
            if pivot.columns.size > 50:
                ax = sns.scatterplot(data=pivot, markers=True)
            else:
                ax = sns.lineplot(data=pivot, markers=True)
                legend = plt.figlegend(
                    loc='upper center', ncol=6,
                    title_fontsize="14", fontsize="12", labelcolor='#ffffff',
                    framealpha=0.5
                )
                legend.set_title(title="Products")
            ax.legend().set_visible(False)

            ax.set_xlabel(xlabel, fontsize="14")
            ax.set_ylabel(ylabel, fontsize="14")
            ax.axes.tick_params(labelsize="12", which='both')
            for _, spine in ax.spines.items():
                spine.set_color('#ffffff')
            
            progress.update({ "stage": "Rendering chart", "percent": "50"})
            yield template("done") + template("progress", **progress)
            
            f = StringIO()
            plt.savefig(f, format='svg')
            progress.update({ "stage": "Done", "percent": "100" })
            yield template("done") + template("progress", **progress)
            
            form = get_form(action, 'post', _filter, organic, data)
            
            yield template("trend", end=True, form=form, svg=f.getvalue())

    except HTTPError as e:
        if 'data' not in locals():
            data = DataFrame()
        if 'form' not in locals():
            form = get_form(action, 'post', _filter, organic, data)
        yield template("done") + template("trend", end=True, form=form, error=e.body)

    finally:
        conn.commit()

def volume(queue: Queue, conn: Connection[TupleRow], path: str, query: FormsDict):
    for item in volume_internal(conn, path, query):
        queue.put(item, block=True)
    queue.put(None)

def volume_internal(conn: Connection[TupleRow], path: str, query: FormsDict):
    progress = {
        'stage': None,
        'percent': None,
    }
    action = path.split('/')[-1]
    organic = BOOLEAN.get(query.organic, None)
    _filter = get_filter(query, allow=PARAMS)
    yield template("trend", start=True)
    try:
        with conn.cursor() as cur:
            query_manager = QueryManager(cur, display_mapper)
            fields = {
                k: get_query_param(*_filter[k])
                for k in sorted(_filter) if k not in ('organic', 'unit') and _filter[k]
            }
            unit = fields['unit'] = query.unit or 'kg'
            fields['organic'] = BOOLEAN.get(query.organic, None)
            if unit and unit not in ALL_UNITS:
                abort(400, f"Unsupported unit: {unit}")

            progress.update({ "stage": "Querying database", "percent": "10"})
            yield template("done") + template("progress", **progress)
            data = get_data(query_manager, **fields)

            if data.empty:
                abort(404, f"No data.")
            
            progress.update({ "stage": "Preparing data", "percent": "30"})
            yield template("done") + template("progress", **progress)
            
            now = datetime.now().date()
            prev_month = date(now.year,now.month-1,1) if now.month > 1 else date(now.year-1, 12, 1)
            data = data[data['ts_month'] == prev_month ]
            group = 'group'
            for g, _g in zip(
                ('category', 'group'),
                ('product', 'category')
            ):
                if g and len(data[g].unique()) != 1:
                    continue
                group = _g
                break

            pivot = data[~(data['quantity'].isnull())].groupby([group,])[['price', 'quantity']].sum()
            
            if pivot.empty:
                abort(404, f"No data.")

            sns.set_theme(style='darkgrid', palette='pastel', context="talk")
            plt.style.use("dark_background")
            plt.rcParams.update(plot_style)
            plt.rcParams.update({"grid.linewidth":0.2, "grid.alpha":0.5})
            svg = []
            for title, col, (pre, fmt, suf) in zip((
                f"Expenditure ${pivot['price'].sum():0.2f}",
                f"Quantity {pivot['quantity'].sum():0.1f} {unit}",
            ), (
                'price',
                'quantity'
            ), [
                ('$', "{0:0.2f}", ''),
                ('', "{0:0.1f}", f' {unit}')
            ]):
                plt.figure(figsize=[16, 9], layout="tight")
                ax = plt.axes()
                wedges, *_ = ax.pie(pivot[col].values, startangle=90)
                bbox_props = dict(boxstyle="square,pad=0.3", lw=0.72, fc='#5f5f5f', ec=plot_style["axes.edgecolor"])
                kw = dict(arrowprops=dict(arrowstyle="-"),
                    bbox=bbox_props, zorder=0, va="center")
                for i, p in enumerate(wedges):
                    ang = (p.theta2 - p.theta1)/2. + p.theta1
                    y = np.sin(np.deg2rad(ang))
                    x = np.cos(np.deg2rad(ang))
                    horizontalalignment = {-1: "right", 1: "left"}[int(np.sign(x))]
                    connectionstyle = f"angle,angleA=0,angleB={ang}"
                    kw["arrowprops"].update({"connectionstyle": connectionstyle})
                    label = pivot.index[i]
                    val = pivot.loc[label, col]
                    label = f"{label} {pre}{fmt.format(val)}{suf} ({val*100/pivot[col].sum():0.1f}%)"
                    ax.annotate(label, xy=(x, y), xytext=(1.35*np.sign(x), 1.4*y),
                                horizontalalignment=horizontalalignment, **kw)
                ax.set_title(title)
                xlabel=''
                ylabel=''
                ax.legend().set_visible(False)

                ax.set_xlabel(xlabel, fontsize="14")
                ax.set_ylabel(ylabel, fontsize="14")
                ax.axes.tick_params(labelsize="12", which='both')
            
                progress.update({ "stage": "Rendering chart", "percent": "50"})
                yield template("done") + template("progress", **progress)
                
                f = StringIO()
                plt.savefig(f, format='svg')
                svg.append(f.getvalue())

            progress.update({ "stage": "Done", "percent": "100" })
            yield template("done") + template("progress", **progress)
            
            form = get_form(action, 'post', _filter, organic, data)
            
            yield template("volume", end=True, form=form, svg=svg)

    except HTTPError as e:
        if 'data' not in locals():
            data = DataFrame()
        if 'form' not in locals():
            form = get_form(action, 'post', _filter, organic, data)
        yield template("done") + template("trend", end=True, form=form, error=e.body)

    finally:
        conn.commit()
