#%%
from matplotlib.pyplot import title
import tensorflow as tf
from tensorflow import keras
from huggingface_hub import from_pretrained_keras
import pandas as pd
import matplotlib.pyplot as plt
import streamlit as st
from zipfile import ZipFile
import os

import warnings
warnings.filterwarnings("ignore")

uri = "https://storage.googleapis.com/tensorflow/tf-keras-datasets/jena_climate_2009_2016.csv.zip"
zip_path = keras.utils.get_file(origin=uri, fname="jena_climate_2009_2016.csv.zip")
zip_file = ZipFile(zip_path)
zip_file.extractall()
csv_path = "jena_climate_2009_2016.csv"
df = pd.read_csv(csv_path)

#%%

title = "Timeseries forecasting for weather prediction"

st.title('Timeseries forecasting for weather prediction')

st.write("Demonstrates how to do timeseries forecasting using a [LSTM model.](https://keras.io/api/layers/recurrent_layers/lstm/#lstm-class)This space demonstration is forecasting for weather prediction. *n* observation is selected from validation dataset." )
st.write("Keras example authors: [Prabhanshu Attri, Yashika Sharma, Kristi Takach, Falak Shah](https://keras.io/examples/timeseries/timeseries_weather_forecasting/)")


# %% model 

titles = [
    "Pressure",
    "Temperature",
    "Temperature in Kelvin",
    "Temperature (dew point)",
    "Relative Humidity",
    "Saturation vapor pressure",
    "Vapor pressure",
    "Vapor pressure deficit",
    "Specific humidity",
    "Water vapor concentration",
    "Airtight",
    "Wind speed",
    "Maximum wind speed",
    "Wind direction in degrees",
]

feature_keys = [
    "p (mbar)",
    "T (degC)",
    "Tpot (K)",
    "Tdew (degC)",
    "rh (%)",
    "VPmax (mbar)",
    "VPact (mbar)",
    "VPdef (mbar)",
    "sh (g/kg)",
    "H2OC (mmol/mol)",
    "rho (g/m**3)",
    "wv (m/s)",
    "max. wv (m/s)",
    "wd (deg)",
]

date_time_key = "Date Time"
split_fraction = 0.715
train_split = int(split_fraction * int(df.shape[0]))
step = 6

past = 720
future = 72
learning_rate = 0.001
batch_size = 256
epochs = 10


def normalize(data, train_split):
    data_mean = data[:train_split].mean(axis=0)
    data_std = data[:train_split].std(axis=0)
    return (data - data_mean) / data_std


print(
    "The selected parameters are:",
    ", ".join([titles[i] for i in [0, 1, 5, 7, 8, 10, 11]]),
)
selected_features = [feature_keys[i] for i in [0, 1, 5, 7, 8, 10, 11]]
features = df[selected_features]
features.index = df[date_time_key]
features.head()

features = normalize(features.values, train_split)
features = pd.DataFrame(features)
features.head()

train_data = features.loc[0 : train_split - 1]
val_data = features.loc[train_split:]


split_fraction = 0.715
train_split = int(split_fraction * int(df.shape[0]))
step = 6

past = 720
future = 72
learning_rate = 0.001
batch_size = 256
epochs = 10


def normalize(data, train_split):
    data_mean = data[:train_split].mean(axis=0)
    data_std = data[:train_split].std(axis=0)
    return (data - data_mean) / data_std
print(
    "The selected parameters are:",
    ", ".join([titles[i] for i in [0, 1, 5, 7, 8, 10, 11]]),
)
selected_features = [feature_keys[i] for i in [0, 1, 5, 7, 8, 10, 11]]
features = df[selected_features]
features.index = df[date_time_key]
features.head()

features = normalize(features.values, train_split)
features = pd.DataFrame(features)
features.head()

train_data = features.loc[0 : train_split - 1]
val_data = features.loc[train_split:]
start = past + future
end = start + train_split

x_train = train_data[[i for i in range(7)]].values
y_train = features.iloc[start:end][[1]]

sequence_length = int(past / step)
x_end = len(val_data) - past - future

label_start = train_split + past + future

x_val = val_data.iloc[:x_end][[i for i in range(7)]].values
y_val = features.iloc[label_start:][[1]]

dataset_val = keras.preprocessing.timeseries_dataset_from_array(
    x_val,
    y_val,
    sequence_length=sequence_length,
    sampling_rate=step,
    batch_size=batch_size,
)
#%%
model = from_pretrained_keras("keras-io/timeseries_forecasting_for_weather")

#%% 
st.set_option('deprecation.showPyplotGlobalUse', False)
def plot():
    n = st.sidebar.slider("Step", min_value = 1, max_value=5, value = 1)
    def show_plot(plot_data, delta, title):
        labels = ["History", "True Future", "Model Prediction"]
        marker = [".-", "rx", "go"]
        time_steps = list(range(-(plot_data[0].shape[0]), 0))
        if delta:
            future = delta
        else:
            future = 0

        plt.title(title)
        for i, val in enumerate(plot_data):
            if i:
                plt.plot(future, plot_data[i], marker[i], markersize=10, label=labels[i])
            else:
                plt.plot(time_steps, plot_data[i].flatten(), marker[i], label=labels[i])
        plt.legend(loc='lower center', bbox_to_anchor=(0.5, 1.05),
          ncol=3, fancybox=True, shadow=True)
        plt.xlim([time_steps[0], (future + 5) * 2])
        plt.xlabel("Time-Step")
        plt.show()
        return


    for x, y in dataset_val.take(n):
        show_plot(
            [x[0][:, 1].numpy(), y[0].numpy(), model.predict(x)[0]],
            12,
            f"{n} Step Prediction",
        )

fig = plot()
st.pyplot(fig)

# %%