import gradio as gr
import models
import results
import theme
    
text = "<h1 style='text-align: center; color: #333333; font-size: 40px;'>AI TCO Comparison Calculator"
text2 = "Please note that the cost/request only defines the infrastructure cost for deployment. The labor cost must be added for the whole AI model service deployment TCO."
description=f"""
<p>In this calculator, we help you compare different AI model services, such as SaaS or "Deploy yourself" solutions, based on the Total Cost of Ownership for their deployment. 😊</p> 
<p>Please note that we focus on getting the AI model service up and running, but not the additional maintenance costs that follow.🚀</p>
<p>If you want to <strong>contribute to the calculator</strong> by adding your own AI service option, follow this <a href="https://huggingface.co/spaces/mithril-security/TCO_calculator/blob/main/How_to_contribute.md">tutorial</a> 👈. </p>
"""
formula = r"""
$CR = \frac{CIT_{1K} \times IT + COT_{1K} \times OT}{1000}$  <br>
with: <br>
$CR$ = Cost per Request <br>
$CIT_{1K}$ = Cost per 1000 Input Tokens <br>
$COT_{1K}$ = Cost per 1000 Output Tokens <br>
$IT$ = Input Tokens <br>
$OT$ = Output Tokens
"""
def set_shared_data(page1, page2):
    return page1, page2

def on_use_case_change(use_case):
    if use_case == "Summarize":
        return gr.update(value=500), gr.update(value=200)
    elif use_case == "Question-Answering":
        return gr.update(value=300), gr.update(value=300)
    else:
        return gr.update(value=50), gr.update(value=10)
    
style = theme.Style()

with gr.Blocks(theme=style) as demo:
    Models: list[models.BaseTCOModel] = [models.OpenAIModelGPT4, models.OpenAIModelGPT3_5, models.CohereModel, models.DIYLlama2Model]
    model_names = [Model().get_name() for Model in Models]
    gr.Markdown(value=text)
    gr.Markdown(value=description)
    
    with gr.Row():
        with gr.Column():
            with gr.Row():
                use_case = gr.Dropdown(["Summarize", "Question-Answering", "Classification"], value="Question-Answering", label=" Describe your use case ")
            with gr.Accordion("Click here if you want to customize the number of input and output tokens per request", open=False):    
                with gr.Row():
                    input_tokens = gr.Slider(minimum=1, maximum=1000, value=300, step=1, label=" Input tokens per request", info="We suggest a value that we believe best suit your use case choice but feel free to adjust", interactive=True)
                    output_tokens = gr.Slider(minimum=1, maximum=1000, value=300, step=1, label=" Output tokens per request", info="We suggest a value that we believe best suit your use case choice but feel free to adjust", interactive=True)
                with gr.Row(visible=False):    
                    num_users = gr.Number(value="1000", interactive = True, label=" Number of users for your service ")
    
    use_case.change(on_use_case_change, inputs=use_case, outputs=[input_tokens, output_tokens])
    
    with gr.Row():
        with gr.Column():
            page1 = models.ModelPage(Models)
            dropdown = gr.Dropdown(model_names, interactive=True, label=" First AI service option ")
            with gr.Accordion("Click here for more information on the computation parameters for your first AI service option", open=False):    
                page1.render()

        with gr.Column():
            page2 = models.ModelPage(Models)
            dropdown2 = gr.Dropdown(model_names, interactive=True, label=" Second AI service option ")
            with gr.Accordion("Click here for more information on the computation parameters for your second AI service option", open=False):        
                page2.render()
    
    shared_page1, shared_page2 = set_shared_data(page1, page2)
    results.set_shared_pages(shared_page1, shared_page2)
    
    dropdown.change(page1.make_model_visible, inputs=[dropdown, use_case], outputs=page1.get_all_components())
    dropdown2.change(page2.make_model_visible, inputs=[dropdown2, use_case], outputs=page2.get_all_components())
    
    compute_tco_btn = gr.Button("Compute & Compare", size="lg", variant="primary", scale=1) 
    tco1, tco2, labor_cost1, labor_cost2, latency, latency2 = [gr.State() for _ in range(6)]
    
    with gr.Row():
        with gr.Accordion("Click here to see the cost/request computation formula", open=False):
            tco_formula = gr.Markdown(formula)
        
    with gr.Row(variant='panel'):
        with gr.Column():
            with gr.Row():
                table = gr.Markdown()
            with gr.Row():
                info = gr.Markdown(text2)
            with gr.Row():
                with gr.Column(scale=1):
                    image = gr.Image(visible=False)
                    ratio = gr.Markdown()
                with gr.Column(scale=2):
                    plot = gr.LinePlot(visible=False)
    
    compute_tco_btn.click(results.compute_cost_per_request, inputs=page1.get_all_components_for_cost_computing() + page2.get_all_components_for_cost_computing() + [dropdown, dropdown2, input_tokens, output_tokens], outputs=[tco1, latency, labor_cost1, tco2, latency2, labor_cost2])\
        .then(results.create_table, inputs=[tco1, tco2, labor_cost1, labor_cost2, dropdown, dropdown2, latency, latency2], outputs=table)\
        .then(results.compare_info, inputs=[tco1, tco2, dropdown, dropdown2], outputs=[image, ratio])\
        .then(results.update_plot, inputs=[tco1, tco2, dropdown, dropdown2, labor_cost1, labor_cost2], outputs=plot)

demo.launch(debug=True)