## Imports
import pickle
import warnings
import streamlit as st
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import datetime

# import torch
from torch.distributions import Normal
from pytorch_forecasting import (
    TimeSeriesDataSet,
    TemporalFusionTransformer,
)

## Functions 
def raw_preds_to_df(raw,quantiles = None):
    """
    raw is output of model.predict with return_index=True
    quantiles can be provided like [0.1,0.5,0.9] to get interpretable quantiles
    in the output, time_idx is the first prediction time index (one step after knowledge cutoff)
    pred_idx the index of the predicted date i.e. time_idx + h - 1
    """
    index = raw[2]
    preds = raw[0].prediction
    dec_len = preds.shape[1]
    n_quantiles = preds.shape[-1]
    preds_df = pd.DataFrame(index.values.repeat(dec_len * n_quantiles, axis=0),columns=index.columns)
    preds_df = preds_df.assign(h=np.tile(np.repeat(np.arange(1,1+dec_len),n_quantiles),len(preds_df)//(dec_len*n_quantiles)))
    preds_df = preds_df.assign(q=np.tile(np.arange(n_quantiles),len(preds_df)//n_quantiles))
    preds_df = preds_df.assign(pred=preds.flatten().cpu().numpy())
    if quantiles is not None:
        preds_df['q'] = preds_df['q'].map({i:q for i,q in enumerate(quantiles)})

    preds_df['pred_idx'] = preds_df['time_idx'] + preds_df['h'] - 1
    return preds_df

def prepare_dataset(parameters, df, rain, temperature, datepicker):
    if rain != "Default":
        df["MTXWTH_Day_precip"] = rain_mapping[rain]
    
    df["MTXWTH_Temp_min"] = df["MTXWTH_Temp_min"] + temperature
    df["MTXWTH_Temp_max"] = df["MTXWTH_Temp_max"] + temperature

    lowerbound = datepicker - datetime.timedelta(days = 35) 
    upperbound = datepicker + datetime.timedelta(days = 30) 

    df = df.loc[(df["Date"]>lowerbound) & (df["Date"]<=upperbound)]
    
    df = TimeSeriesDataSet.from_parameters(parameters, df)
    return df.to_dataloader(train=False, batch_size=256,num_workers = 0)

def predict(model, dataloader): 
    return model.predict(dataloader, mode="raw", return_x=True, return_index=True)
    
## Initiate Data
with open('data/parameters.pkl', 'rb') as f:
    parameters = pickle.load(f)
model = TemporalFusionTransformer.load_from_checkpoint('model/tft_check.ckpt', map_location=torch.device('cpu'))

df = pd.read_pickle('data/test_data.pkl')
df = df.loc[(df["Branch"] == 15) & (df["Group"].isin(["6","7","4","1"]))]

rain_mapping = {
    "Yes" : 1,
    "No" : 0
}

# Start App
st.title("Temporal Fusion Transformers for Interpretable Multi-horizon Time Series Forecasting")

st.markdown(body = """
    ### Abstract
    Multi-horizon forecasting often contains a complex mix of inputs – including
    static (i.e. time-invariant) covariates, known future inputs, and other exogenous
    time series that are only observed in the past – without any prior information
    on how they interact with the target. Several deep learning methods have been
    proposed, but they are typically ‘black-box’ models which do not shed light on
    how they use the full range of inputs present in practical scenarios. In this pa-
    per, we introduce the Temporal Fusion Transformer (TFT) – a novel attention-
    based architecture which combines high-performance multi-horizon forecasting
    with interpretable insights into temporal dynamics. To learn temporal rela-
    tionships at different scales, TFT uses recurrent layers for local processing and
    interpretable self-attention layers for long-term dependencies. TFT utilizes spe-
    cialized components to select relevant features and a series of gating layers to
    suppress unnecessary components, enabling high performance in a wide range of
    scenarios. On a variety of real-world datasets, we demonstrate significant per-
    formance improvements over existing benchmarks, and showcase three practical
    interpretability use cases of TFT.
""")

rain = st.radio("Rain Indicator", ('Default', 'Yes', 'No'))

temperature = st.slider('Change in Temperature', min_value=-10, max_value=+10, value=0, step=0.25)

datepicker = st.date_input("Start of Forecast", datetime.date(2022, 12, 24), min_value=datetime.date(2022, 6, 26) + datetime.timedelta(days = 35), max_value=datetime.date(2023, 6, 26) - datetime.timedelta(days = 30))

arr = np.random.normal(1, 1, size=100)
fig, ax = plt.subplots()
ax.hist(arr, bins=20)

st.pyplot(fig)

st.button("Forecast Sales", type="primary") #on_click=None,

# %%
preds = raw_preds_to_df(out, quantiles = None)

preds = preds.merge(data_selected[['time_idx','Group','Branch','sales','weight','Date','MTXWTH_Day_precip','MTXWTH_Temp_max','MTXWTH_Temp_min']],how='left',left_on=['pred_idx','Group','Branch'],right_on=['time_idx','Group','Branch'])
preds.rename(columns={'time_idx_x':'time_idx'},inplace=True)
preds.drop(columns=['time_idx_y'],inplace=True)