# Imports
import pandas as pd
import matplotlib.pyplot as plt
from prophet import Prophet
import gradio as gr

# Load and preprocess data
url = 'https://data.open-power-system-data.org/time_series/2020-10-06/time_series_60min_singleindex.csv'
data = pd.read_csv(url)

df = data[['utc_timestamp', 'DE_load_actual_entsoe_transparency']].copy()
df.rename(columns={'utc_timestamp': 'ds', 'DE_load_actual_entsoe_transparency': 'y'}, inplace=True)
df.dropna(inplace=True)
df['ds'] = pd.to_datetime(df['ds']).dt.tz_localize(None)

# Forecast function with dropdown support
def forecast_energy(days, graph_type):
    model = Prophet()
    model.fit(df)

    future = model.make_future_dataframe(periods=24 * days, freq='H')
    forecast = model.predict(future)
    forecast_future = forecast.tail(24 * days)

    # Paths to store each plot
    paths = {}

    # Forecast Plot
    plt.figure(figsize=(12, 5))
    plt.plot(forecast_future['ds'], forecast_future['yhat'], label='Predicted Load (MW)', color='blue')
    plt.fill_between(forecast_future['ds'], forecast_future['yhat_lower'], forecast_future['yhat_upper'], color='skyblue', alpha=0.3, label='Confidence Interval')
    plt.xticks(rotation=45)
    plt.xlabel("Date")
    plt.ylabel("Load (MW)")
    plt.title(f"Forecasted Load (Next {days} Days)")
    plt.legend()
    plt.tight_layout()
    paths["Forecast Plot"] = f"forecast_plot_{days}.png"
    plt.savefig(paths["Forecast Plot"])
    plt.close()

    # Actual vs Forecast
    plt.figure(figsize=(12, 4))
    plt.plot(df['ds'].tail(24 * 7), df['y'].tail(24 * 7), label='Actual Load (Past Week)', color='black')
    plt.plot(forecast_future['ds'], forecast_future['yhat'], label='Forecast Load', color='green')
    plt.xticks(rotation=45)
    plt.title("Actual vs Forecast")
    plt.xlabel("Date")
    plt.ylabel("Load (MW)")
    plt.legend()
    plt.tight_layout()
    paths["Actual vs Forecast"] = f"actual_vs_forecast_{days}.png"
    plt.savefig(paths["Actual vs Forecast"])
    plt.close()

    # Hourly Average
    df['hour'] = df['ds'].dt.hour
    hourly_avg = df.groupby('hour')['y'].mean()
    plt.figure(figsize=(10, 4))
    plt.plot(hourly_avg.index, hourly_avg.values, marker='o', linestyle='-', color='orange')
    plt.title("Average Load by Hour of Day")
    plt.xlabel("Hour")
    plt.ylabel("Average Load (MW)")
    plt.grid(True)
    paths["Hourly Pattern"] = f"hourly_pattern.png"
    plt.savefig(paths["Hourly Pattern"])
    plt.close()

    # Seasonality
    seasonality_fig = model.plot_components(forecast)
    paths["Seasonality"] = f"seasonality_{days}.png"
    seasonality_fig.savefig(paths["Seasonality"])
    plt.close()

    # Peak demand
    peak_row = forecast_future.loc[forecast_future['yhat'].idxmax()]
    peak_time = peak_row['ds']
    peak_value = round(peak_row['yhat'], 2)
    peak_info = f"🔺 Peak Demand Time: {peak_time} — {peak_value} MW"

    return paths[graph_type], peak_info

# Gradio Interface
iface = gr.Interface(
    fn=forecast_energy,
    inputs=[
        gr.Radio([3, 7, 14], label="Select Forecast Period (Days)"),
        gr.Dropdown(["Forecast Plot", "Actual vs Forecast", "Hourly Pattern", "Seasonality"], label="Select Graph Type")
    ],
    outputs=[
        gr.Image(type="filepath", label="Selected Plot"),
        gr.Textbox(label="Peak Demand Info")
    ],
    title="Smart Energy Load Forecasting",
    description="Choose forecast days and which graph you want to visualize."
)

iface.launch(share=True)