import pandas as pd
import numpy as np
from datetime import datetime
import matplotlib.pyplot as plt
import gradio as gr
from xgboost import XGBRegressor
from sklearn.model_selection import train_test_split

def simulate_factory_data(days=7, freq='H'):
    date_rng = pd.date_range(end=datetime.now(), periods=24 * days, freq=freq)
    df = pd.DataFrame(date_rng, columns=['ds'])
    df['temperature'] = np.random.normal(loc=70, scale=4, size=(len(date_rng)))
    anomaly_indices = np.random.choice(len(df), size=8, replace=False)
    df.loc[anomaly_indices, 'temperature'] += np.random.uniform(10, 20, size=8)
    return df

def create_lag_features(df, lags=6):
    for i in range(1, lags + 1):
        df[f'lag_{i}'] = df['temperature'].shift(i)
    df = df.dropna().reset_index(drop=True)
    return df

def forecast_temperature(hours):
    df = simulate_factory_data()
    df = create_lag_features(df)
    X = df[[f'lag_{i}' for i in range(1, 7)]]
    y = df['temperature']
    model = XGBRegressor(n_estimators=100, learning_rate=0.1)
    model.fit(X, y)

    last_row = df.iloc[-1][[f'lag_{i}' for i in range(1, 7)]].values
    preds = []
    for _ in range(int(hours)):
        pred = model.predict([last_row])[0]
        preds.append(pred)
        last_row = np.roll(last_row, -1)
        last_row[-1] = pred

    future_dates = pd.date_range(start=df['ds'].iloc[-1], periods=int(hours)+1, freq='H')[1:]
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(df['ds'], df['temperature'], label='Actual')
    ax.plot(future_dates, preds, label='Forecast', color='orange')
    ax.axhspan(85, 100, color='red', alpha=0.1, label='Danger Zone')
    ax.set_title("Machine Temperature Forecast (XGBoost)")
    ax.set_xlabel("Time")
    ax.set_ylabel("Temperature (°C)")
    ax.legend()
    plt.tight_layout()
    return fig

demo = gr.Interface(
    fn=forecast_temperature,
    inputs=gr.Slider(6, 48, value=24, step=1, label="Forecast Hours"),
    outputs=gr.Plot(label="Forecasted Temperature Chart"),
    title="Smart Factory AI Pipeline (XGBoost)",
    description="Forecast machine temperature using XGBoost. Red zone shows potential overheating."
)

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