import gradio as gr
from fastai.data.all import *
from fastai.vision.all import *
import pandas as pd

learn_fc = load_learner('fc_model.pkl')

#train_std = pd.Series(data=[1.550274e+01, 1.574251e+01, 1.526607e+01, 1.550392e+01, 1.066174e+01, 1.960896e+06 ], 
#                      index=['Open', 'High', 'Low', 'Close', 'Adj Close', 'Volume' ])

prediction_horizon = 32
look_back = 64
feature_selected = ['Open', 'Volume']
label_selected = 'Open'

class MyDataset:
    """Great for big datasets and long sequencies"""
    def __init__(self, df, features, target, look_back=1, prediction_horizon=1): 
        self.df, self.look_back, self.prediction_horizon = df, look_back, prediction_horizon
        self.features, self.target = features, target
        train_mean = 0
        train_std = self.df.std()
        self.df = (self.df - train_mean) / train_std
    def __len__(self): 
        n = len(self.df)-self.look_back+1-self.prediction_horizon
        return n
    def new_empty(self): return self[0]
    def __getitem__(self, i):
        p = self.look_back
        q = self.prediction_horizon
        X = self.df.iloc[i:(i+p)][self.features].values
        y = self.df.iloc[(i+p):(i+p + q)][self.target].values
        return tensor(X), tensor(y)

def predict(df):
    df.set_index('Date', inplace=True)
    # df = (df - train_mean) / train_std
    test_ds = MyDataset(df, feature_selected, label_selected, look_back, prediction_horizon)
    pred,_,_ = learn_fc.predict(test_ds[len(test_ds)-1])
    fig = plt.figure()
    plt.plot(pred)
    return fig

title = "Predict time series"
description = "Predict next 32 values for stock time series"
# article="<p style='text-align: center'><a href='https://tmabraham.github.io/blog/gradio_hf_spaces_tutorial' target='_blank'>Blog post</a></p>"
examples = ['BA-test.csv', 'BA-test-2.csv', 'BA-test-3.csv']

iface = gr.Interface(fn=predict, inputs=gr.inputs.Timeseries(x=None, y=['Date','Open','High','Low','Close','Adj Close','Volume']), 
                     outputs=gr.outputs.Plot(type='matplotlib'), examples=examples, title=title, description=description)
iface.launch()