import gradio as gr
import pandas as pd

from gluonts.dataset.pandas import PandasDataset
from gluonts.dataset.split import split
from gluonts.torch.model.deepar import DeepAREstimator
from gluonts.torch.distributions import (
    NegativeBinomialOutput,
    StudentTOutput,
    NormalOutput,
)
from gluonts.evaluation import Evaluator, make_evaluation_predictions

from make_plot import plot_forecast, plot_train_test


def offset_calculation(prediction_length, rolling_windows, length):
    row_offset = -1 * prediction_length * rolling_windows
    if abs(row_offset) > 0.95 * length:
        raise gr.Error("Reduce prediction_length * rolling_windows")
    return row_offset


def preprocess(
    input_data,
    prediction_length,
    rolling_windows,
    progress=gr.Progress(track_tqdm=True),
):
    df = pd.read_csv(input_data.name, index_col=0, parse_dates=True)
    df.sort_index(inplace=True)
    row_offset = offset_calculation(prediction_length, rolling_windows, len(df))
    return plot_train_test(df.iloc[:row_offset], df.iloc[row_offset:])


def train_and_forecast(
    input_data,
    prediction_length,
    rolling_windows,
    epochs,
    distribution,
    progress=gr.Progress(track_tqdm=True),
):
    if not input_data:
        raise gr.Error("Upload a file with the Upload button")
    try:
        df = pd.read_csv(input_data.name, index_col=0, parse_dates=True)
        df.sort_index(inplace=True)
    except AttributeError:
        raise gr.Error("Upload a file with the Upload button")

    row_offset = offset_calculation(prediction_length, rolling_windows, len(df))

    try:
        gluon_df = PandasDataset(df, target=df.columns[0])
    except TypeError:
        freq = pd.infer_freq(df.index[:3])
        date_range = pd.date_range(df.index[0], df.index[-1], freq=freq)
        new_df = df.reindex(date_range)
        gluon_df = PandasDataset(new_df, target=new_df.columns[0], freq=freq)

    training_data, test_gen = split(gluon_df, offset=row_offset)

    if distribution == "StudentT":
        distr_output = StudentTOutput()
    elif distribution == "Normal":
        distr_output = NormalOutput()
    else:
        distr_output = NegativeBinomialOutput()
    estimator = DeepAREstimator(
        distr_output=distr_output,
        prediction_length=prediction_length,
        freq=gluon_df.freq,
        trainer_kwargs=dict(max_epochs=epochs),
    )

    predictor = estimator.train(
        training_data=training_data,
    )

    test_data = test_gen.generate_instances(
        prediction_length=prediction_length, windows=rolling_windows
    )

    evaluator = Evaluator(num_workers=0)
    forecast_it, ts_it = make_evaluation_predictions(
        dataset=test_data.input, predictor=predictor
    )
    agg_metrics, _ = evaluator(ts_it, forecast_it)

    forecasts = list(predictor.predict(test_data.input))

    return plot_forecast(df, forecasts), agg_metrics


with gr.Blocks() as demo:
    gr.Markdown(
        """
    # How to use

    Upload a *univariate* csv with the where the first column contains your dates and the second column is your data for example:

    | ds    | y        | 
    |------------|---------------|
    | 2007-12-10 | 9.590761      |
    | 2007-12-11 | 8.519590      |
    | 2007-12-12 | 8.183677      |
    | 2007-12-13 | 8.072467      |
    | 2007-12-14 | 7.893572      |

    ## Steps

    1. Click **Upload** to upload your data and visualize it.
    2. Click **Run**
        - This app will then train an estimator and show its predictions as well as evaluation metrics.
    """
    )
    with gr.Accordion(label="Hyperparameters"):
        with gr.Row():
            prediction_length = gr.Number(
                value=12, label="Prediction Length", precision=0
            )
            windows = gr.Number(value=3, label="Number of Windows", precision=0)
            epochs = gr.Number(value=10, label="Number of Epochs", precision=0)
            distribution = gr.Radio(
                choices=["StudentT", "Negative Binomial", "Normal"],
                value="StudentT",
                label="Distribution",
            )

    with gr.Row(label="Dataset"):
        upload_btn = gr.UploadButton(label="Upload")
        train_btn = gr.Button(label="Train and Forecast")
    plot = gr.Plot()
    json = gr.JSON(label="Evaluation Metrics")

    upload_btn.upload(
        fn=preprocess,
        inputs=[upload_btn, prediction_length, windows],
        outputs=plot,
    )
    train_btn.click(
        fn=train_and_forecast,
        inputs=[upload_btn, prediction_length, windows, epochs, distribution],
        outputs=[plot, json],
    )

if __name__ == "__main__":
    demo.queue().launch()