import gradio as gr
import pandas as pd
import random
from datetime import datetime, timedelta
from groq import Groq

# Initialize Groq API
client = Groq(api_key="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")  # Replace with your Groq API key

# Predefined resource inference logic
def infer_resources(schedule):
    resource_map = {
        "Excavation": {"labor": 10, "equipment": "Excavator", "material": "Soil"},
        "Foundation": {"labor": 15, "equipment": "Concrete Mixer", "material": "Concrete"},
        "Framing": {"labor": 20, "equipment": "Cranes", "material": "Steel"},
        "Finishing": {"labor": 5, "equipment": "Hand Tools", "material": "Paint"}
    }

    inferred_resources = []
    for _, row in schedule.iterrows():
        task = row["task"]
        resources = resource_map.get(task, {"labor": 5, "equipment": "General", "material": "Standard"})
        inferred_resources.append({
            "task": task,
            "labor": resources["labor"],
            "equipment": resources["equipment"],
            "material": resources["material"]
        })

    return pd.DataFrame(inferred_resources)

# Fill missing columns
def fill_missing_columns(schedule):
    # Generate random dates if missing
    if "start_date" not in schedule.columns:
        schedule["start_date"] = [
            (datetime.now() + timedelta(days=random.randint(1, 30))).strftime("%Y-%m-%d")
            for _ in range(len(schedule))
        ]
    if "end_date" not in schedule.columns:
        schedule["end_date"] = [
            (datetime.strptime(start, "%Y-%m-%d") + timedelta(days=random.randint(5, 15))).strftime("%Y-%m-%d")
            for start in schedule["start_date"]
        ]
    return schedule

# Mock optimization logic
def mock_optimize_schedule(schedule_with_resources):
    optimized_schedule = []
    conflicts = []

    for _, row in schedule_with_resources.iterrows():
        task = row["task"]
        start_date = row["start_date"]
        end_date = row["end_date"]
        labor = row["labor"]
        equipment = row["equipment"]
        material = row["material"]

        # Check for conflicts (mock logic)
        if labor > 20:  # Example conflict condition
            conflicts.append(f"Task '{task}' exceeds labor capacity.")
        
        optimized_schedule.append({
            "task": task,
            "start_date": start_date,
            "end_date": end_date,
            "labor": labor,
            "equipment": equipment,
            "material": material,
            "conflict": "Yes" if f"Task '{task}' exceeds labor capacity." in conflicts else "No"
        })

    return pd.DataFrame(optimized_schedule), conflicts

# Main function for resource optimization
def optimize_resources(schedule_file):
    try:
        # Load schedule file
        schedule = pd.read_csv(schedule_file.name)

        # Ensure the 'task' column exists
        if "task" not in schedule.columns:
            raise ValueError("The uploaded schedule must contain a 'task' column.")

        # Fill missing columns
        schedule = fill_missing_columns(schedule)

        # Infer resources
        inferred_resources = infer_resources(schedule)
        schedule_with_resources = pd.concat([schedule, inferred_resources], axis=1)

        # Perform optimization (mocked for now)
        optimized_schedule_df, conflicts = mock_optimize_schedule(schedule_with_resources)

        return optimized_schedule_df, "\n".join(conflicts) if conflicts else "No conflicts detected."
    except Exception as e:
        return f"Error: {e}"

# Define Gradio interface
interface = gr.Interface(
    fn=optimize_resources,
    inputs=[
        gr.File(label="Upload Schedule File (CSV)")
    ],
    outputs=[
        gr.Dataframe(label="Optimized Schedule"),  # Tabular output
        gr.Textbox(label="Conflicts")  # Text output for conflict details
    ],
    title="Smart Construction Resource Loading",
    description="Upload a construction schedule with at least a 'task' column. The app will dynamically infer other details and optimize the schedule."
)

# Launch the app
if __name__ == "__main__":
    interface.launch()