app.py

import pandas as pd
import numpy as np
import pulp as pl
import matplotlib.pyplot as plt
import gradio as gr
from itertools import product
import io
import base64
import tempfile
import os
from datetime import datetime

def am_pm(hour):
    """Converts 24-hour time to AM/PM format."""
    period = "AM"
    if hour >= 12:
        period = "PM"
    if hour > 12:
        hour -= 12
    elif hour == 0:
        hour = 12  # Midnight
    return f"{int(hour):02d}:00 {period}"

def show_dataframe(csv_path):
    """Reads a CSV file and returns a Pandas DataFrame."""
    try:
        df = pd.read_csv(csv_path)
        return df
    except Exception as e:
        return f"Error loading CSV: {e}"

def optimize_staffing(
    csv_file,
    beds_per_staff,
    max_hours_per_staff,  # This will now be interpreted as hours per 28-day period
    hours_per_cycle,
    rest_days_per_week,
    clinic_start,
    clinic_end,
    overlap_time,
    max_start_time_change,
    exact_staff_count=None,
    overtime_percent=100
):
    # Load data
    try:
        if isinstance(csv_file, str):
            # Handle the case when a filepath is passed directly
            data = pd.read_csv(csv_file)
        else:
            # Handle the case when file object is uploaded through Gradio
            data = pd.read_csv(io.StringIO(csv_file.decode('utf-8')))
    except Exception as e:
        print(f"Error loading CSV file: {e}")
        return f"Error loading CSV file: {e}", None, None, None, None
    
    # Rename the index column if necessary
    if data.columns[0] not in ['day', 'Day', 'DAY']:
        data = data.rename(columns={data.columns[0]: 'day'})
    
    # Fill missing values
    for col in data.columns:
        if col.startswith('cycle'):
            data[col] = data[col].fillna(0)
    
    # Calculate clinic hours
    if clinic_end < clinic_start:  # overnight clinic (e.g., 7 AM to 3 AM next day)
        clinic_hours = 24 - clinic_start + clinic_end
    else:
        clinic_hours = clinic_end - clinic_start
    
    # Get number of days in the dataset
    num_days = len(data)
    
    # Parameters
    BEDS_PER_STAFF = float(beds_per_staff)
    STANDARD_PERIOD_DAYS = 30  # Standard month period (changed from 28 to 30)
    
    # Scale MAX_HOURS_PER_STAFF based on the ratio of actual days to standard month
    BASE_MAX_HOURS = float(max_hours_per_staff)  # This is for a 30-day period
    MAX_HOURS_PER_STAFF = BASE_MAX_HOURS * (num_days / STANDARD_PERIOD_DAYS)
    
    # Log the adjustment for transparency
    original_results = f"Input max hours per staff (30-day period): {BASE_MAX_HOURS}\n"
    original_results += f"Adjusted max hours for {num_days}-day period: {MAX_HOURS_PER_STAFF:.1f}\n"
    original_results += f"(Adjustment ratio: {num_days}/{STANDARD_PERIOD_DAYS} = {(num_days/STANDARD_PERIOD_DAYS):.2f})\n\n"
    
    HOURS_PER_CYCLE = float(hours_per_cycle)
    REST_DAYS_PER_WEEK = int(rest_days_per_week)
    SHIFT_TYPES = [6, 8, 10, 12]  # Standard shift types
    OVERLAP_TIME = float(overlap_time)
    CLINIC_START = int(clinic_start)
    CLINIC_END = int(clinic_end)
    CLINIC_HOURS = clinic_hours
    MAX_START_TIME_CHANGE = int(max_start_time_change)
    OVERTIME_ALLOWED = 1 + (overtime_percent / 100)  # Convert percentage to multiplier
    
    # Calculate staff needed per cycle (beds/BEDS_PER_STAFF, rounded up)
    for col in data.columns:
        if col.startswith('cycle') and not col.endswith('_staff'):
            data[f'{col}_staff'] = np.ceil(data[col] / BEDS_PER_STAFF)
    
    # Get cycle names and number of cycles
    cycle_cols = [col for col in data.columns if col.startswith('cycle') and not col.endswith('_staff')]
    num_cycles = len(cycle_cols)
    
    # Define cycle times - adjusted for overnight clinic
    cycle_times = {}
    for i, cycle in enumerate(cycle_cols):
        # Ensure first cycle starts exactly at clinic start time
        cycle_start = CLINIC_START if i == 0 else (CLINIC_START + i * HOURS_PER_CYCLE) % 24
        cycle_end = (cycle_start + HOURS_PER_CYCLE) % 24
        cycle_times[cycle] = (cycle_start, cycle_end)
    
    # Get staff requirements
    max_staff_needed = max([data[f'{cycle}_staff'].max() for cycle in cycle_cols])
    
    # Define possible shift start times for overnight clinic
    shift_start_times = []
    if CLINIC_END < CLINIC_START:  # overnight clinic
        # Always include clinic start time first to ensure coverage
        shift_start_times.append(CLINIC_START)
        # Add remaining morning shifts
        shift_start_times.extend([t for t in range(CLINIC_START + 1, 24)])
        # Add evening shifts that end next day
        shift_start_times.extend(range(0, CLINIC_END + 1))
    else:
        # Always include clinic start time first
        shift_start_times.append(CLINIC_START)
        # Add remaining times
        shift_start_times.extend([t for t in range(CLINIC_START + 1, CLINIC_END - min(SHIFT_TYPES) + 1)])
    
    # Generate all possible shifts with better overnight handling
    possible_shifts = []
    # First generate shifts starting at clinic start time
    for duration in sorted(SHIFT_TYPES, reverse=True):  # Try longer shifts first
        start_time = CLINIC_START
        end_time = (start_time + duration) % 24
        
        shift = {
            'id': f"{duration}hr_{start_time:02d}",
            'start': start_time,
            'end': end_time,
            'duration': duration,
            'cycles_covered': set()
        }
        
        # Determine which cycles this shift covers
        for cycle, (cycle_start, cycle_end) in cycle_times.items():
            # Handle overnight cycles
            if cycle_end < cycle_start:  # overnight cycle
                if start_time >= cycle_start or end_time <= cycle_end:
                    shift['cycles_covered'].add(cycle)
                elif start_time < end_time and end_time > cycle_start:
                    shift['cycles_covered'].add(cycle)
                elif end_time < start_time and (start_time < cycle_end or end_time > cycle_start):
                    shift['cycles_covered'].add(cycle)
            else:  # normal cycle
                shift_end = end_time if end_time > start_time else end_time + 24
                cycle_end_adj = cycle_end if cycle_end > cycle_start else cycle_end + 24
                
                # Check for overlap
                if not (shift_end <= cycle_start or start_time >= cycle_end_adj):
                    shift['cycles_covered'].add(cycle)
        
        if shift['cycles_covered']:  # Only add shifts that cover at least one cycle
            possible_shifts.append(shift)
    
    # Then generate remaining shifts
    for duration in SHIFT_TYPES:
        for start_time in shift_start_times:
            if start_time == CLINIC_START:  # Skip as we already handled clinic start time
                continue
                
            end_time = (start_time + duration) % 24
            
            # Skip shifts that don't align with clinic hours
            if CLINIC_END < CLINIC_START:  # overnight clinic
                if start_time < CLINIC_START and start_time > CLINIC_END:
                    continue
                if (start_time + duration) % 24 < CLINIC_START and (start_time + duration) % 24 > CLINIC_END:
                    continue
            else:
                if start_time < CLINIC_START or end_time > CLINIC_END:
                    continue
            
            shift = {
                'id': f"{duration}hr_{start_time:02d}",
                'start': start_time,
                'end': end_time,
                'duration': duration,
                'cycles_covered': set()
            }
            
            # Determine which cycles this shift covers
            for cycle, (cycle_start, cycle_end) in cycle_times.items():
                if cycle_end < cycle_start:  # overnight cycle
                    if start_time >= cycle_start or end_time <= cycle_end:
                        shift['cycles_covered'].add(cycle)
                    elif start_time < end_time and end_time > cycle_start:
                        shift['cycles_covered'].add(cycle)
                    elif end_time < start_time and (start_time < cycle_end or end_time > cycle_start):
                        shift['cycles_covered'].add(cycle)
                else:  # normal cycle
                    shift_end = end_time if end_time > start_time else end_time + 24
                    cycle_end_adj = cycle_end if cycle_end > cycle_start else cycle_end + 24
                    
                    if not (shift_end <= cycle_start or start_time >= cycle_end_adj):
                        shift['cycles_covered'].add(cycle)
            
            if shift['cycles_covered']:  # Only add shifts that cover at least one cycle
                possible_shifts.append(shift)
    
    # Estimate minimum number of staff needed - more precise calculation
    total_staff_hours = 0
    for _, row in data.iterrows():
        for cycle in cycle_cols:
            total_staff_hours += row[f'{cycle}_staff'] * HOURS_PER_CYCLE
    
    # Calculate theoretical minimum staff with perfect utilization
    theoretical_min_staff = np.ceil(total_staff_hours / MAX_HOURS_PER_STAFF)
    
    # Add a small buffer for rest day constraints
    min_staff_estimate = np.ceil(theoretical_min_staff * (7 / (7 - REST_DAYS_PER_WEEK)))
    
    # Use exact_staff_count if provided, otherwise estimate
    if exact_staff_count is not None and exact_staff_count > 0:
        # When exact staff count is provided, use it regardless of minimum required
        if exact_staff_count < min_staff_estimate:
            original_results += f"\nWarning: Provided staff count ({exact_staff_count}) is below estimated minimum ({min_staff_estimate:.1f}). Solution may not be feasible.\n"
        estimated_staff = exact_staff_count
        num_staff_to_create = exact_staff_count
    else:
        # Add some buffer for constraints like rest days and shift changes
        estimated_staff = max(min_staff_estimate, max_staff_needed + 1)
        num_staff_to_create = int(estimated_staff)
    
    def optimize_schedule(num_staff, time_limit=600):
        try:
            # Create a binary linear programming model
            model = pl.LpProblem("Staff_Scheduling", pl.LpMinimize)
            
            # Decision variables
            x = pl.LpVariable.dicts("shift", 
                                   [(s, d, shift['id']) for s in range(1, num_staff+1) 
                                                        for d in range(1, num_days+1) 
                                                        for shift in possible_shifts],
                                   cat='Binary')
            
            # Staff usage variable (1 if staff s is used at all, 0 otherwise)
            staff_used = pl.LpVariable.dicts("staff_used", range(1, num_staff+1), cat='Binary')
            
            # Total hours worked by all staff
            total_hours = pl.LpVariable("total_hours", lowBound=0)
            
            # Individual staff hours variables for balancing
            staff_hours = pl.LpVariable.dicts("staff_hours", range(1, num_staff+1), lowBound=0)
            
            # Objective function modification for exact staff count
            if exact_staff_count is not None:
                # When exact staff count is specified, focus on balancing hours between staff
                avg_hours = total_staff_hours / num_staff
                model += pl.lpSum(staff_hours[s] for s in range(1, num_staff+1))
                
                # Add penalty for deviation from average
                for s in range(1, num_staff+1):
                    model += staff_hours[s] >= avg_hours * 0.8  # Each staff must get at least 80% of average hours
                    model += staff_hours[s] <= avg_hours * 1.2  # Each staff must not exceed 120% of average hours
            else:
                # Original objective for minimizing staff and total hours
                model += 10**10 * pl.lpSum(staff_used[s] for s in range(1, num_staff+1)) + total_hours
            
            # Link staff_hours to actual hours worked
            for s in range(1, num_staff+1):
                model += staff_hours[s] == pl.lpSum(x[(s, d, shift['id'])] * shift['duration'] 
                                                  for d in range(1, num_days+1) 
                                                  for shift in possible_shifts)
            
            # Link total_hours to sum of staff_hours
            model += total_hours == pl.lpSum(staff_hours[s] for s in range(1, num_staff+1))
            
            # When exact staff count is provided, ensure all staff are used
            if exact_staff_count is not None:
                for s in range(1, num_staff+1):
                    # Ensure each staff works at least some minimum shifts
                    min_shifts = max(1, int(num_days / (num_staff * 2)))  # At least this many shifts per staff
                    model += pl.lpSum(x[(s, d, shift['id'])] 
                                    for d in range(1, num_days+1) 
                                    for shift in possible_shifts) >= min_shifts
                    
                    # Maximum shifts per staff (to prevent overloading)
                    max_shifts = int(num_days * 0.8)  # At most 80% of days
                    model += pl.lpSum(x[(s, d, shift['id'])] 
                                    for d in range(1, num_days+1) 
                                    for shift in possible_shifts) <= max_shifts
            else:
                # Original staff usage constraints
                for s in range(1, num_staff+1):
                    model += pl.lpSum(x[(s, d, shift['id'])] 
                                     for d in range(1, num_days+1) 
                                     for shift in possible_shifts) <= num_days * staff_used[s]
                    model += pl.lpSum(x[(s, d, shift['id'])] 
                                     for d in range(1, num_days+1) 
                                     for shift in possible_shifts) >= staff_used[s]
                
                # Maintain staff ordering only when not using exact staff count
                for s in range(1, num_staff):
                    model += staff_used[s] >= staff_used[s+1]
            
            # Each staff works at most one shift per day
            for s in range(1, num_staff+1):
                for d in range(1, num_days+1):
                    model += pl.lpSum(x[(s, d, shift['id'])] for shift in possible_shifts) <= 1
        
            # Rest day constraints (with some flexibility)
            min_rest_days = max(1, REST_DAYS_PER_WEEK - 1)
            for s in range(1, num_staff+1):
                for w in range((num_days + 6) // 7):
                    week_start = w*7 + 1
                    week_end = min(week_start + 6, num_days)
                    days_in_this_week = week_end - week_start + 1
                    
                    if days_in_this_week < 7:
                        adjusted_rest_days = max(1, int(min_rest_days * days_in_this_week / 7))
                    else:
                        adjusted_rest_days = min_rest_days
                    
                    model += pl.lpSum(x[(s, d, shift['id'])] 
                                     for d in range(week_start, week_end+1) 
                                        for shift in possible_shifts) <= days_in_this_week - adjusted_rest_days
        
            # HARD CONSTRAINT: No overtime allowed - strict limit at MAX_HOURS_PER_STAFF
            for s in range(1, num_staff+1):
                # Calculate total hours worked by this staff
                staff_hours_value = pl.lpSum(x[(s, d, shift['id'])] * shift['duration'] 
                             for d in range(1, num_days+1) 
                                      for shift in possible_shifts)
                
                # STRICT constraint: No overtime allowed
                model += staff_hours_value <= MAX_HOURS_PER_STAFF
        
            # HARD CONSTRAINT: Full coverage required
            for d in range(1, num_days+1):
                day_index = d - 1  # 0-indexed for DataFrame
                
                for cycle in cycle_cols:
                    staff_needed = data.iloc[day_index][f'{cycle}_staff']
                    cycle_start, cycle_end = cycle_times[cycle]
                    
                    # Get all shifts that cover this cycle
                    covering_shifts = [shift for shift in possible_shifts if cycle in shift['cycles_covered']]
                    
                    # For the first cycle of the day (starting at clinic start time)
                    if cycle_start == CLINIC_START:
                        # Only consider shifts that start at clinic start time
                        early_shifts = [shift for shift in covering_shifts if shift['start'] == CLINIC_START]
                        
                        # Must have enough staff starting at clinic start time
                        model += (pl.lpSum(x[(s, d, shift['id'])] 
                                     for s in range(1, num_staff+1) 
                                         for shift in early_shifts) >= staff_needed)
                    
                    # General coverage constraint for all cycles
                    model += (pl.lpSum(x[(s, d, shift['id'])] 
                                 for s in range(1, num_staff+1) 
                                     for shift in covering_shifts) >= staff_needed)
            
            # HARD CONSTRAINT: Maximum 60 hours per week for each staff
            for s in range(1, num_staff+1):
                for w in range((num_days + 6) // 7):
                    week_start = w*7 + 1
                    week_end = min(week_start + 6, num_days)
                    
                    # Calculate total hours worked by this staff in this week
                    weekly_hours = pl.lpSum(x[(s, d, shift['id'])] * shift['duration'] 
                         for d in range(week_start, week_end+1) 
                         for shift in possible_shifts)
                    
                    # STRICT constraint: No more than 60 hours per week
                    model += weekly_hours <= 60
            
            # Solve with extended time limit
            solver = pl.PULP_CBC_CMD(timeLimit=time_limit, msg=1, gapRel=0.01)  # Tighter gap for better solutions
            model.solve(solver)
            
            # Check if a feasible solution was found
            if model.status == pl.LpStatusOptimal or model.status == pl.LpStatusNotSolved:
                # Extract the solution
                schedule = []
                for s in range(1, num_staff+1):
                    for d in range(1, num_days+1):
                        for shift in possible_shifts:
                            if pl.value(x[(s, d, shift['id'])]) == 1:
                                # Find the shift details
                                shift_details = next((sh for sh in possible_shifts if sh['id'] == shift['id']), None)
                                
                                schedule.append({
                                    'staff_id': s,
                                    'day': d,
                                    'shift_id': shift['id'],
                                    'start': shift_details['start'],
                                    'end': shift_details['end'],
                                    'duration': shift_details['duration'],
                                    'cycles_covered': list(shift_details['cycles_covered'])
                                })
                
                return schedule, model.objective.value()
            else:
                return None, None
        except Exception as e:
            print(f"Error in optimization: {e}")
            return None, None
    
    # Try to solve with estimated number of staff
    if exact_staff_count is not None and exact_staff_count > 0:
        # If exact staff count is specified, only try with that count
        staff_count = int(exact_staff_count)
        results = original_results  # Include the hours adjustment information
        results += f"\nUsing exactly {staff_count} staff as specified"
        if staff_count < min_staff_estimate:
            results += f" (Warning: This is below estimated minimum of {min_staff_estimate:.1f})"
        results += "...\n"
        
        # Try to solve with exactly this many staff
        schedule, objective = optimize_schedule(staff_count)
        
        if schedule is None:
            results += f"Failed to find a feasible solution with exactly {staff_count} staff.\n"
            if staff_count < min_staff_estimate:
                results += f"This is likely because the staff count is below the estimated minimum of {min_staff_estimate:.1f}.\n"
            results += "Try increasing the staff count or adjusting other parameters.\n"
            return results, None, None, None, None
    else:
        # Start from theoretical minimum and work up
        min_staff = max(1, int(theoretical_min_staff))  # Start from theoretical minimum
        max_staff = int(min_staff_estimate) + 5  # Allow some buffer
        
        results = original_results  # Include the hours adjustment information
        results += f"Theoretical minimum staff needed: {theoretical_min_staff:.1f}\n"
        results += f"Searching for minimum staff count starting from {min_staff}...\n"
        
        # Try each staff count from min to max
        for staff_count in range(min_staff, max_staff + 1):
            results += f"Trying with {staff_count} staff...\n"
            
            # Increase time limit for each attempt to give the solver more time
            time_limit = 300 + (staff_count - min_staff) * 100  # More time for larger staff counts
            schedule, objective = optimize_schedule(staff_count, time_limit)
            
            if schedule is not None:
                results += f"Found feasible solution with {staff_count} staff.\n"
                break
        
        if schedule is None:
            results += "Failed to find a feasible solution with the attempted staff counts.\n"
            results += "Try increasing the staff count manually or relaxing constraints.\n"
            return results, None, None, None, None
    
    results += f"Optimal solution found with {staff_count} staff\n"
    results += f"Total staff hours: {objective}\n"
    
    # Convert to DataFrame for analysis
    schedule_df = pd.DataFrame(schedule)
    
    # Analyze staff workload
    staff_hours = {}
    for s in range(1, staff_count+1):
        staff_shifts = schedule_df[schedule_df['staff_id'] == s]
        total_hours = staff_shifts['duration'].sum()
        staff_hours[s] = total_hours
    
    # Handle staff hours display based on whether exact count was specified
    if exact_staff_count is not None:
        # When exact count is specified, show all staff including those with 0 hours
        active_staff_hours = staff_hours
    else:
        # Otherwise, only show active staff
        active_staff_hours = {s: hours for s, hours in staff_hours.items() if hours > 0}
    
    results += "\nStaff Hours:\n"
    total_active_hours = sum(active_staff_hours.values())
    avg_hours = total_active_hours / len(active_staff_hours) if active_staff_hours else 0
    
    for staff_id, hours in active_staff_hours.items():
        utilization = (hours / MAX_HOURS_PER_STAFF) * 100
        deviation_from_avg = ((hours - avg_hours) / avg_hours * 100) if avg_hours > 0 else 0
        results += f"Staff {staff_id}: {hours:.1f} hours ({utilization:.1f}% utilization)"
        if exact_staff_count is not None:
            results += f" [Deviation from avg: {deviation_from_avg:+.1f}%]"
        results += "\n"
        
        # Add overtime information
        if hours > MAX_HOURS_PER_STAFF:
            overtime = hours - MAX_HOURS_PER_STAFF
            overtime_percent = (overtime / MAX_HOURS_PER_STAFF) * 100
            results += f"  Overtime: {overtime:.1f} hours ({overtime_percent:.1f}%)\n"
    
    if exact_staff_count is not None:
        results += f"\nWorkload Distribution Stats:\n"
        results += f"Average hours per staff: {avg_hours:.1f}\n"
        if active_staff_hours:
            max_deviation = max(abs((hours - avg_hours) / avg_hours * 100) for hours in active_staff_hours.values()) if avg_hours > 0 else 0
            results += f"Maximum deviation from average: {max_deviation:.1f}%\n"
    
    # Use active_staff_hours for average utilization calculation
    active_staff_count = len(active_staff_hours)
    avg_utilization = sum(active_staff_hours.values()) / (active_staff_count * MAX_HOURS_PER_STAFF) * 100
    results += f"\nAverage staff utilization: {avg_utilization:.1f}%\n"
    
    # Check coverage for each day and cycle
    coverage_check = []
    for d in range(1, num_days+1):
        day_index = d - 1  # 0-indexed for DataFrame
        
        day_schedule = schedule_df[schedule_df['day'] == d]
        
        for cycle in cycle_cols:
            required = data.iloc[day_index][f'{cycle}_staff']
            
            # Count staff covering this cycle
            assigned = sum(1 for _, shift in day_schedule.iterrows() 
                          if cycle in shift['cycles_covered'])
            
            coverage_check.append({
                'day': d,
                'cycle': cycle,
                'required': required,
                'assigned': assigned,
                'satisfied': assigned >= required
            })
    
    coverage_df = pd.DataFrame(coverage_check)
    satisfaction = coverage_df['satisfied'].mean() * 100
    results += f"Coverage satisfaction: {satisfaction:.1f}%\n"
    
    if satisfaction < 100:
        results += "Warning: Not all staffing requirements are met!\n"
        unsatisfied = coverage_df[~coverage_df['satisfied']]
        results += unsatisfied.to_string() + "\n"
    
    # Generate detailed schedule report
    detailed_schedule = "Detailed Schedule:\n"
    for d in range(1, num_days+1):
        day_schedule = schedule_df[schedule_df['day'] == d]
        day_schedule = day_schedule.sort_values(['start'])
        
        detailed_schedule += f"\nDay {d}:\n"
        for _, shift in day_schedule.iterrows():
            start_hour = shift['start']
            end_hour = shift['end']

            start_str = am_pm(start_hour)
            end_str = am_pm(end_hour)

            cycles = ", ".join(shift['cycles_covered'])
            detailed_schedule += f"  Staff {shift['staff_id']}: {start_str}-{end_str} ({shift['duration']} hrs), Cycles: {cycles}\n"
    
    # Generate schedule visualization
    fig, ax = plt.subplots(figsize=(15, 8))
    
    # Prepare schedule for plotting
    staff_days = {}
    for s in range(1, staff_count+1):
        staff_days[s] = [0] * num_days  # 0 means off duty
    
    for _, shift in schedule_df.iterrows():
        staff_id = shift['staff_id']
        day = shift['day'] - 1  # 0-indexed
        staff_days[staff_id][day] = shift['duration']
    
    # Plot the schedule
    for s, hours in staff_days.items():
        ax.bar(range(1, num_days+1), hours, label=f'Staff {s}')
    
    ax.set_xlabel('Day')
    ax.set_ylabel('Shift Hours')
    ax.set_title('Staff Schedule')
    ax.set_xticks(range(1, num_days+1))
    ax.legend()
    
    # Save the figure to a temporary file
    plot_path = None
    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
        plt.savefig(f.name)
        plt.close(fig)
        plot_path = f.name
    
    # Create a Gantt chart with advanced visuals and alternating labels - only showing active staff
    gantt_path = create_gantt_chart(schedule_df, num_days, staff_count)

    # Convert schedule to CSV data
    schedule_df['start_ampm'] = schedule_df['start'].apply(am_pm)
    schedule_df['end_ampm'] = schedule_df['end'].apply(am_pm)
    schedule_csv = schedule_df[['staff_id', 'day', 'start_ampm', 'end_ampm', 'duration', 'cycles_covered']].to_csv(index=False)

    # Create a temporary file and write the CSV data into it
    with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix=".csv") as temp_file:
        temp_file.write(schedule_csv)
        schedule_csv_path = temp_file.name

    # Create staff assignment table
    staff_assignment_data = []
    for d in range(1, num_days + 1):
        cycle_staff = {}
        for cycle in cycle_cols:
            # Get staff IDs assigned to this cycle on this day
            staff_ids = schedule_df[(schedule_df['day'] == d) & (schedule_df['cycles_covered'].apply(lambda x: cycle in x))]['staff_id'].tolist()
            cycle_staff[cycle] = len(staff_ids)
        staff_assignment_data.append([d] + [cycle_staff[cycle] for cycle in cycle_cols])

    staff_assignment_df = pd.DataFrame(staff_assignment_data, columns=['Day'] + cycle_cols)
    
    # Create CSV files for download
    staff_assignment_csv_path = None
    with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix=".csv") as temp_file:
        staff_assignment_df.to_csv(temp_file.name, index=False)
        staff_assignment_csv_path = temp_file.name
    
    # Return all required values in the correct order
    return results, staff_assignment_df, gantt_path, schedule_df, plot_path, schedule_csv_path, staff_assignment_csv_path

