import gradio as gr
import pandas as pd

# Define the LLM models and their properties
models = {
    "gpt-4o-2024-08-06": {
        "input_price_per_1M": 2.50,
        "output_price_per_1M": 10.00,
        "max_input_tokens": 128_000,
    },
    "gpt-4o-mini-2024-07-18": {
        "input_price_per_1M": 0.15,
        "output_price_per_1M": 0.600,
        "max_input_tokens": 128_000,
    },
    "Claude 3.5 Sonnet": {
        "input_price_per_1M": 3.0,
        "output_price_per_1M": 15.0,
        "max_input_tokens": 200_000,
    },
    "GPT-3.5-turbo": {
        "input_price_per_1M": 0.5,
        "output_price_per_1M": 1.5,
        "max_input_tokens": 4096,
    },
    "GPT-4": {
        "input_price_per_1M": 30.0,
        "output_price_per_1M": 60.0,
        "max_input_tokens": 8192,
    },
}

def calculate_cost(model, input_tokens, output_tokens, num_requests):
    if model not in models:
        return "Invalid model selected", 0, 0, 0

    if input_tokens > models[model]["max_input_tokens"]:
        return f"Input tokens exceed the maximum limit for {model}", 0, 0, 0

    input_cost = (input_tokens / 1_000_000) * models[model]["input_price_per_1M"] * num_requests
    output_cost = (output_tokens / 1_000_000) * models[model]["output_price_per_1M"] * num_requests
    total_cost = input_cost + output_cost

    return f"${total_cost:.6f}", input_cost, output_cost, total_cost

def compare_models(input_tokens, output_tokens, num_requests):
    results = []
    for model in models:
        total_cost_str, input_cost, output_cost, total_cost = calculate_cost(
            model, input_tokens, output_tokens, num_requests
        )
        results.append(
            {
                "Model": model,
                "Input Cost": f"${input_cost:.6f}",
                "Output Cost": f"${output_cost:.6f}",
                "Total Cost": total_cost_str,
                "Max Input Tokens": models[model]["max_input_tokens"],
                "Input Price (1M)": f"${models[model]['input_price_per_1M']:.2f}",
                "Output Price (1M)": f"${models[model]['output_price_per_1M']:.2f}",
            }
        )
    return pd.DataFrame(results)

def create_interface():
    with gr.Blocks() as interface:
        gr.Markdown("# LLM Price Comparison Tool")
        with gr.Row():
            input_tokens = gr.Number(label="Input Tokens", value=100)
            output_tokens = gr.Number(label="Output Tokens", value=100)
            num_requests = gr.Number(label="Number of Requests", value=1, step=1)
        compare_btn = gr.Button("Compare Models")
        output_table = gr.DataFrame(label="Comparison Results")

        compare_btn.click(
            fn=compare_models,
            inputs=[input_tokens, output_tokens, num_requests],
            outputs=output_table,
        )

    return interface

# Create and launch the interface
demo = create_interface()

# Hugging Face specific launch
if __name__ == "__main__":
    demo.launch()
else:
    demo.launch(share=True)