import gradio as gr
from task import tasks_config
from transformers import pipeline


def review_training_choices(choice):
    print(choice)
    if choice == "Use Pipeline":
        return gr.Row(visible=True)
    else:
        return gr.Row(visible=False)

def task_dropdown_choices():
    return [(task["name"], task_id)
                                 for task_id, task in tasks_config.items()]
def handle_task_change(task):
    visibility = task == "question-answering"
    models = tasks_config[task]["config"]["models"]
    model_choices = [(model, model) for model in models]
    return gr.update(visible=visibility), gr.Dropdown(
        choices=model_choices,
        label="Model",
        allow_custom_value=True,
        interactive=True
    ), gr.Dropdown(info=tasks_config[task]["info"])


def test_pipeline(task, model=None, prompt=None, context=None):
    # configure additional options for each model
    options = {"ner": {"grouped_entities": True}, "question-answering": {},
               "text-generation": {}, "fill-mask": {}, "summarization": {}}
    # configure pipeline
    test = pipeline(task, model=model, **
                    options[task]) if model else pipeline(task, **options[task])
    # call pipeline
    if task == "question-answering":
        if not context:
            return "Context is required"
        else:
            result = test(question=prompt, context=context)
    else:
        result = test(prompt)

    # generated ouput based on task and return
    output_mapping = {
        "text-generation": lambda x: x[0]["generated_text"],
        "fill-mask": lambda x: x[0]["sequence"],
        "summarization": lambda x: x[0]["summary_text"],
        "ner": lambda x: "\n".join(f"{k}={v}" for item in x for k, v in item.items() if k not in ["start", "end", "index"]).rstrip("\n"),
        "question-answering": lambda x: x
    }

    return gr.TextArea(output_mapping[task](result))