##################################### Imports ######################################
# Generic imports
import gradio as gr

# Module imports
from utilities.setup import get_json_cfg
from utilities.templates import prompt_template

########################### Global objects and functions ###########################

conf = get_json_cfg()

def textbox_visibility(radio):
    value = radio
    if value == "Hugging Face Hub Dataset":
        return gr.Dropdown(visible=bool(1))
    else:
        return gr.Dropdown(visible=bool(0))

def upload_visibility(radio):
    value = radio
    if value == "Upload Your Own":
        return gr.UploadButton(visible=bool(1)) #make it visible
    else:
        return gr.UploadButton(visible=bool(0))

def greet(model_name, inject_prompt, dataset, pefts):
    """The model call"""
    return f"Hello!! Using model: {model_name} with template: {inject_prompt}"

##################################### App UI #######################################

def main():
    with gr.Blocks() as demo:
    
        ##### Title Block #####
        gr.Markdown("# Instruction Tuning with Unsloth")
    
        ##### Model Inputs #####
    
        # Select Model
        modelnames = conf['model']['choices']
        model_name = gr.Dropdown(label="Supported Models", 
                                 choices=modelnames, 
                                 value=modelnames[0])
        # Prompt template
        inject_prompt = gr.Textbox(label="Prompt Template", 
                                     value=prompt_template())
        # Dataset choice
        dataset_choice = gr.Radio(label="Choose Dataset", 
                                  choices=["Hugging Face Hub Dataset", "Upload Your Own"], 
                                  value="Hugging Face Hub Dataset")
        dataset_predefined = gr.Textbox(label="Hugging Face Hub Dataset", 
                                        value='yahma/alpaca-cleaned', 
                                        visible=True)
        dataset_upload = gr.UploadButton(label="Upload Dataset (csv, jsonl, or txt)", 
                                         file_types=[".csv",".jsonl", ".txt"], 
                                         visible=False)
        dataset_choice.change(textbox_visibility, 
                              dataset_choice, 
                              dataset_predefined)
        dataset_choice.change(upload_visibility, 
                              dataset_choice, 
                              dataset_upload)
        # Hyperparameters (allow selection, but hide in accordion.)
        with gr.Accordion("Advanced Tuning", open=False):
            # config
            peftparams = conf['model']['peft']
            sftparams = conf['model']['sft']
            # accordion container content
            gr.Markdown("### PEFT Parameters")
            r = gr.Textbox(label="r", 
                           value=peftparams['r'])
            alpha = gr.Textbox(label="LoRA alpha", 
                               value=peftparams['alpha'])
            dropout = gr.Textbox(label="LoRA dropout", 
                                 value=peftparams['dropout'])
            bias = gr.Textbox(label="Bias", 
                              value=peftparams['bias'])
            seed = gr.Textbox(label="Random State", 
                              value=peftparams['seed'])
            rslora = gr.Textbox(label="Use R-S LoRA", 
                                value=peftparams['rslora'])
            gr.Markdown("### Supervised Fine-Tuning Parameters")
            gr.Markdown("List of items")
        
        ##### Execution #####
    
        # Setup button
        tune_btn = gr.Button("Start Fine Tuning")
        # Text output (for now)
        output = gr.Textbox(label="Output") 
        # Execute button
        tune_btn.click(fn=greet, 
                       inputs=[model_name, 
                               inject_prompt, 
                               dataset_predefined, 
                               # peft
                               r, alpha, dropout, bias, seed, rslora
                               # sft
                              ],
                       outputs=output)
        # Launch baby
        demo.launch()

##################################### Launch #######################################

if __name__ == "__main__":
    main()