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

# Load the model
model = joblib.load("model.joblib")

# Load the original dataset to extract real specialty mappings
df = pd.read_csv("dataset.csv")  # <-- Make sure this file is uploaded to Hugging Face

# Get actual teams and specialties
teams = sorted(df['Team'].unique())
specialty_map = dict(zip(df['Team'], df['Specialty']))

# Sample task examples (all feature-engineered already)
sample_tasks = [
    {
        'ProductType': 'Mothball',
        'TaskType': 'Packaging',
        'OrderQuantity': 120,
        'DeadlineDays': 4,
        'ExperienceYears': 6,
        'AvgTaskTime_Minutes': 28.0,
        'ErrorRate': 0.05,
        'TrainingHours': 20.0,
        'DayNumber': 2
    },
    {
        'ProductType': 'Perfume',
        'TaskType': 'Boxing',
        'OrderQuantity': 80,
        'DeadlineDays': 2,
        'ExperienceYears': 2,
        'AvgTaskTime_Minutes': 45.0,
        'ErrorRate': 0.12,
        'TrainingHours': 10.0,
        'DayNumber': 3
    },
    {
        'ProductType': 'Soap',
        'TaskType': 'Molding',
        'OrderQuantity': 200,
        'DeadlineDays': 6,
        'ExperienceYears': 4,
        'AvgTaskTime_Minutes': 32.0,
        'ErrorRate': 0.07,
        'TrainingHours': 15.0,
        'DayNumber': 4
    },
]

# Prediction function
def predict():
    task = random.choice(sample_tasks)

    # Feature engineering
    task['ThroughputRate'] = task['OrderQuantity'] / task['AvgTaskTime_Minutes']
    task['TimePressure'] = task['OrderQuantity'] / ((task['DeadlineDays'] or 1) * task['AvgTaskTime_Minutes'])
    task['PriorityLevel'] = 3  # high priority

    # Predict for each team
    results = []
    for team in teams:
        row = task.copy()
        row['Team'] = team
        row['Specialty'] = specialty_map.get(team, "Unknown")
        input_df = pd.DataFrame([row])
        pred = model.predict(input_df)[0]
        results.append((team, round(pred, 4)))

    # Sort results
    results.sort(key=lambda x: x[1], reverse=True)

    # Plot
    teams_sorted, scores_sorted = zip(*results)
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.barh(teams_sorted, scores_sorted, color="skyblue")
    ax.set_xlabel("Predicted Productivity")
    ax.set_title("Team Ranking for Random Task")
    ax.invert_yaxis()
    plt.tight_layout()

    return fig, json.dumps(task, indent=2)

# Gradio UI
demo = gr.Interface(
    fn=predict,
    inputs=[],
    outputs=[gr.Plot(label="Team Productivity Rankings"), gr.Textbox(label="Task Details")],
    live=False,
    title="Team Productivity Predictor",
    description="Click Generate to predict the best team for a randomly selected task."
)

demo.launch()