import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
import gradio as gr
import matplotlib.pyplot as plt

def train_model(df):
    # Extract year (first column) and target (second column)
    years = df.iloc[:, 0].values.reshape(-1, 1)
    target = df.iloc[:, 1].values

    # Determine model type based on number of columns
    if df.shape[1] > 2:
        # Multiple regression (year + additional features)
        features = df.iloc[:, 2:].values
        X = np.hstack((years, features))
        model = LinearRegression()
        model.fit(X, target)
        return model, 'multiple', None, features
    else:
        # Polynomial regression (only year as feature)
        poly = PolynomialFeatures(degree=2)
        X_poly = poly.fit_transform(years)
        model = LinearRegression()
        model.fit(X_poly, target)
        return model, 'poly', poly, None

def predict(dataset, years_to_predict):
    # Read dataset
    if dataset.name.endswith('.csv'):
        df = pd.read_csv(dataset.name)
    else:
        df = pd.read_excel(dataset.name)
    
    # Validate dataset
    if df.shape[1] < 2:
        raise gr.Error("Dataset must have at least 2 columns: Year and Target!")
    
    # Train model
    model, model_type, poly_features, features = train_model(df)
    
    # Prepare future years
    last_year = df.iloc[-1, 0]
    future_years = np.arange(last_year + 1, last_year + 1 + years_to_predict).reshape(-1, 1)
    
    # Generate predictions
    if model_type == 'poly':
        future_X = poly_features.transform(future_years)
    else:
        # Use last available features for future predictions
        last_features = df.iloc[-1, 2:].values.reshape(1, -1)
        repeated_features = np.repeat(last_features, years_to_predict, axis=0)
        future_X = np.hstack((future_years, repeated_features))
    
    predictions = model.predict(future_X)
    
    # Create visualization
    plt.figure(figsize=(10, 5))
    plt.plot(df.iloc[:, 0], df.iloc[:, 1], 'bo-', label='Historical Data')
    plt.plot(future_years, predictions, 'ro--', label='Predictions')
    plt.xlabel('Year')
    plt.ylabel('Target Value')
    plt.title('Time Series Forecast')
    plt.legend()
    plt.grid(True)
    
    # Create output DataFrame
    result_df = pd.DataFrame({
        'Year': future_years.flatten(),
        'Predicted Value': predictions.round(2)
    })
    
    return plt, result_df

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 Time Series Forecasting Tool")
    with gr.Row():
        file_input = gr.File(label="Upload Dataset (CSV/Excel)")
        years_input = gr.Dropdown([1, 2, 5, 10], value=5, label="Years to Predict")
    btn = gr.Button("PREDICT")
    with gr.Row():
        plot_output = gr.Plot()
        table_output = gr.DataFrame()
    
    btn.click(
        fn=predict,
        inputs=[file_input, years_input],
        outputs=[plot_output, table_output]
    )

if __name__ == "__main__":
    demo.launch()