import gradio as gr
from markup import highlight, get_text
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 highlighted code diff, and the title for the section.
    """
    code, explanation, docs = get_text(inp)
    if inp == "Basic":
        return (highlight(code), "## Accelerate Code (Base Integration)", explanation, docs)
    elif inp == "Calculating Metrics":
        return (highlight(code), f"## Accelerate Code ({inp})", explanation, docs)
    else:
        return (highlight(code), f"## Accelerate Code ({inp})", explanation, docs)

default = change("Basic")

with gr.Blocks() as demo:
    inp = gr.Radio(
        ["Basic", "Calculating Metrics", "Checkpointing", "Experiment Tracking", "Gradient Accumulation"], 
        label="Select a feature you would like to integrate",
        value="Basic"
    )
    with gr.Row():
        with gr.Column():
            feature = gr.Markdown("## Accelerate Code")
            out = gr.Markdown(default[0])
    with gr.Row():
        with gr.Column():    
            gr.Markdown(default[1])
            explanation = gr.Markdown(default[2])
    with gr.Row():
        with gr.Column():    
            gr.Markdown("## Documentation Links")
            docs = gr.Markdown(default[3])
    inp.change(
        fn=change, 
        inputs=inp,
        outputs=[out, feature, explanation, docs]
    )
demo.launch()