def convert_to_24h(time_str):
    """Converts AM/PM time string to 24-hour format."""
    try:
        time_obj = datetime.strptime(time_str, "%I:00 %p")
        return time_obj.hour
    except ValueError:
        return None

def gradio_wrapper(
    csv_file, beds_per_staff, max_hours_per_staff, hours_per_cycle,
    rest_days_per_week, clinic_start_ampm, clinic_end_ampm, overlap_time, max_start_time_change,
    exact_staff_count=None, overtime_percent=100
):
    try:
        # Convert AM/PM times to 24-hour format
        clinic_start = convert_to_24h(clinic_start_ampm)
        clinic_end = convert_to_24h(clinic_end_ampm)
        
        # Call the optimization function
        results, staff_assignment_df, gantt_path, schedule_df, plot_path, schedule_csv_path, staff_assignment_csv_path = optimize_staffing(
            csv_file, beds_per_staff, max_hours_per_staff, hours_per_cycle,
            rest_days_per_week, clinic_start, clinic_end, overlap_time, max_start_time_change,
            exact_staff_count, overtime_percent
        )
        
        # Return the results
        return staff_assignment_df, gantt_path, schedule_df, plot_path, staff_assignment_csv_path, schedule_csv_path
    except Exception as e:
        # If there's an error in the optimization process, return a meaningful error message
        empty_staff_df = pd.DataFrame(columns=["Day"])
        error_message = f"Error during optimization: {str(e)}\n\nPlease try with different parameters or a simpler dataset."
        # Return error in the first output
        return empty_staff_df, None, None, None, None, None

