import gradio as gr

def calculate_queue_penalty(num_gpus, wait_days, price_per_hour):
    """
    Calculate the queue penalty based on the formula:
    Queue Penalty = ROUND(num_gpus * wait_days * price_per_hour * 24)
    """
    if num_gpus is None or wait_days is None or price_per_hour is None:
        return "—", "—", "—", "—"
    
    try:
        # Calculate penalty
        penalty = round(num_gpus * wait_days * price_per_hour * 24)
        
        # Calculate breakdown
        hourly_cost = num_gpus * price_per_hour
        daily_cost = hourly_cost * 24
        total_hours = wait_days * 24
        
        return (
            f"${penalty:,.2f}",
            f"${hourly_cost:.2f}/hour",
            f"${daily_cost:.2f}/day",
            f"{total_hours:.1f} hours"
        )
    except Exception as e:
        return f"Error: {str(e)}", "—", "—", "—"

# Custom CSS for better styling on white background
custom_css = """
.gradio-container {
    max-width: 900px !important;
    background: white !important;
}
.result-box {
    font-size: 48px !important;
    font-weight: bold !important;
    text-align: center !important;
    padding: 30px !important;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    color: white !important;
    border-radius: 12px !important;
    box-shadow: 0 4px 12px rgba(102, 126, 234, 0.3) !important;
    border: 3px solid #5a67d8 !important;
}
.result-box label {
    color: white !important;
}
.result-box label span {
    color: white !important;
}
.result-box textarea {
    color: white !important;
    text-align: center !important;
    font-size: 48px !important;
    font-weight: bold !important;
    background: transparent !important;
    border: none !important;
}
.stat-box {
    text-align: center !important;
    font-size: 18px !important;
    font-weight: 600 !important;
    color: #1f2937 !important;
    padding: 15px !important;
    background: #f9fafb !important;
    border-radius: 8px !important;
    border: 2px solid #e5e7eb !important;
}
.stat-box textarea {
    text-align: center !important;
    font-size: 18px !important;
    font-weight: 600 !important;
    color: #1f2937 !important;
    background: transparent !important;
    border: none !important;
}
.input-section {
    background: #ffffff !important;
    padding: 20px !important;
    border-radius: 12px !important;
    border: 2px solid #e5e7eb !important;
}
body {
    background: white !important;
}
#component-0, .gradio-container, .main {
    background: white !important;
}
"""

# Create the Gradio interface
with gr.Blocks(title="Queue Penalty Calculator", css=custom_css, theme=gr.themes.Default()) as demo:
    
    # Header
    gr.Markdown("""
    # 💰 GPU Queue Penalty Calculator
    ### Calculate the cost of waiting for GPU resources
    """)
    
    with gr.Row():
        # Left column - Inputs
        with gr.Column(scale=1, elem_classes="input-section"):
            gr.Markdown("#### 📊 Input Parameters")
            
            num_gpus = gr.Slider(
                label="🖥️ Number of GPUs",
                value=7,
                minimum=1,
                maximum=100,
                step=1,
                info="How many GPUs do you need?"
            )
            wait_days = gr.Slider(
                label="⏱️ Wait Time (Days)",
                value=3,
                minimum=0.5,
                maximum=30,
                step=0.5,
                info="How long will you wait in queue?"
            )
            price_per_hour = gr.Number(
                label="💵 GPU Price per Hour ($)",
                value=2.49,
                minimum=0,
                step=0.01,
                info="Hourly cost per GPU"
            )
        
        # Right column - Results
        with gr.Column(scale=1):
            gr.Markdown("#### 📈 Results")
            
            result = gr.Textbox(
                label="Total Queue Penalty",
                interactive=False,
                elem_classes="result-box",
                show_label=True
            )
            
            gr.Markdown("##### Cost Breakdown")
            with gr.Row():
                hourly_cost = gr.Textbox(
                    label="Hourly Cost",
                    interactive=False,
                    elem_classes="stat-box"
                )
                daily_cost = gr.Textbox(
                    label="Daily Cost",
                    interactive=False,
                    elem_classes="stat-box"
                )
            
            with gr.Row():
                total_hours = gr.Textbox(
                    label="Total Wait Time",
                    interactive=False,
                    elem_classes="stat-box"
                )
    
    # Formula explanation
    with gr.Accordion("📐 How is this calculated?", open=False):
        gr.Markdown("""
        ### Formula
        ```
        Queue Penalty = ROUND(Number of GPUs × Wait Days × Price per Hour × 24)
        ```
        
        **Breakdown:**
        - **Hourly Cost** = Number of GPUs × Price per Hour
        - **Daily Cost** = Hourly Cost × 24 hours
        - **Total Penalty** = Daily Cost × Wait Days (rounded)
        
        **Example:**
        - 7 GPUs × 3 days × $2.49/hour × 24 hours = **$1,255**
        
        This represents the total cost you incur while waiting in the queue, 
        based on the resources you need and their hourly rate.
        """)
    
    # Set up live calculation on any input change
    for input_component in [num_gpus, wait_days, price_per_hour]:
        input_component.change(
            fn=calculate_queue_penalty,
            inputs=[num_gpus, wait_days, price_per_hour],
            outputs=[result, hourly_cost, daily_cost, total_hours]
        )
    
    # Initial calculation
    demo.load(
        fn=calculate_queue_penalty,
        inputs=[num_gpus, wait_days, price_per_hour],
        outputs=[result, hourly_cost, daily_cost, total_hours]
    )

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