import gradio as gr
import torch
import numpy as np
import matplotlib.pyplot as plt
import io
import base64
import pickle
import torch.nn as nn
from sklearn.preprocessing import MinMaxScaler
import pandas as pd

# Load your dataset
df_weekly_km = pd.read_csv('Weekly_Km.csv')

# Model parameters
INPUT_SIZE = 1
HIDDEN_LAYER_SIZE = 100
OUTPUT_SIZE = 1

# Normalize the data
data = df_weekly_km['Total Kilometers'].values.astype(float)
scaler = MinMaxScaler(feature_range=(-1, 1))
data_normalized = scaler.fit_transform(data.reshape(-1, 1))

seq_length = 4

# Define LSTM model
class LSTM(nn.Module):
    def __init__(self, input_size=INPUT_SIZE, hidden_layer_size=HIDDEN_LAYER_SIZE, output_size=OUTPUT_SIZE):
        super(LSTM, self).__init__()
        self.hidden_layer_size = hidden_layer_size
        self.lstm = nn.LSTM(input_size, hidden_layer_size)
        self.linear = nn.Linear(hidden_layer_size, output_size)
        self.hidden_cell = (torch.zeros(1, 1, self.hidden_layer_size),
                            torch.zeros(1, 1, self.hidden_layer_size))

    def forward(self, input_seq):
        lstm_out, self.hidden_cell = self.lstm(input_seq.view(len(input_seq), 1, -1), self.hidden_cell)
        predictions = self.linear(lstm_out.view(len(input_seq), -1))
        return predictions[-1]

# Function to prepare input
def prepare_custom_input(last_values, seq_length, scaler):
    last_values_normalized = scaler.transform(np.array(last_values).reshape(-1, 1))
    input_seq = torch.from_numpy(last_values_normalized).float()
    return input_seq.view(-1)

# Load the trained model
model_path = 'lstm_model.pth'
model = LSTM(INPUT_SIZE, HIDDEN_LAYER_SIZE, OUTPUT_SIZE)
model.load_state_dict(torch.load(model_path))

# Prediction and plot function
def predict_and_plot(week4, week3, week2, week1):
    last_four_weeks = [week4, week3, week2, week1]

    custom_input = prepare_custom_input(last_four_weeks, seq_length, scaler)

    model.eval()
    with torch.no_grad():
        model.hidden_cell = (torch.zeros(1, 1, model.hidden_layer_size),
                             torch.zeros(1, 1, model.hidden_layer_size))
        prediction = model(custom_input)

    predicted_kilometers = scaler.inverse_transform(prediction.numpy().reshape(-1, 1))
    predicted_value = predicted_kilometers[0][0]

    weeks = ['Week -4', 'Week -3', 'Week -2', 'Week -1', 'Predicted Week']
    actual_values = last_four_weeks + [predicted_value]
    
    plt.figure(figsize=(10, 6))
    plt.plot(weeks, actual_values, marker='o', label='Total Kilometers')
    plt.axvline(x='Predicted Week', color='red', linestyle='--', label='Predicted Week')
    plt.title('Total Kilometers for Last 4 Weeks and Prediction')
    plt.xlabel('Weeks')
    plt.ylabel('Total Kilometers')
    plt.ylim(bottom=0)
    plt.grid()
    plt.tight_layout()
    plt.legend()

    # Saving the plot to a BytesIO object
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    plt.close()

    # Converting to base64 for Gradio
    img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    img_tag = f'<img src="data:image/png;base64,{img_base64}"/>'
    
    return f"{predicted_value:.2f}", img_tag

# Gradio interface
inputs = [gr.Number(label='Week -4'), 
          gr.Number(label='Week -3'), 
          gr.Number(label='Week -2'), 
          gr.Number(label='Week -1')]
outputs = [gr.Textbox(label='Predicted Week'), 
           gr.HTML(label='Plot')]

gr.Interface(fn=predict_and_plot, 
             inputs=inputs, 
             outputs=outputs,
             title="LSTM Model Prediction",
             description="Enter the total kilometers for the last 4 weeks to get the prediction for the next week along with a plot.").launch(share=True)