import gradio as gr
from markup import highlight, get_text
from template import get_templates

templates = get_templates()

def change(inp, components=[]):
    """Based on an `inp`, render and highlight the appropriate code sample.

    Args:
        inp (`str`):
            The input button from the interface.

    Returns:
        `tuple`: A tuple of the initial code, the highlighted code, and the title for the section.
    """
    code, explanation, docs = get_text(inp)
    if inp == "Basic":
        return (templates["initial"], highlight(code), "## Accelerate Code (Base Integration)", explanation, docs)
    elif inp == "Calculating Metrics":
        return (templates["initial_with_metrics"], highlight(code), f"## Accelerate Code ({inp})", explanation, docs)
    else:
        return (templates["accelerate"], highlight(code), f"## Accelerate Code ({inp})", explanation, docs)

initial_md = gr.Markdown("## Initial Code")
initial_code = gr.Markdown(templates["initial"])

with gr.Blocks() as demo:
    gr.Markdown(f'''# Accelerate Training Code Template Generator
Here is a very basic Python training loop.
Select how you would like to introduce an Accelerate capability to add to it.''')
    inp = gr.Radio(
        ["Basic", "Calculating Metrics", "Checkpointing", "Experiment Tracking", "Gradient Accumulation"], 
        label="Select a feature"
    )
    with gr.Row():
        with gr.Column():
            initial_md.render()
            initial_code.render()
        with gr.Column():
            feature = gr.Markdown("## Accelerate Code")
            out = gr.Markdown()
    with gr.Row():
        with gr.Column():    
            gr.Markdown("## Documentation Links")
            docs = gr.Markdown()
        with gr.Column():    
            gr.Markdown("## Explanation")
            explanation = gr.Markdown()
    inp.change(
        fn=change, 
        inputs=inp, 
        outputs=[initial_code, out, feature, explanation, docs]
    )
demo.launch()