app.py

import gradio as gr
import plotly.express as px
import pandas as pd
from datetime import datetime, timedelta

def get_schedule_data(start_hour, end_hour, uptake_minutes, scan_minutes):
    """
    Calculates the schedule and returns a Pandas DataFrame 
    ready for Plotly timeline plotting.
    """
    # Setup base dates (using arbitrary date for time calculation)
    base_date = datetime.now().date()
    # Handle edge case where start/end might be float or string
    start_h = int(start_hour)
    end_h = int(end_hour)
    
    work_start = datetime.combine(base_date, datetime.min.time()) + timedelta(hours=start_h)
    work_end = datetime.combine(base_date, datetime.min.time()) + timedelta(hours=end_h)
    
    # Logic: Validate inputs
    if start_h >= end_h:
        return pd.DataFrame() # Return empty if invalid

    schedule = []
    current_scan_end = None
    patient_num = 1
    
    while True:
        # 1. Determine Scan Start Time
        if patient_num == 1:
            # First patient: Scan starts after their uptake is done (uptake starts at open)
            scan_start = work_start + timedelta(minutes=uptake_minutes)
            uptake_start = work_start
        else:
            # Subsequent patients: Scan starts immediately when previous scan finishes
            scan_start = current_scan_end
            # Their uptake is back-calculated
            uptake_start = scan_start - timedelta(minutes=uptake_minutes)
        
        scan_end = scan_start + timedelta(minutes=scan_minutes)
        uptake_end = scan_start # Uptake ends when scan begins

        # 2. Check if we are past closing time
        if scan_end > work_end:
            break
            
        # 3. Add Uptake Phase to list
        schedule.append({
            "Patient": f"Patient {patient_num}",
            "Task": "Uptake",
            "Start": uptake_start,
            "Finish": uptake_end,
            "Duration": uptake_minutes
        })
        
        # 4. Add Scan Phase to list
        schedule.append({
            "Patient": f"Patient {patient_num}",
            "Task": "Scanning",
            "Start": scan_start,
            "Finish": scan_end,
            "Duration": scan_minutes
        })

        current_scan_end = scan_end
        patient_num += 1
        
    return pd.DataFrame(schedule)

def plot_schedule(start_hour, end_hour, uptake_minutes, scan_minutes):
    # Get the data
    df = get_schedule_data(start_hour, end_hour, uptake_minutes, scan_minutes)
    
    if df.empty:
        # Return a blank figure with a warning title if input is invalid
        fig = px.bar(title="Invalid Settings: Start Time must be before End Time")
        return fig

    # Calculate stats for the title
    total_patients = df['Patient'].nunique()
    # Calculate overlap rooms needed
    rooms_needed = int((uptake_minutes / scan_minutes) + 0.99) # Ceiling division

    # Create the Gantt Chart using Plotly Express
    fig = px.timeline(
        df, 
        x_start="Start", 
        x_end="Finish", 
        y="Patient", 
        color="Task",
        title=f"Throughput: {total_patients} Patients | Needs ~{rooms_needed} Uptake Rooms",
        color_discrete_map={"Uptake": "lightgreen", "Scanning": "red"},
        height=600 # Fixed height
    )

    # Customize the Layout
    fig.update_yaxes(autorange="reversed") # Patient 1 at top
    fig.update_layout(
        xaxis_title="Time of Day",
        yaxis_title="Patient Sequence",
        xaxis=dict(
            tickformat="%H:%M",  # Format as 14:00
            dtick=3600000.0,     # Tick every 1 hour (in milliseconds)
        ),
        legend_title="Phase"
    )
    
    return fig

# --- Launch Gradio ---
# We removed 'allow_flagging' to fix the error
with gr.Interface(
    fn=plot_schedule,
    inputs=[
        gr.Slider(minimum=5, maximum=12, step=1, value=7, label="Start Hour (24h)"),
        gr.Slider(minimum=13, maximum=22, step=1, value=16, label="End Hour (24h)"),
        gr.Slider(minimum=10, maximum=120, step=5, value=60, label="Uptake Time (min)"),
        gr.Slider(minimum=5, maximum=60, step=1, value=30, label="Scan Slot Time (min)"),
    ],
    outputs=gr.Plot(label="Interactive Schedule"),
    title="PET Scanner Workload Estimator"
) as app:
    app.launch()