__author__ = "Lukas Leufen, Felix Kleinert"
__date__ = '2019-12-17'

import os
import logging

import xarray as xr
import matplotlib
import seaborn as sns
import matplotlib.pyplot as plt

import dask.array as da
import dask
import dask.dataframe as ddf
logging.getLogger('matplotlib').setLevel(logging.WARNING)


def plot_monthly_summary(stations, data_path, name: str, window_lead_time, target_var, plot_folder="."):

    logging.debug("run plot_monthly_summary()")
    forecasts = None

    for station in stations:
        logging.debug(f"... preprocess station {station}")
        file_name = os.path.join(data_path, name % station)
        data = xr.open_dataarray(file_name)

        data_cnn = data.sel(type="CNN").squeeze()
        data_cnn.coords["ahead"].values = [f"{days}d" for days in data_cnn.coords["ahead"].values]

        data_orig = data.sel(type="orig", ahead=1).squeeze()
        data_orig.coords["ahead"] = "orig"

        data_concat = xr.concat([data_orig, data_cnn], dim="ahead")
        data_concat = data_concat.drop("type")

        data_concat.index.values = data_concat.index.values.astype("datetime64[M]").astype(int) %12 +1
        data_concat = data_concat.clip(min=0)

        forecasts = xr.concat([forecasts, data_concat], 'index') if forecasts is not None else data_concat

    forecasts = forecasts.to_dataset(name='values').to_dask_dataframe()
    logging.debug("... start plotting")
    ax = sns.boxplot(x='index', y='values',  hue='ahead', data=forecasts.compute(), whis=1.,
                     palette=[matplotlib.colors.cnames["green"]]+sns.color_palette("Blues_d", window_lead_time).as_hex(),
                     flierprops={'marker':'.', 'markersize':1},
                     showmeans=True, meanprops={'markersize':1, 'markeredgecolor':'k'})
    ax.set(xlabel='month', ylabel=f'{target_var}')
    plt.tight_layout()
    plot_path = os.path.join(os.path.abspath(plot_folder), 'test_monthly_box.pdf')
    logging.debug(f"... save plot to {plot_path}")
    plt.savefig(plot_path)
    plt.close('all')