
import argparse as ap
import matplotlib.pyplot as plt
import json
import pandas as pd

from bokeh.io import show as bshow
from bokeh.models import HoverTool
from bokeh.plotting import figure
from datetime import datetime as dt
from gluonts.dataset.common import Dataset, ListDataset
from gluonts.dataset.field_names import FieldName
from gluonts.dataset.repository.datasets import default_dataset_path, get_dataset, get_download_path
from gluonts.evaluation import Evaluator, make_evaluation_predictions
from gluonts.model.forecast import Forecast
from gluonts.model.deepar import DeepAREstimator
from gluonts.model.deepstate import DeepStateEstimator
from gluonts.model.predictor import Predictor
from gluonts.mx.distribution import NegativeBinomialOutput
from gluonts.mx.trainer import Trainer
from gluonts.mx.trainer.callback import Callback, TrainingHistory
from pandas import DataFrame
from pathlib import Path
from random import randint
from typing import Iterable, List, Optional, Union



default_model_path = get_download_path() / "models"


def create_arg_parser():
    default_quantiles = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
    parser = ap.ArgumentParser(
            description="Reproduce results presented in [Deep State Space Models for Time Series Forecasting]" + 
                        "(https://proceedings.neurips.cc/paper/2018/hash/5cf68969fb67aa6082363a6d4e6468e2-Abstract.html)" +
                        "over the `parts` dataset (a.k.a. `car_parts_without_missing` in the GluonTS package)."
        )
    datasetargs = parser.add_argument_group("Dataset")
    datasetargs.add_argument("--dataset_path", type=str,
                             help="Where the dataset is stored or should be saved.",
                             default=default_dataset_path)
    datasetargs.add_argument("--asis", type=bool, nargs="?",
                             help="The `parts` dataset is composed by 2674 time series. If this flag is " +
                                  "False the dataset will be preprocessed as explained in " +
                                  "[Forecasting the intermittent demand for slow-moving inventories: " +
                                  "A modelling approach]" +
                                  "(https://doi.org/10.1016/j.ijforecast.2011.03.009) and aligned with the "
                                  "Deep State Space Models paper; the resulting dataset will " +
                                  "retain 1046 time series.",
                             const=True,
                             default=False)
    datasetargs.add_argument("--regenerate", type=bool, nargs="?",
                             help="Whether to regenerate the dataset even if a local file is present. " + 
                                  "If this flag is False and the file is present, the dataset will not " +
                                  "be downloaded again.",
                             const=True,
                             default=False)
    estimatorargs = parser.add_argument_group("Estimator")
    estimatorargs.add_argument("--model_path", type=str,
                               help="Where the model is stored or should be saved. Any non-existing parent " +
                                    "along the path will be created automatically (if applicable). If --train " +
                                    " is specified, a timestamp directory will be created at the specified path.",
                               default=default_model_path)
    estimatorargs.add_argument("--train", type=bool, nargs="?",
                               help="If this flag is False, the training phase is skipped.",
                               const=True,
                               default=False)
    estimatorargs.add_argument("--estimator", type=str,
                               help="Which estimator to use.", choices=["deepar", "deepssm"],
                               default="deepar")
    estimatorargs.add_argument("--quantiles", type=float, nargs="*",
                               help="List of floats in [0, 1] with the quantile levels.",
                               default=default_quantiles)
    
    return parser


def get_parts_dataset(
    asis: Optional[bool] = True,
    path: Optional[Path] = default_dataset_path,
    regenerate: Optional[bool] = False
):
    parts = get_dataset("car_parts_without_missing",
                        path=path, regenerate=regenerate)

    if asis:
        return parts.train, parts.test

    samples = []
    for sample in iter(parts.test):
        target = sample['target']

        positive = len(target[target > 0])
        first_period = sum(target[0:15])
        last_period = sum(target[-15:])

        if positive >= 10 and first_period > 0 and last_period > 0:
            samples.append(sample)
  
    train = ListDataset(
        [
            {
                FieldName.ITEM_ID: sample['item_id'],
                FieldName.START: sample['start'],
                FieldName.TARGET: sample['target'][:-parts.metadata.prediction_length],
                FieldName.FEAT_STATIC_CAT: [idx]
            }
            for (sample, idx) in zip(samples, range(len(samples)))
        ],
        freq=parts.metadata.freq
    )
    test = ListDataset(
        [
            {
                FieldName.ITEM_ID: sample['item_id'],
                FieldName.START: sample['start'],
                FieldName.TARGET: sample['target'],
                FieldName.FEAT_STATIC_CAT: [idx]
            }
            for (sample, idx) in zip(samples, range(len(samples)))
        ],
        freq=parts.metadata.freq
    )
    return train, test


