import gradio as gr
import numpy as np
import pandas as pd
from tensorflow import keras
from sklearn.preprocessing import StandardScaler

model = keras.models.load_model('productivity_model.h5')
scaler = StandardScaler()
scaler.mean_ = np.load('scaler_mean.npy')
scaler.scale_ = np.load('scaler_scale.npy')
feature_columns = open('feature_columns.txt').read().splitlines()

def predict_csv(file):
    df = pd.read_csv(file.name)
    df = df.drop(columns=['user_id', 'Unnamed: 15', 'productivity_0_100'], errors='ignore').dropna()
    df_enc = pd.get_dummies(df, columns=['gender', 'occupation', 'work_mode'], drop_first=True)
    df_enc = df_enc.reindex(columns=feature_columns, fill_value=0).astype('float32')

    preds = model.predict(scaler.transform(df_enc), verbose=0).flatten().round(2)
    result = df.copy()
    result['predicted_productivity'] = preds
    return result

gr.Interface(
    fn=predict_csv,
    inputs=gr.File(label="Upload CSV"),
    outputs=gr.Dataframe(label="Predictions"),
    title="Productivity Predictor",
    description="Upload a CSV with the same columns as training data. Returns predicted productivity for each row."
).launch()