import gradio as gr
import pandas as pd
from io import BytesIO
from fpdf import FPDF
import matplotlib.pyplot as plt
from groq import Groq

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

# Function to interact with GROQ API for optimization
def optimize_schedule_with_groq(schedule):
    try:
        optimized_schedule = client.optimize_schedule(schedule.to_dict(orient="records"))
        return pd.DataFrame(optimized_schedule)
    except Exception as e:
        print(f"Error using GROQ API: {e}")
        raise e

# Function to generate PDF report
def generate_pdf_report(schedule, conflicts, cost_estimates):
    pdf = FPDF()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.add_page()

    # Title
    pdf.set_font("Arial", "B", 16)
    pdf.cell(200, 10, txt="Project Schedule Optimization Report", ln=True, align="C")

    # Schedule Table
    pdf.set_font("Arial", "B", 12)
    pdf.cell(200, 10, txt="Optimized Schedule", ln=True, align="L")
    pdf.set_font("Arial", size=10)
    for index, row in schedule.iterrows():
        pdf.cell(0, 10, txt=f"{row['tasks']}: {row['start_date']} - {row['end_date']}", ln=True)

    # Conflict Details
    pdf.set_font("Arial", "B", 12)
    pdf.cell(200, 10, txt="\nConflict Details", ln=True, align="L")
    pdf.set_font("Arial", size=10)
    for conflict in conflicts:
        pdf.cell(0, 10, txt=f"{conflict}", ln=True)

    # Cost Estimates
    pdf.set_font("Arial", "B", 12)
    pdf.cell(200, 10, txt="\nCost Estimates", ln=True, align="L")
    pdf.set_font("Arial", size=10)
    for resource, cost in cost_estimates.items():
        pdf.cell(0, 10, txt=f"{resource}: ${cost:.2f}", ln=True)

    # Save PDF to a buffer
    buffer = BytesIO()
    pdf.output(buffer)
    buffer.seek(0)
    return buffer

# Generate a sample visualization
def generate_visualization(schedule):
    fig, ax = plt.subplots(figsize=(10, 6))
    tasks = schedule['tasks']
    start_dates = pd.to_datetime(schedule['start_date'])
    end_dates = pd.to_datetime(schedule['end_date'])
    durations = (end_dates - start_dates).dt.days

    ax.barh(tasks, durations, left=start_dates.map(lambda x: x.toordinal()), color='skyblue')
    ax.set_xlabel("Dates")
    ax.set_ylabel("Tasks")
    ax.set_title("Gantt Chart (Simplified)")

    # Convert figure to BytesIO object
    buf = BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    return buf

# Main function
def process_schedule(file):
    schedule = pd.read_csv(file)

    # Basic checks
    if 'tasks' not in schedule.columns:
        return "Error: 'tasks' column is mandatory.", None

    # Infer missing columns (dummy inference for demonstration)
    if 'start_date' not in schedule.columns:
        schedule['start_date'] = pd.date_range("2025-01-01", periods=len(schedule))
    if 'end_date' not in schedule.columns:
        schedule['end_date'] = schedule['start_date'] + pd.to_timedelta(7, unit='d')
    if 'required_resources' not in schedule.columns:
        schedule['required_resources'] = ["Labor"] * len(schedule)

    # Use GROQ API to optimize the schedule
    optimized_schedule = optimize_schedule_with_groq(schedule)

    # Placeholder conflict and cost calculations
    conflicts = ["Task 1 and Task 2 overlap.", "Resource 'Crane' exceeds availability."]
    cost_estimates = {"Labor": 5000, "Equipment": 2000}

    # Generate PDF Report
    pdf_report = generate_pdf_report(optimized_schedule, conflicts, cost_estimates)

    # Generate Visualization
    gantt_chart = generate_visualization(optimized_schedule)

    return pdf_report, gantt_chart

# Gradio interface
iface = gr.Interface(
    fn=process_schedule,
    inputs=gr.File(label="Upload Schedule File (CSV)"),
    outputs=[
        gr.File(label="Download PDF Report"),
        gr.Image(label="Visualization (Gantt Chart)")
    ],
    title="Intelligent Resource Loading in Construction Schedule",
    description="Upload a schedule file to generate a PDF report with optimized schedule, conflict details, cost estimates, and visualizations."
)

if __name__ == "__main__":
    iface.launch()