app.py

import gradio as gr
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from sklearn.cluster import KMeans
from transformers import pipeline
from datasets import load_dataset
import matplotlib.pyplot as plt
import os

# --- 1. Load Data and Models ---
print("Loading dataset...")
dataset = load_dataset("nprak26/remote-worker-productivity")
df = dataset["train"].to_pandas()

print("Creating new features...")
# Create Work_Life_Balance_Score
df['Work_Life_Balance_Score'] = df['Work_Hours_Per_Week'] * df['Work_Life_Balance']

# Create Communication_Overhead
collaboration_mapping = {
    'Monthly': 1,
    'Weekly': 2,
    'Few times per week': 3,
    'Daily': 4
}
df['Team_Collaboration_Numerical'] = df['Team_Collaboration_Frequency'].map(collaboration_mapping)
df['Communication_Overhead'] = df['Meetings_Per_Week'] + df['Team_Collaboration_Numerical']

# Define features for clustering
features = [
    "Task_Completion_Rate",
    "Stress_Level",
    "Job_Satisfaction",
    "Manager_Support_Level",
    "Team_Collaboration_Frequency",
    "Work_Life_Balance_Score",
    "Communication_Overhead"
]

# Redefine df_processed and handle missing values
df_processed = pd.get_dummies(df[features], columns=["Manager_Support_Level", "Team_Collaboration_Frequency"], drop_first=True)
# Handle missing values in 'Communication_Overhead'
df_processed['Communication_Overhead'].fillna(df_processed['Communication_Overhead'].median(), inplace=True)
# Handle potential NaNs in 'Work_Life_Balance_Score' if any
df_processed['Work_Life_Balance_Score'].fillna(df_processed['Work_Life_Balance_Score'].median(), inplace=True)


print("Scaling features...")
scaler = MinMaxScaler()
scaled_features = scaler.fit_transform(df_processed)

print("Fitting KMeans model with optimal k=3...")
# Optimal k=3 was determined from the elbow method plot
kmeans = KMeans(n_clusters=3, random_state=42, n_init=10)
kmeans.fit(scaled_features)

print("Loading text generation model...")
# Using `pipeline` to load the text generation model
generator = pipeline("text2text-generation", model="google/flan-t5-base")

# Placeholder descriptions for the 3 clusters. These would ideally be derived from an analysis of the clusters.
# For this example, we will use a generic description.
cluster_descriptions = {
    0: "This cluster is characterized by high productivity metrics and good work-life balance.",
    1: "This cluster shows signs of moderate stress and varying productivity levels.",
    2: "This cluster may be struggling with high stress, low satisfaction, or other challenges."
}

# --- 2. Define Helper Functions ---
def generate_tip(task_completion, stress, satisfaction, support, collaboration, cluster_description):
    """
    Generates an actionable productivity tip.
    """
    prompt = f"""
    Based on the following productivity profile:
    - Task Completion Rate: {task_completion}%
    - Stress Level: {stress}/10
    - Job Satisfaction: {satisfaction}%
    - Manager Support Level: {support}
    - Team Collaboration Frequency: {collaboration}

    This profile aligns with a cluster characterized by: {cluster_description}

    Suggest one highly actionable tip to improve productivity, specifically tailored to this profile and cluster characteristics. The tip should be concise and practical.
    """
    tip = generator(prompt, max_new_tokens=50)[0]["generated_text"]
    return tip

def analyze_productivity(task_completion, stress, satisfaction, support, collaboration):
    """
    Analyzes user productivity metrics, predicts a cluster, and generates a tip.
    """
    user_data = pd.DataFrame([{
        "Task_Completion_Rate": task_completion,
        "Stress_Level": stress,
        "Job_Satisfaction": satisfaction,
        "Manager_Support_Level": support,
        "Team_Collaboration_Frequency": collaboration,
        "Work_Life_Balance_Score": 0,  # Placeholder, not used for tip generation
        "Communication_Overhead": 0    # Placeholder
    }])

    # One-hot encode the user input
    user_processed = pd.get_dummies(user_data, columns=["Manager_Support_Level", "Team_Collaboration_Frequency"], drop_first=True)

    # Align columns with the training data
    user_aligned = user_processed.reindex(columns=df_processed.columns, fill_value=0)

    # Scale the user input
    scaled_input = scaler.transform(user_aligned)

    # Predict the cluster
    cluster = kmeans.predict(scaled_input)[0]
    cluster_description = cluster_descriptions[cluster]

    # Generate the tip
    tip = generate_tip(task_completion, stress, satisfaction, support, collaboration, cluster_description)

    return f"Cluster {cluster}: {cluster_description}", tip