# Create a Gantt chart with advanced visuals and alternating labels - only showing active staff
def create_gantt_chart(schedule_df, num_days, staff_count):
    # Get the list of active staff IDs (staff who have at least one shift)
    active_staff_ids = sorted(schedule_df['staff_id'].unique())
    active_staff_count = len(active_staff_ids)
    
    # Create a mapping from original staff ID to position in the chart
    staff_position = {staff_id: i+1 for i, staff_id in enumerate(active_staff_ids)}
    
    # Create a larger figure with higher DPI
    plt.figure(figsize=(max(30, num_days * 1.5), max(12, active_staff_count * 0.8)), dpi=200)
    
    # Use a more sophisticated color palette - only for active staff
    colors = plt.cm.viridis(np.linspace(0.1, 0.9, active_staff_count))
    
    # Set a modern style
    plt.style.use('seaborn-v0_8-whitegrid')
    
    # Create a new axis with a slight background color
    ax = plt.gca()
    ax.set_facecolor('#f8f9fa')
    
    # Sort by staff then day
    schedule_df = schedule_df.sort_values(['staff_id', 'day'])
    
    # Plot Gantt chart - only for active staff
    for i, staff_id in enumerate(active_staff_ids):
        staff_shifts = schedule_df[schedule_df['staff_id'] == staff_id]
        
        y_pos = active_staff_count - i  # Position based on index in active staff list
        
        # Add staff label with a background box
        ax.text(-0.7, y_pos, f"Staff {staff_id}", fontsize=12, fontweight='bold', 
                ha='right', va='center', bbox=dict(facecolor='white', edgecolor='gray', 
                                                  boxstyle='round,pad=0.5', alpha=0.9))
        
        # Add a subtle background for each staff row
        ax.axhspan(y_pos-0.4, y_pos+0.4, color='white', alpha=0.4, zorder=-5)
        
        # Track shift positions to avoid label overlap
        shift_positions = []
        
        for idx, shift in enumerate(staff_shifts.iterrows()):
            _, shift = shift
            day = shift['day']
            start_hour = shift['start']
            end_hour = shift['end']
            duration = shift['duration']
            
            # Format times for display
            start_ampm = am_pm(start_hour)
            end_ampm = am_pm(end_hour)
            
            # Calculate shift position
            shift_start_pos = day-1+start_hour/24
            
            # Handle overnight shifts
            if end_hour < start_hour:  # Overnight shift
                # First part of shift (until midnight)
                rect1 = ax.barh(y_pos, (24-start_hour)/24, left=shift_start_pos, 
                       height=0.6, color=colors[i], alpha=0.9, 
                       edgecolor='black', linewidth=1, zorder=10)
                
                # Add gradient effect
                for r in rect1:
                    r.set_edgecolor('black')
                    r.set_linewidth(1)
                
                # Second part of shift (after midnight)
                rect2 = ax.barh(y_pos, end_hour/24, left=day, 
                       height=0.6, color=colors[i], alpha=0.9,
                       edgecolor='black', linewidth=1, zorder=10)
                
                # Add gradient effect
                for r in rect2:
                    r.set_edgecolor('black')
                    r.set_linewidth(1)
                
                # For overnight shifts, we'll place the label in the first part if it's long enough
                shift_width = (24-start_hour)/24
                if shift_width >= 0.1:  # Only add label if there's enough space
                    label_pos = shift_start_pos + shift_width/2
                    
                    # Alternate labels above and below
                    y_offset = 0.35 if idx % 2 == 0 else -0.35
                    
                    # Add label with background for better readability
                    label = f"{start_ampm}-{end_ampm}"
                    text = ax.text(label_pos, y_pos + y_offset, label, 
                           ha='center', va='center', fontsize=9, fontweight='bold',
                           color='black', bbox=dict(facecolor='white', alpha=0.9, pad=3, 
                                                   boxstyle='round,pad=0.3', edgecolor='gray'),
                           zorder=20)
                    
                    shift_positions.append(label_pos)
            else:
                # Regular shift
                shift_width = duration/24
                rect = ax.barh(y_pos, shift_width, left=shift_start_pos, 
                       height=0.6, color=colors[i], alpha=0.9,
                       edgecolor='black', linewidth=1, zorder=10)
                
                # Add gradient effect
                for r in rect:
                    r.set_edgecolor('black')
                    r.set_linewidth(1)
                
                # Only add label if there's enough space
                if shift_width >= 0.1:
                    label_pos = shift_start_pos + shift_width/2
                    
                    # Alternate labels above and below
                    y_offset = 0.35 if idx % 2 == 0 else -0.35
                    
                    # Add label with background for better readability
                    label = f"{start_ampm}-{end_ampm}"
                    text = ax.text(label_pos, y_pos + y_offset, label, 
                           ha='center', va='center', fontsize=9, fontweight='bold',
                           color='black', bbox=dict(facecolor='white', alpha=0.9, pad=3, 
                                                   boxstyle='round,pad=0.3', edgecolor='gray'),
                           zorder=20)
                    
                    shift_positions.append(label_pos)
    
    # Add weekend highlighting with a more sophisticated look
    for day in range(1, num_days + 1):
        # Determine if this is a weekend (assuming day 1 is Monday)
        is_weekend = (day % 7 == 0) or (day % 7 == 6)  # Saturday or Sunday
        
        if is_weekend:
            ax.axvspan(day-1, day, alpha=0.15, color='#ff9999', zorder=-10)
            day_label = "Saturday" if day % 7 == 6 else "Sunday"
            ax.text(day-0.5, 0.2, day_label, ha='center', fontsize=10, color='#cc0000',
                   fontweight='bold', bbox=dict(facecolor='white', alpha=0.7, pad=2, boxstyle='round'))
    
    # Set x-axis ticks for each day with better formatting
    ax.set_xticks(np.arange(0.5, num_days, 1))
    day_labels = [f"Day {d}" for d in range(1, num_days+1)]
    ax.set_xticklabels(day_labels, rotation=0, ha='center', fontsize=10)
    
    # Add vertical lines between days with better styling
    for day in range(1, num_days):
        ax.axvline(x=day, color='#aaaaaa', linestyle='-', alpha=0.5, zorder=-5)
    
    # Set y-axis ticks for each staff
    ax.set_yticks(np.arange(1, active_staff_count+1))
    ax.set_yticklabels([])  # Remove default labels as we've added custom ones
    
    # Set axis limits with some padding
    ax.set_xlim(-0.8, num_days)
    ax.set_ylim(0.5, active_staff_count + 0.5)
    
    # Add grid for hours (every 6 hours) with better styling
    for day in range(num_days):
        for hour in [6, 12, 18]:
            ax.axvline(x=day + hour/24, color='#cccccc', linestyle=':', alpha=0.5, zorder=-5)
            # Add small hour markers at the bottom
            hour_label = "6AM" if hour == 6 else "Noon" if hour == 12 else "6PM"
            ax.text(day + hour/24, 0, hour_label, ha='center', va='bottom', fontsize=7, 
                   color='#666666', rotation=90, alpha=0.7)
    
    # Add title and labels with more sophisticated styling
    plt.title(f'Staff Schedule ({active_staff_count} Active Staff)', fontsize=24, fontweight='bold', pad=20, color='#333333')
    plt.xlabel('Day', fontsize=16, labelpad=10, color='#333333')
    
    # Add a legend for time reference with better styling
    time_box = plt.figtext(0.01, 0.01, "Time Reference:", ha='left', fontsize=10, 
                          fontweight='bold', color='#333333')
    time_markers = ['6 AM', 'Noon', '6 PM', 'Midnight']
    for i, time in enumerate(time_markers):
        plt.figtext(0.08 + i*0.06, 0.01, time, ha='left', fontsize=9, color='#555555')
    
    # Remove spines
    for spine in ['top', 'right', 'left']:
        ax.spines[spine].set_visible(False)
    
    # Add a note about weekends with better styling
    weekend_note = plt.figtext(0.01, 0.97, "Red areas = Weekends", fontsize=12, 
                              color='#cc0000', fontweight='bold',
                              bbox=dict(facecolor='white', alpha=0.7, pad=5, boxstyle='round'))
    
    # Add a subtle border around the entire chart
    plt.box(False)
    
    # Save the Gantt chart with high quality
    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
        plt.tight_layout()
        plt.savefig(f.name, dpi=200, bbox_inches='tight', facecolor='white')
        plt.close()
        return f.name

