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

# Load the model and dataset
model = joblib.load("xgb_model.joblib")
df = pd.read_csv("worker_productivity.csv")  # Ensure this is uploaded to Hugging Face Space

# Get unique teams
teams = sorted(df['team'].unique())

# Define task generation function
def generate_task():
    task = {
        'quarter': random.choice(['Q1', 'Q2', 'Q3', 'Q4']),
        'department': random.choice(['sewing', 'finishing']),
        'day': random.choice(['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday']),
        'no_of_workers': random.randint(30, 60),
        'incentive': round(random.uniform(1.0, 4.0), 2),
        'idle_time': round(random.uniform(0.0, 1.0), 2),
        'idle_men': random.randint(0, 10),
        'smv': round(random.uniform(20.0, 40.0), 2),
        'month': random.randint(1, 12),
        'day_of_week': random.randint(0, 6),
        'is_weekend': random.choice([0, 1]),
        'no_of_style_change': random.randint(0, 2),
        'targeted_productivity': round(random.uniform(0.5, 0.95), 2)
    }
    task['smv_per_worker'] = task['smv'] / task['no_of_workers']
    task['effort_index'] = task['smv'] + task['incentive'] + 1.0 - task['idle_time']
    task['log_wip'] = np.log1p(random.randint(20, 100))
    task['log_overtime'] = np.log1p(random.uniform(0.5, 3.0))
    return task

# Prediction function
def predict():
    task = generate_task()
    team_scores = []

    for team in teams:
        t = task.copy()
        t['team'] = team
        task_df = pd.DataFrame([t])
        pred = model.predict(task_df)[0]
        team_scores.append((team, pred))

    # Sort results
    team_scores_df = pd.DataFrame(team_scores, columns=["Team", "Predicted Productivity"])
    team_scores_df = team_scores_df.sort_values(by="Predicted Productivity", ascending=False)

    # Plot results
    fig, ax = plt.subplots(figsize=(10, 6))
    bars = ax.barh(team_scores_df["Team"].astype(str), team_scores_df["Predicted Productivity"], color='skyblue')
    ax.set_xlabel("Predicted Productivity")
    ax.set_title("Predicted Productivity by Team for Custom Task")
    ax.invert_yaxis()

    for bar in bars:
        width = bar.get_width()
        ax.text(width + 0.01, bar.get_y() + bar.get_height()/2, f"{width:.2f}", va='center')

    plt.tight_layout()

    task_info = "\n".join([f"{k}: {v}" for k, v in task.items()])
    return fig, task_info

# Gradio UI
demo = gr.Interface(
    fn=predict,
    inputs=[],
    outputs=[
        gr.Plot(label="Team Productivity Rankings"),
        gr.Textbox(label="Task Details", lines=20)
    ],
    live=False,
    title="Worker Productivity Predictor",
    description="Generates predicted productivity scores for each team on a random custom task."
)

demo.launch()