def plot_agg(
    forecasts: Iterable[Forecast],
    series: Optional[Iterable[Union[pd.DataFrame, pd.Series]]] = None,
    axis: Optional[int] = 0
):
    fdf = DataFrame(index=forecasts[0].index)
    if not (series is None):
        sdf = DataFrame(index=series[0].index)
        for (sample, fsample) in zip(series, forecasts):
            sdf = pd.concat([sdf, DataFrame(
                                      data=sample.values,
                                      index=sample.index,
                                      columns=[f'{fsample.item_id}']
                                  )],
                                  axis=1)
            fdf = pd.concat([fdf, DataFrame(
                                      data=fsample.median,
                                      index=fsample.index,
                                      columns=[f'{fsample.item_id}']
                                  )],
                                  axis=1)
    else:
        for fsample in forecasts:
            fdf = pd.concat([fdf, DataFrame(
                                      data=fsample.median,
                                      index=fsample.index,
                                      columns=[f'{fsample.item_id}']
                                  )],
                                  axis=1)
    agg = fdf.sum(axis=axis)
    print(f"Test Aggregation: {agg.values}")
    fg = figure(plot_height=768, plot_width=1024, x_axis_type="datetime")
    cat = fg.line(agg.index,
                  agg.values,
                  color="green", name="Forecast",
                  line_dash="dashed", legend_label="Forecast")
    fg.add_tools(HoverTool(renderers=[cat], 
                           tooltips=[('date', '@x{%Y-%m}'),
                                     ('name', '$name'),
                                     ('value', '@y')],
                           formatters={'@x': 'datetime'},
                           mode='vline'))
    if not (series is None):
        tagg = sdf.sum(axis=axis)
        tcat = fg.line(tagg.index,
                       tagg.values,
                       name="Qty", legend_label="Qty")
        fg.add_tools(HoverTool(renderers=[tcat],
                            tooltips=[('date', '@x{%Y-%m}'),
                                      ('name', '$name'),
                                      ('value', '@y')],
                            formatters={'@x': 'datetime'},
                            mode='vline'))
    fg.legend.location = "top_right"
    fg.legend.click_policy="mute"
    bshow(fg)


def plot_prob_forecasts(ts_entry, forecast_entry):
    plot_length = 150
    prediction_intervals = (50.0, 90.0)
    legend = ["observations", "median prediction"] + [f"{k}% prediction interval" for k in prediction_intervals][::-1]

    _, ax = plt.subplots(1, 1, figsize=(10, 7))
    ts_entry[-plot_length:].plot(ax=ax)  # plot the time series
    forecast_entry.plot(prediction_intervals=prediction_intervals, color='g')
    plt.grid(which="both")
    plt.legend(legend, loc="upper left")
    plt.show()


def train(data: Dataset,
          estimator: Optional[str] = 'deepar',
          callbacks: Optional[List[Callback]] = []
):
    if estimator == 'deepar':
        _estimator = DeepAREstimator(
            freq='M',
            prediction_length=12,
            trainer=Trainer(
                epochs=80,
                learning_rate=0.001,
                callbacks=callbacks
            ),
            num_layers=3,
            num_cells=40,
            dropoutcell_type='RNNZoneoutCell',
            use_feat_static_cat=True,
            cardinality=[1046],
            embedding_dimension=[1046],
            distr_output= NegativeBinomialOutput(),
            batch_size=64
        )
    if estimator == 'deepssm':
        _estimator = DeepStateEstimator(
            freq='M',
            prediction_length=12,
            trainer=Trainer(
                callbacks=callbacks
            ),
            num_layers=3,
            num_cells=40,
            use_feat_static_cat=True,
            cardinality=[1046],
            embedding_dimension=[1],
            batch_size=64
        )
    
    predictor = _estimator.train(data)
    return predictor


if __name__ == "__main__":
    parser = create_arg_parser()
    args = parser.parse_args()
    dpath = Path(args.dataset_path).resolve()
    mpath = Path(args.model_path).resolve()


    traind, testd = get_parts_dataset(args.asis,
                                      path=args.dataset_path,
                                      regenerate=args.regenerate)
    history = TrainingHistory()

    if args.train:
        start = dt.now()
        predictor=train(traind,
                        estimator = args.estimator.lower(),
                        callbacks=[history])
        end = dt.now()
        suffix = str(end - start).replace(":", "'")
        mpath = mpath / (start.isoformat().replace(":", "'") + "ET" + suffix)
        mpath.mkdir(parents=True, exist_ok=True)
        predictor.serialize(mpath)
        plt.plot(history.loss_history)
    else:
        predictor = Predictor.deserialize(mpath)

    forecast_it, ts_it = make_evaluation_predictions(
        dataset=testd,
        predictor=predictor,
        num_samples=200
    )
    forecasts = list(forecast_it)
    tss = list(ts_it)
    rand_idx = randint(0, len(forecasts) - 1)
    plot_prob_forecasts(tss[rand_idx], forecasts[rand_idx])

    evaluator = Evaluator(quantiles=args.quantiles)
    agg_metrics, item_metrics = evaluator(iter(tss), iter(forecasts))
    print(json.dumps(agg_metrics, indent=4))

    plot_agg(forecasts, series=tss, axis=1)