import pandas as pd
import numpy as np
import torch
from torch import nn
from torch.optim import Adam
import streamlit as st
import matplotlib.pyplot as plt

df = pd.read_csv("data.csv")
df = df[5::6]
df.index = pd.to_datetime(df['Date Time'], format='%d.%m.%Y %H:%M:%S')
temp = df['T (degC)']


def df_to_X_y(df, window_size=5):
    df_as_np = df.to_numpy()
    X = []
    y = []
    for i in range(len(df_as_np) - window_size):
        row = [[a] for a in df_as_np[i:i + window_size]]
        X.append(row)
        label = df_as_np[i + window_size]
        y.append(label)
    return np.array(X), np.array(y)


WINDOW_SIZE = 5
X1, y1 = df_to_X_y(temp, WINDOW_SIZE)
X_train1, y_train1 = X1[:60000], y1[:60000]
X_val1, y_val1 = X1[60000:65000], y1[60000:65000]
X_test1, y_test1 = X1[65000:], y1[65000:]

DEVICE = "cpu"
X_train = torch.tensor(X_train1, dtype=torch.float32).to(DEVICE)
y_train = torch.tensor(y_train1, dtype=torch.float32).to(DEVICE)
X_val = torch.tensor(X_val1, dtype=torch.float32).to(DEVICE)
y_val = torch.tensor(y_val1, dtype=torch.float32).to(DEVICE)
X_test = torch.tensor(X_test1, dtype=torch.float32).to(DEVICE)
y_test = torch.tensor(y_test1, dtype=torch.float32).to(DEVICE)


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.lstm1 = nn.LSTM(input_size=1, hidden_size=64, batch_first=True)
        self.lstm2 = nn.LSTM(input_size=64, hidden_size=64, batch_first=True)
        self.fc1 = nn.Linear(64, 8)
        self.fc2 = nn.Linear(8, 1)
        self.relu = nn.ReLU()

    def forward(self, x):
        x, _ = self.lstm1(x)
        x, _ = self.lstm2(x)
        x = self.fc1(x[:, -1, :])
        x = self.relu(x)
        x = self.fc2(x)
        x = x.view(-1)
        return x


model = Model().to(DEVICE)
criterion = nn.MSELoss().to(DEVICE)
optimizer = Adam(model.parameters(), lr=0.01)

model.load_state_dict(torch.load("timeseries.pt", map_location=torch.device(DEVICE)))
model.eval()

train_predictions = model(X_train)
train_results = pd.DataFrame(data={'Train Predictions': train_predictions.cpu().detach().numpy(), 'Actuals': y_train1})

val_predictions = model(X_val)
val_results = pd.DataFrame(data={'Val Predictions': val_predictions.cpu().detach().numpy(), 'Actuals': y_val1})

test_predictions = model(X_test)
test_results = pd.DataFrame(data={'Test Predictions': test_predictions.cpu().detach().numpy(), 'Actuals': y_test1})


def plot_train():
    fig, ax = plt.subplots()
    ax.set_title('Train Results')
    plt.plot(train_results['Train Predictions'][:100])
    plt.plot(train_results['Actuals'][:100])
    return fig


def plot_test():
    fig, ax = plt.subplots()
    ax.set_title('Train Results')
    plt.plot(test_results['Test Predictions'][:100])
    plt.plot(test_results['Actuals'][:100])
    return fig


def plot_val():
    fig, ax = plt.subplots()
    ax.set_title('Train Results')
    plt.plot(val_results['Val Predictions'][:100])
    plt.plot(val_results['Actuals'][:100])
    return fig


# Streamlit App
def main():
    st.title('Time Series Forecasting using LSTM')

    st.subheader(f'Training Data')

    col1, col2 = st.columns(2)
    with col1:
        st.write(train_results)
    with col2:
        fig = plot_train()
        st.pyplot(fig)

    st.subheader(f'Testing Data')

    col1, col2 = st.columns(2)
    with col1:
        st.write(test_results)
    with col2:
        fig = plot_test()
        st.pyplot(fig)

    st.subheader(f'Validation Data')

    col1, col2 = st.columns(2)
    with col1:
        st.write(val_results)
    with col2:
        fig = plot_val()
        st.pyplot(fig)


if __name__ == "__main__":
    main()