#
# Copyright (c) Daniel Sheffield 2023
#
# All rights reserved
#
# THIS SOFTWARE IS PROVIDED AS IS WITHOUT WARRANTY
from datetime import date
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from app.data.QueryManager import QueryManager

def get_data(query_manager: QueryManager, unit=None, **kwargs) -> pd.DataFrame:
    d = pd.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 pivot_data(data: QueryManager):
    pivot = data.groupby(['ts_month','group',])['price', 'quantity'].sum()
    pivot = pivot.reset_index().set_index('group')
    return pivot

def pie(p: pd.DataFrame, col=None, title=None):
    ax = p.plot.pie(y=col, figsize=(5, 5))
    ax.get_legend().remove()
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.set_title(title)
    plt.show()

def line(pivot, ylabel=None, xlabel=None):
    ax = sns.lineplot(data=pivot, markers=True)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    plt.show()

def get_selection(
    query_manager: QueryManager,
    fields: dict[str, str],
    units: set[str],
    name: str
):
    options = query_manager.unique_suggestions(name, **fields)
    matches = '\t'.join(list(options))
    print(f'{name.title()} names: {matches}')
    while (len(options) >=1):
        v = fields[name] if name in fields else ''
        fields[name] = v + input(f"{name.title()}: {v}")
        options = query_manager.unique_suggestions(name, **fields)
        if name == 'unit':
            options = sorted(set(options) | units)

        if fields[name] == '':
            choice = ''
            break
        elif len(options) == 1:
            choice = options[0]
            break
        elif fields[name].lower() in map(lambda x: x.lower(), options):
            choice = next(
                filter(lambda x: x.lower() == fields[name].lower(), options)
            )
            break
        elif len(options) == 0 and name == 'unit':
            choice = fields[name]
            break
        matches = '\t'.join(options)
        print(f'Matches ({name}): {matches}')

    return choice

def get_input(query_manager: QueryManager, units: set[str]) -> dict[str, str]:
    fields = dict()
    for k in ('group', 'category', 'product', 'unit'):
        choice = get_selection(query_manager, fields, units, k)

        if k != 'unit':
            fields[k] = (fields[k] and choice) or ''
        else:
            fields[k] = choice
        if fields[k] == '':
            print(f'Ignoring {k} {choice} as it does not exist')
        else:
            print(f'Selected {k}: {fields[k]}')
        print()
    return fields

