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

templates = get_templates()

def change(inp):
    """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.
    """
    if inp == "Basic":
        return (templates["initial"], highlight(inp), "## Accelerate Code (Base Integration)")
    elif inp == "Calculating Metrics":
        return (templates["initial_with_metrics"], highlight(inp), f"## Accelerate Code ({inp})")
    else:
        return (templates["accelerate"], highlight(inp), f"## Accelerate Code ({inp})")

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", "Gradient Accumulation", ], 
        label="Select a feature"
    )
    with gr.Row():
        with gr.Column():
            gr.Markdown("## Initial Code")
            code = gr.Markdown(templates["initial"])
        with gr.Column():
            feature = gr.Markdown("## Accelerate Code")
            out = gr.Markdown()
    inp.change(fn=change, inputs=inp, outputs=[code, out, feature])  
demo.launch()