import argparse
import itertools
import math
import os
from pathlib import Path
from typing import Optional
import subprocess
import sys

import torch


from spanish_medica_llm import run_training, run_training_process

import gradio as gr

#def greet(name):
#    return "Hello " + name + "!!"

#iface = gr.Interface(fn=greet, inputs="text", outputs="text")
#iface.launch()

def generate_model(name):
    return f"Welcome to Gradio HF_ACCES_TOKEN, {os.environ.get('HG_FACE_TOKEN')}!"
    
def generate(prompt):
    #from diffusers import StableDiffusionPipeline
    
    #pipe = StableDiffusionPipeline.from_pretrained("./output_model", torch_dtype=torch.float16)
    pipe = pipe.to("cuda")
    image = pipe(prompt).images[0]  
    return(image)
    
def evaluate_model():
    #from diffusers import StableDiffusionPipeline
    
    #pipe = StableDiffusionPipeline.from_pretrained("./output_model", torch_dtype=torch.float16)
    #pipe = pipe.to("cuda")
    #image = pipe(prompt).images[0]  
    return(f"Evaluate Model {os.environ.get('HF_LLM_MODEL_ID')} from dataset {os.environ.get('HF_LLM_DATASET_ID')}")



def train_model(*inputs):
    if "IS_SHARED_UI" in os.environ:
        raise gr.Error("This Space only works in duplicated instances")
  
    # args_general = argparse.Namespace(
    #             image_captions_filename = True,
    #             train_text_encoder = True,
    #             #stop_text_encoder_training = stptxt,
    #             save_n_steps = 0,
    #             #pretrained_model_name_or_path = model_to_load,
    #             instance_data_dir="instance_images",
    #             #class_data_dir=class_data_dir,
    #             output_dir="output_model",
    #             instance_prompt="",
    #             seed=42,
    #             resolution=512,
    #             mixed_precision="fp16",
    #             train_batch_size=1,
    #             gradient_accumulation_steps=1,
    #             use_8bit_adam=True,
    #             learning_rate=2e-6,
    #             lr_scheduler="polynomial",
    #             lr_warmup_steps = 0,
    #             #max_train_steps=Training_Steps,     
    # )
    # run_training(args_general)
    # torch.cuda.empty_cache()
    # #convert("output_model", "model.ckpt")
    # #shutil.rmtree('instance_images')
    # #shutil.make_archive("diffusers_model", 'zip', "output_model")
    # #with zipfile.ZipFile('diffusers_model.zip', 'w', zipfile.ZIP_DEFLATED) as zipf:
    # #    zipdir('output_model/', zipf)
    # torch.cuda.empty_cache()
    # return [gr.update(visible=True, value=["diffusers_model.zip"]), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)]
    run_training_process()
    
    return f"Train Model Sucessful!!!"

def stop_model(*input):
    return f"Model with Gradio!"


with gr.Blocks() as demo:
    gr.Markdown("Start typing below and then click **Run** to see the output.")
    with gr.Row():
        inp = gr.Textbox(placeholder="What is your name?")
        out = gr.Textbox()
    btn_response = gr.Button("Generate Response")
    btn_response.click(fn=generate_model, inputs=inp, outputs=out)
    btn_train = gr.Button("Train Model")
    btn_train.click(fn=train_model, inputs=[], outputs=out)
    btn_evaluate = gr.Button("Evaluate Model")
    btn_evaluate.click(fn=evaluate_model, inputs=[], outputs=out)
    btn_stop = gr.Button("Stop Model")
    btn_stop.click(fn=stop_model, inputs=[], outputs=out)

demo.launch()