import gradio as gr
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense

def lstm_forecast(file, forecast_months):
    # Load and preprocess data
    df = pd.read_csv(file)
    df['Month'] = pd.to_datetime(df['Month'])
    df.set_index('Month', inplace=True)
    df = df.sort_index()

    if 'Spend' not in df.columns:
        raise ValueError("CSV must contain a 'Spend' column.")

    data = df[['Spend']]
    
    # Scale data
    scaler = MinMaxScaler()
    scaled_data = scaler.fit_transform(data)

    # Create sequences
    def create_sequences(data, seq_length):
        X, y = [], []
        for i in range(len(data) - seq_length):
            X.append(data[i:i+seq_length])
            y.append(data[i+seq_length])
        return np.array(X), np.array(y)

    seq_length = 12
    X, y = create_sequences(scaled_data, seq_length)
    X = X.reshape((X.shape[0], X.shape[1], 1))

    # Build LSTM model
    model = Sequential([
        LSTM(64, return_sequences=False, input_shape=(seq_length, 1)),
        Dense(1)
    ])
    model.compile(optimizer='adam', loss='mse')
    model.fit(X, y, epochs=30, batch_size=1, verbose=0)

    # Forecast future values
    forecast_input = scaled_data[-seq_length:].reshape((1, seq_length, 1))
    forecast = []

    for _ in range(forecast_months):
        pred = model.predict(forecast_input, verbose=0)[0]
        forecast.append(pred)
        forecast_input = np.append(forecast_input[:, 1:, :], [[pred]], axis=1)

    forecast = scaler.inverse_transform(forecast)

    # Create forecast dates
    last_date = df.index[-1]
    forecast_dates = pd.date_range(start=last_date + pd.DateOffset(months=1), periods=forecast_months, freq='M')

    # Plotting
    plt.figure(figsize=(12, 6))
    plt.plot(df.index, df['Spend'], label='Historical Spend', marker='o')
    plt.plot(forecast_dates, forecast, label='LSTM Forecast', marker='o', color='green')
    plt.title("Credit Card Spend Forecast with LSTM")
    plt.xlabel("Month")
    plt.ylabel("Spend")
    plt.grid(True)
    plt.legend()

    return plt.gcf()

# Gradio Interface
interface = gr.Interface(
    fn=lstm_forecast,
    inputs=[
        gr.File(label="Upload credit_card_spend.csv"),
        gr.Slider(minimum=6, maximum=24, step=1, value=12, label="Months to Forecast")
    ],
    outputs=gr.Plot(label="LSTM Forecast Plot"),
    title="📊 Credit Card Spend Forecast using LSTM",
    description="Upload a CSV with 'Month' and 'Spend' columns to forecast future trends using an LSTM model."
)

interface.launch()