# Define Gradio UI
am_pm_times = [f"{i:02d}:00 AM" for i in range(1, 13)] + [f"{i:02d}:00 PM" for i in range(1, 13)]

with gr.Blocks(title="Staff Scheduling Optimizer", css="""
    #staff_assignment_table {
        width: 100% !important;
    }
    #csv_schedule {
        width: 100% !important;
    }
    .container {
        max-width: 100% !important;
        padding: 0 !important;
    }
    .download-btn {
        margin-top: 10px !important;
    }
    """) as iface:
    
    gr.Markdown("# Staff Scheduling Optimizer")
    gr.Markdown("Upload a CSV file with cycle data and configure parameters to generate an optimal staff schedule.")
    
    with gr.Row():
        # LEFT PANEL - Inputs
        with gr.Column(scale=1):
            gr.Markdown("### Input Parameters")
            
            # Input parameters
            csv_input = gr.File(label="Upload CSV")
            beds_per_staff = gr.Number(label="Beds per Staff", value=3)
            max_hours_per_staff = gr.Number(label="Maximum monthly hours", value=160)
            hours_per_cycle = gr.Number(label="Hours per Cycle", value=4)
            rest_days_per_week = gr.Number(label="Rest Days per Week", value=2)
            clinic_start_ampm = gr.Dropdown(label="Clinic Start Hour (AM/PM)", choices=am_pm_times, value="08:00 AM")
            clinic_end_ampm = gr.Dropdown(label="Clinic End Hour (AM/PM)", choices=am_pm_times, value="08:00 PM")
            overlap_time = gr.Number(label="Overlap Time", value=0)
            max_start_time_change = gr.Number(label="Max Start Time Change", value=2)
            exact_staff_count = gr.Number(label="Exact Staff Count (optional)", value=None)
            overtime_percent = gr.Slider(label="Overtime Allowed (%)", minimum=0, maximum=100, value=100, step=10)
            
            optimize_btn = gr.Button("Optimize Schedule", variant="primary", size="lg")
        
        # RIGHT PANEL - Outputs
        with gr.Column(scale=2):
            gr.Markdown("### Results")
            
            # Tabs for different outputs - reordered
            with gr.Tabs():
                with gr.TabItem("Detailed Schedule"):
                    with gr.Row():
                        csv_schedule = gr.Dataframe(label="Detailed Schedule", elem_id="csv_schedule")
                    
                    with gr.Row():
                        schedule_download_file = gr.File(label="Download Detailed Schedule", visible=True)
                
                with gr.TabItem("Gantt Chart"):
                    gantt_chart = gr.Image(label="Staff Schedule Visualization", elem_id="gantt_chart")
                
                with gr.TabItem("Staff Coverage by Cycle"):
                    with gr.Row():
                        staff_assignment_table = gr.Dataframe(label="Staff Count in Each Cycle (Staff May Overlap)", elem_id="staff_assignment_table")
                    
                    with gr.Row():
                        staff_download_file = gr.File(label="Download Coverage Table", visible=True)
                
                with gr.TabItem("Hours Visualization"):
                    schedule_visualization = gr.Image(label="Hours by Day Visualization", elem_id="schedule_visualization")
    
    # Define download functions
    def create_download_link(df, filename="data.csv"):
        """Create a CSV download link for a dataframe"""
        if df is None or df.empty:
            return None
        
        csv_data = df.to_csv(index=False)
        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.csv') as f:
            f.write(csv_data)
            return f.name

    # Update the optimize_and_display function
    def optimize_and_display(csv_file, beds_per_staff, max_hours_per_staff, hours_per_cycle,
                            rest_days_per_week, clinic_start_ampm, clinic_end_ampm, 
                            overlap_time, max_start_time_change, exact_staff_count, overtime_percent):
        try:
            # Convert AM/PM times to 24-hour format
            clinic_start = convert_to_24h(clinic_start_ampm)
            clinic_end = convert_to_24h(clinic_end_ampm)
            
            # Call the optimization function
            results, staff_assignment_df, gantt_path, schedule_df, plot_path, schedule_csv_path, staff_assignment_csv_path = optimize_staffing(
                csv_file, beds_per_staff, max_hours_per_staff, hours_per_cycle,
                rest_days_per_week, clinic_start, clinic_end, overlap_time, max_start_time_change,
                exact_staff_count, overtime_percent
            )
            
            # Return the results
            return staff_assignment_df, gantt_path, schedule_df, plot_path, staff_assignment_csv_path, schedule_csv_path
        except Exception as e:
            # If there's an error in the optimization process, return a meaningful error message
            empty_staff_df = pd.DataFrame(columns=["Day"])
            error_message = f"Error during optimization: {str(e)}\n\nPlease try with different parameters or a simpler dataset."
            # Return error in the first output
            return empty_staff_df, None, None, None, None, None
    
    # Connect the button to the optimization function
    optimize_btn.click(
        fn=optimize_and_display,
        inputs=[
            csv_input, beds_per_staff, max_hours_per_staff, hours_per_cycle,
            rest_days_per_week, clinic_start_ampm, clinic_end_ampm, 
            overlap_time, max_start_time_change, exact_staff_count, overtime_percent
        ],
        outputs=[
            staff_assignment_table, gantt_chart, csv_schedule, schedule_visualization, 
            staff_download_file, schedule_download_file
        ]
    )

# Launch the Gradio app
iface.launch(share=True)