def analyze_what_if(what_if_task_completion, what_if_stress, what_if_satisfaction, what_if_support, what_if_collaboration):
    """
    Analyzes 'what-if' user productivity metrics, predicts a cluster, and generates a tip.
    """
    user_data = pd.DataFrame([{
        "Task_Completion_Rate": what_if_task_completion,
        "Stress_Level": what_if_stress,
        "Job_Satisfaction": what_if_satisfaction,
        "Manager_Support_Level": what_if_support,
        "Team_Collaboration_Frequency": what_if_collaboration,
        "Work_Life_Balance_Score": 0,
        "Communication_Overhead": 0
    }])

    # One-hot encode the user input
    user_processed = pd.get_dummies(user_data, columns=["Manager_Support_Level", "Team_Collaboration_Frequency"], drop_first=True)
    user_aligned = user_processed.reindex(columns=df_processed.columns, fill_value=0)
    scaled_input = scaler.transform(user_aligned)

    # Predict the cluster
    cluster = kmeans.predict(scaled_input)[0]
    cluster_description = cluster_descriptions[cluster]

    # Generate the tip
    tip = generate_tip(what_if_task_completion, what_if_stress, what_if_satisfaction, what_if_support, what_if_collaboration, cluster_description)

    return f"Cluster {cluster}: {cluster_description}", tip

# Function to save user feedback
def save_feedback(tip, feedback):
    """
    Saves the generated tip and user feedback to a CSV file.
    """
    feedback_file = "feedback.csv"
    feedback_data = pd.DataFrame({"tip": [tip], "feedback": [feedback]})

    # Check if the file exists to decide whether to write the header
    if os.path.exists(feedback_file):
        with open(feedback_file, "a") as f:
            feedback_data.to_csv(f, header=False, index=False)
    else:
        feedback_data.to_csv(feedback_file, header=True, index=False)

    return "Feedback submitted!"

# --- 3. Build the Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 Productivity Profiler")
    gr.Markdown("Enter your productivity metrics to get a personalized tip. This tool analyzes your input and provides an AI-generated tip to help improve your productivity.")

    with gr.Tabs():
        with gr.TabItem("Your Productivity Profile"):
            task_completion = gr.Slider(0, 100, label="Task Completion Rate (%)")
            stress = gr.Slider(1, 10, label="Stress Level (1-10)")
            satisfaction = gr.Slider(0, 100, label="Job Satisfaction (0-100)")
            support = gr.Dropdown(
                choices=['Very Low', 'Low', 'Moderate', 'High', 'Very High'],
                label="Manager Support Level"
            )
            collaboration = gr.Dropdown(
                choices=['Monthly', 'Weekly', 'Few times per week', 'Daily'],
                label="Team Collaboration Frequency"
            )
            submit = gr.Button("Analyze")

        with gr.TabItem("Results"):
            profile_output = gr.Textbox(label="Your Productivity Cluster")
            tip_output = gr.Textbox(label="AI-Powered Productivity Tip")
            with gr.Row():
                upvote_button = gr.Button("👍")
                downvote_button = gr.Button("👎")
            feedback_message = gr.Textbox(label="Feedback Status")

        with gr.TabItem("What-If Scenario"):
            gr.Markdown("Explore how changing your metrics could affect your productivity profile.")
            what_if_task_completion = gr.Slider(0, 100, label="Adjusted Task Completion Rate (%)")
            what_if_stress = gr.Slider(1, 10, label="Adjusted Stress Level (1-10)")
            what_if_satisfaction = gr.Slider(0, 100, label="Adjusted Job Satisfaction (0-100)")
            what_if_support = gr.Dropdown(
                choices=['Very Low', 'Low', 'Moderate', 'High', 'Very High'],
                label="Adjusted Manager Support Level"
            )
            what_if_collaboration = gr.Dropdown(
                choices=['Monthly', 'Weekly', 'Few times per week', 'Daily'],
                label="Adjusted Team Collaboration Frequency"
            )
            what_if_submit = gr.Button("Analyze What-If Scenario")

        with gr.TabItem("What-If Results"):
            what_if_profile_output = gr.Textbox(label="What-If Productivity Cluster")
            what_if_tip_output = gr.Textbox(label="What-If AI-Powered Productivity Tip")
            with gr.Row():
                what_if_upvote_button = gr.Button("👍")
                what_if_downvote_button = gr.Button("👎")
            what_if_feedback_message = gr.Textbox(label="Feedback Status")


    submit.click(
        fn=analyze_productivity,
        inputs=[task_completion, stress, satisfaction, support, collaboration],
        outputs=[profile_output, tip_output]
    )

    what_if_submit.click(
        fn=analyze_what_if,
        inputs=[what_if_task_completion, what_if_stress, what_if_satisfaction, what_if_support, what_if_collaboration],
        outputs=[what_if_profile_output, what_if_tip_output]
    )

    upvote_button.click(
        fn=save_feedback,
        inputs=[tip_output, gr.Textbox(value="👍", visible=False)],
        outputs=[feedback_message]
    )
    downvote_button.click(
        fn=save_feedback,
        inputs=[tip_output, gr.Textbox(value="👎", visible=False)],
        outputs=[feedback_message]
    )
    what_if_upvote_button.click(
        fn=save_feedback,
        inputs=[what_if_tip_output, gr.Textbox(value="👍", visible=False)],
        outputs=[what_if_feedback_message]
    )
    what_if_downvote_button.click(
        fn=save_feedback,
        inputs=[what_if_tip_output, gr.Textbox(value="👎", visible=False)],
        outputs=[what_if_feedback_message]
    )


# The Gradio app is launched using `demo.launch()`. For Hugging Face Spaces,
# this is handled by the platform. You just need to have the `demo` object.
# The `launch()` call is fine for local testing.
if __name__ == "__main__":
    demo.launch(share=True)