import gradio as gr
import models
import time
    
text = "<h1 style='text-align: center; color: midnightblue; font-size: 30px;'>TCO Comparison Calculator"
text1 = "<h1 style='text-align: center; color: midnightblue; font-size: 20px;'>First solution"
text2 = "<h1 style='text-align: center; color: midnightblue; font-size: 20px;'>Second solution"
text3 = "<h1 style='text-align: center; color: midnightblue; font-size: 20px;'>Compute and compare TCOs"
description=f"""
<p>In this demo application, we help you compare different solutions for your AI incorporation plans, such as OpenSource or SaaS.</p>
<p>First, you'll have to choose the two solutions you'd like to compare. Then, follow the instructions to select your configurations for each solution and we will compute the cost/token accordingly to them. Eventually, both solutions are compared to evaluate which one best suits your needs.</p>
"""

def compute_ratio(tco1, tco2):
        time.sleep(2)
        try:
            r = tco1 / tco2

            if r < 1:
                comparison_result = f"Second solution is {1/r:.5f} times more expensive than the first"
            elif r > 1:
                comparison_result = f"Second solution is {r:.5f} times cheaper than the first"
            else:
                comparison_result = "Both solutions will cost the same."
            return comparison_result

        except ValueError as e:
            return f"Error: {str(e)}"

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    Models: list[models.BaseTCOModel] = [models.OpenAIModel, models.OpenSourceLlama2Model, models.CohereModel]
    model_names = [Model().get_name() for Model in Models]
    gr.Markdown(value=text)
    gr.Markdown(value=description)
    
    with gr.Row():
        with gr.Column():
            gr.Markdown(value=text1)
            page1 = models.ModelPage(Models)
            dropdown = gr.Dropdown(model_names, interactive=True, label="First solution")
            page1.render()

        with gr.Column():
            gr.Markdown(value=text2)
            page2 = models.ModelPage(Models)
            dropdown2 = gr.Dropdown(model_names, interactive=True, label="Second solution")
            page2.render()
            
    dropdown.change(page1.make_model_visible, inputs=dropdown, outputs=page1.get_all_components())
    dropdown2.change(page2.make_model_visible, inputs=dropdown2, outputs=page2.get_all_components())
    
    gr.Markdown(value=text3)
    compute_tco_btn = gr.Button("Compute TCOs", size="lg")
    tco1 = gr.State()
    tco2 = gr.State()
    
    with gr.Row():  
        with gr.Column():
            tco_output = gr.Text("Output 1: ", label="TCO for the first solution")
            with gr.Accordion("Open to see the formula", open=False):
                tco_formula = gr.Markdown()
    
        with gr.Column():
            tco_output2 = gr.Text("Output 2: ", label="TCO for the second solution")
            with gr.Accordion("Open to see the formula", open=False):
                tco_formula2 = gr.Markdown()
    
    ratio = gr.Text("Ratio: ", label="Ratio of cost/token for both solutions")
    
    compute_tco_btn.click(page1.compute_cost_per_token, inputs=page1.get_all_components_for_cost_computing() + [dropdown], outputs=[tco_output, tco1, tco_formula]).then(page2.compute_cost_per_token, inputs=page2.get_all_components_for_cost_computing() + [dropdown2], outputs=[tco_output2, tco2, tco_formula2]).then(compute_ratio, inputs=[tco1, tco2], outputs=ratio)

demo.launch(debug=True)