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

templates = get_templates()


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

    Args:
        inp (`str`):
            The input button from the interface.
        textbox (`str`):
            The textbox specifying the tab name from the interface.

    Returns:
        `tuple`: A tuple of the highlighted code diff, and the title for the section.
    """
    if textbox == "base":
        code, explanation, docs = get_text(inp, textbox)
        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)
    elif textbox == "large_scale_training":
        config, code, explanation, docs = get_text(inp, textbox)
        return (highlight(config), highlight(code), f"## Accelerate Code ({inp})", explanation, docs)


default = change("Basic", "base")


def base_features(textbox):
    # textbox.value = "base"
    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("## Explanation")
            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, textbox], outputs=[out, feature, explanation, docs])


def large_scale_training(textbox):
    # textbox.value = "large_scale_training"
    inp = gr.Radio(
        ["Multi GPU", "Multi Node Multi GPU", "AWS SageMaker", "DeepSpeed", "PyTorch FSDP", "Megatron-LM"],
        label="Select a feature you would like to integrate",
        value="Basic",
    )
    with gr.Row():
        with gr.Column():
            feature = gr.Markdown("## Accelerate Config")
            config = gr.Markdown("")
    with gr.Row():
        with gr.Column():
            feature = gr.Markdown("## Accelerate Code")
            out = gr.Markdown("")
    with gr.Row():
        with gr.Column():
            gr.Markdown("## Explanation")
            explanation = gr.Markdown("")
    with gr.Row():
        with gr.Column():
            gr.Markdown("## Documentation Links")
            docs = gr.Markdown("")
    inp.change(fn=change, inputs=[inp, textbox], outputs=[config, out, feature, explanation, docs])


# def big_model_inference():
#     inp = gr.Radio(
#         ["Accelerate's Big Model Inference",], # "DeepSpeed ZeRO Stage-3 Offload"
#         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, "big_model_inference"], outputs=[out, feature, explanation, docs])


# def notebook_launcher():
#     inp = gr.Radio(
#         ["Colab GPU", "Colab TPU", "Kaggle GPU", "Kaggle Multi GPU", "Kaggle TPU", "Multi GPU VMs"],
#         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, "notebook_launcher"], outputs=[out, feature, explanation, docs])


with gr.Blocks() as demo:

    with gr.Tabs():
        with gr.TabItem("Simplify your code and improve efficieny"):
            textbox = gr.Textbox(label="tab_name", visible=False, value="base")
            base_features(textbox)
        with gr.TabItem("Large Scale Training"):
            textbox = gr.Textbox(label="tab_name", visible=False, value="large_scale_training")
            large_scale_training(textbox)
        with gr.TabItem("Big Model Inference"):
            # big_model_inference()
            pass
        with gr.TabItem("Notebook Launcher Intergation"):
            # notebook_launcher()
            pass


demo.launch()