from tkinter import filedialog, Tk
import os
import gradio as gr
from easygui import msgbox
import shutil

folder_symbol = '\U0001f4c2'  # 📂
refresh_symbol = '\U0001f504'  # 🔄
save_style_symbol = '\U0001f4be'  # 💾
document_symbol = '\U0001F4C4'   # 📄


def update_optimizer(my_data):
    if my_data.get('use_8bit_adam', False):
        my_data['optimizer'] = 'AdamW8bit'
        my_data['use_8bit_adam'] = False
    return my_data


def get_dir_and_file(file_path):
    dir_path, file_name = os.path.split(file_path)
    return (dir_path, file_name)


def has_ext_files(directory, extension):
    # Iterate through all the files in the directory
    for file in os.listdir(directory):
        # If the file name ends with extension, return True
        if file.endswith(extension):
            return True
    # If no extension files were found, return False
    return False


def get_file_path(
    file_path='', defaultextension='.json', extension_name='Config files'
):
    current_file_path = file_path
    # print(f'current file path: {current_file_path}')

    initial_dir, initial_file = get_dir_and_file(file_path)

    root = Tk()
    root.wm_attributes('-topmost', 1)
    root.withdraw()
    file_path = filedialog.askopenfilename(
        filetypes=(
            (f'{extension_name}', f'{defaultextension}'),
            ('All files', '*'),
        ),
        defaultextension=defaultextension,
        initialfile=initial_file,
        initialdir=initial_dir,
    )
    root.destroy()

    if file_path == '':
        file_path = current_file_path

    return file_path


def get_any_file_path(file_path=''):
    current_file_path = file_path
    # print(f'current file path: {current_file_path}')

    initial_dir, initial_file = get_dir_and_file(file_path)

    root = Tk()
    root.wm_attributes('-topmost', 1)
    root.withdraw()
    file_path = filedialog.askopenfilename(
        initialdir=initial_dir,
        initialfile=initial_file,
    )
    root.destroy()

    if file_path == '':
        file_path = current_file_path

    return file_path


def remove_doublequote(file_path):
    if file_path != None:
        file_path = file_path.replace('"', '')

    return file_path


def set_legacy_8bitadam(optimizer, use_8bit_adam):
    if optimizer == 'AdamW8bit':
        # use_8bit_adam = True
        return gr.Dropdown.update(value=optimizer), gr.Checkbox.update(
            value=True, interactive=False, visible=True
        )
    else:
        # use_8bit_adam = False
        return gr.Dropdown.update(value=optimizer), gr.Checkbox.update(
            value=False, interactive=False, visible=True
        )


def get_folder_path(folder_path=''):
    current_folder_path = folder_path

    initial_dir, initial_file = get_dir_and_file(folder_path)

    root = Tk()
    root.wm_attributes('-topmost', 1)
    root.withdraw()
    folder_path = filedialog.askdirectory(initialdir=initial_dir)
    root.destroy()

    if folder_path == '':
        folder_path = current_folder_path

    return folder_path


def get_saveasfile_path(
    file_path='', defaultextension='.json', extension_name='Config files'
):
    current_file_path = file_path
    # print(f'current file path: {current_file_path}')

    initial_dir, initial_file = get_dir_and_file(file_path)

    root = Tk()
    root.wm_attributes('-topmost', 1)
    root.withdraw()
    save_file_path = filedialog.asksaveasfile(
        filetypes=(
            (f'{extension_name}', f'{defaultextension}'),
            ('All files', '*'),
        ),
        defaultextension=defaultextension,
        initialdir=initial_dir,
        initialfile=initial_file,
    )
    root.destroy()

    # print(save_file_path)

    if save_file_path == None:
        file_path = current_file_path
    else:
        print(save_file_path.name)
        file_path = save_file_path.name

    # print(file_path)

    return file_path


def get_saveasfilename_path(
    file_path='', extensions='*', extension_name='Config files'
):
    current_file_path = file_path
    # print(f'current file path: {current_file_path}')

    initial_dir, initial_file = get_dir_and_file(file_path)

    root = Tk()
    root.wm_attributes('-topmost', 1)
    root.withdraw()
    save_file_path = filedialog.asksaveasfilename(
        filetypes=((f'{extension_name}', f'{extensions}'), ('All files', '*')),
        defaultextension=extensions,
        initialdir=initial_dir,
        initialfile=initial_file,
    )
    root.destroy()

    if save_file_path == '':
        file_path = current_file_path
    else:
        # print(save_file_path)
        file_path = save_file_path

    return file_path


def add_pre_postfix(
    folder='', prefix='', postfix='', caption_file_ext='.caption'
):
    if not has_ext_files(folder, caption_file_ext):
        msgbox(
            f'No files with extension {caption_file_ext} were found in {folder}...'
        )
        return

    if prefix == '' and postfix == '':
        return

    files = [f for f in os.listdir(folder) if f.endswith(caption_file_ext)]
    if not prefix == '':
        prefix = f'{prefix} '
    if not postfix == '':
        postfix = f' {postfix}'

    for file in files:
        with open(os.path.join(folder, file), 'r+') as f:
            content = f.read()
            content = content.rstrip()
            f.seek(0, 0)
            f.write(f'{prefix}{content}{postfix}')
    f.close()


def find_replace(folder='', caption_file_ext='.caption', find='', replace=''):
    print('Running caption find/replace')
    if not has_ext_files(folder, caption_file_ext):
        msgbox(
            f'No files with extension {caption_file_ext} were found in {folder}...'
        )
        return

    if find == '':
        return

    files = [f for f in os.listdir(folder) if f.endswith(caption_file_ext)]
    for file in files:
        with open(os.path.join(folder, file), 'r', errors='ignore') as f:
            content = f.read()
            f.close
        content = content.replace(find, replace)
        with open(os.path.join(folder, file), 'w') as f:
            f.write(content)
            f.close()


def color_aug_changed(color_aug):
    if color_aug:
        msgbox(
            'Disabling "Cache latent" because "Color augmentation" has been selected...'
        )
        return gr.Checkbox.update(value=False, interactive=False)
    else:
        return gr.Checkbox.update(value=True, interactive=True)


def save_inference_file(output_dir, v2, v_parameterization, output_name):
    # List all files in the directory
    files = os.listdir(output_dir)

    # Iterate over the list of files
    for file in files:
        # Check if the file starts with the value of output_name
        if file.startswith(output_name):
            # Check if it is a file or a directory
            if os.path.isfile(os.path.join(output_dir, file)):
                # Split the file name and extension
                file_name, ext = os.path.splitext(file)

                # Copy the v2-inference-v.yaml file to the current file, with a .yaml extension
                if v2 and v_parameterization:
                    print(
                        f'Saving v2-inference-v.yaml as {output_dir}/{file_name}.yaml'
                    )
                    shutil.copy(
                        f'./v2_inference/v2-inference-v.yaml',
                        f'{output_dir}/{file_name}.yaml',
                    )
                elif v2:
                    print(
                        f'Saving v2-inference.yaml as {output_dir}/{file_name}.yaml'
                    )
                    shutil.copy(
                        f'./v2_inference/v2-inference.yaml',
                        f'{output_dir}/{file_name}.yaml',
                    )


def set_pretrained_model_name_or_path_input(model_list, pretrained_model_name_or_path, v2, v_parameterization):
    # define a list of substrings to search for
    substrings_v2 = [
        'stabilityai/stable-diffusion-2-1-base',
        'stabilityai/stable-diffusion-2-base',
    ]

    # check if $v2 and $v_parameterization are empty and if $pretrained_model_name_or_path contains any of the substrings in the v2 list
    if str(model_list) in substrings_v2:
        print('SD v2 model detected. Setting --v2 parameter')
        v2 = True
        v_parameterization = False

        return model_list, v2, v_parameterization

    # define a list of substrings to search for v-objective
    substrings_v_parameterization = [
        'stabilityai/stable-diffusion-2-1',
        'stabilityai/stable-diffusion-2',
    ]

    # check if $v2 and $v_parameterization are empty and if $pretrained_model_name_or_path contains any of the substrings in the v_parameterization list
    if str(model_list) in substrings_v_parameterization:
        print(
            'SD v2 v_parameterization detected. Setting --v2 parameter and --v_parameterization'
        )
        v2 = True
        v_parameterization = True

        return model_list, v2, v_parameterization

    # define a list of substrings to v1.x
    substrings_v1_model = [
        'CompVis/stable-diffusion-v1-4',
        'runwayml/stable-diffusion-v1-5',
    ]

    if str(model_list) in substrings_v1_model:
        v2 = False
        v_parameterization = False

        return model_list, v2, v_parameterization

    if model_list == 'custom':
        if str(pretrained_model_name_or_path) in substrings_v1_model or str(pretrained_model_name_or_path) in substrings_v2 or str(pretrained_model_name_or_path) in substrings_v_parameterization:
            pretrained_model_name_or_path = ''
            v2 = False
            v_parameterization = False
        return pretrained_model_name_or_path, v2, v_parameterization

    ###
    ### Gradio common GUI section
    ###


def gradio_config():
    with gr.Accordion('Configuration file', open=False):
        with gr.Row():
            button_open_config = gr.Button('Open 📂', elem_id='open_folder')
            button_save_config = gr.Button('Save 💾', elem_id='open_folder')
            button_save_as_config = gr.Button(
                'Save as... 💾', elem_id='open_folder'
            )
            config_file_name = gr.Textbox(
                label='',
                placeholder="type the configuration file path or use the 'Open' button above to select it...",
                interactive=True,
            )
    return (
        button_open_config,
        button_save_config,
        button_save_as_config,
        config_file_name,
    )


def gradio_source_model():
    with gr.Tab('Source model'):
        # Define the input elements
        with gr.Row():
            pretrained_model_name_or_path = gr.Textbox(
                label='Pretrained model name or path',
                placeholder='enter the path to custom model or name of pretrained model',
                value='runwayml/stable-diffusion-v1-5'
            )
            pretrained_model_name_or_path_file = gr.Button(
                document_symbol, elem_id='open_folder_small'
            )
            pretrained_model_name_or_path_file.click(
                get_any_file_path,
                inputs=pretrained_model_name_or_path,
                outputs=pretrained_model_name_or_path,
            )
            pretrained_model_name_or_path_folder = gr.Button(
                folder_symbol, elem_id='open_folder_small'
            )
            pretrained_model_name_or_path_folder.click(
                get_folder_path,
                inputs=pretrained_model_name_or_path,
                outputs=pretrained_model_name_or_path,
            )
            model_list = gr.Dropdown(
                label='Model Quick Pick',
                choices=[
                    'custom',
                    'stabilityai/stable-diffusion-2-1-base',
                    'stabilityai/stable-diffusion-2-base',
                    'stabilityai/stable-diffusion-2-1',
                    'stabilityai/stable-diffusion-2',
                    'runwayml/stable-diffusion-v1-5',
                    'CompVis/stable-diffusion-v1-4',
                ],
                value='runwayml/stable-diffusion-v1-5'
            )
            save_model_as = gr.Dropdown(
                label='Save trained model as',
                choices=[
                    'same as source model',
                    'ckpt',
                    'diffusers',
                    'diffusers_safetensors',
                    'safetensors',
                ],
                value='safetensors',
            )

        with gr.Row():
            v2 = gr.Checkbox(label='v2', value=False)
            v_parameterization = gr.Checkbox(
                label='v_parameterization', value=False
            )
        model_list.change(
            set_pretrained_model_name_or_path_input,
            inputs=[model_list, pretrained_model_name_or_path, v2, v_parameterization],
            outputs=[
                pretrained_model_name_or_path,
                v2,
                v_parameterization,
            ],
        )
    return (
        pretrained_model_name_or_path,
        v2,
        v_parameterization,
        save_model_as,
        model_list,
    )


def gradio_training(
    learning_rate_value='1e-6',
    lr_scheduler_value='constant',
    lr_warmup_value='0',
):
    with gr.Row():
        train_batch_size = gr.Slider(
            minimum=1,
            maximum=64,
            label='Train batch size',
            value=1,
            step=1,
        )
        epoch = gr.Textbox(label='Epoch', value=1)
        save_every_n_epochs = gr.Textbox(label='Save every N epochs', value=1)
        caption_extension = gr.Textbox(
            label='Caption Extension',
            placeholder='(Optional) Extension for caption files. default: .caption',
        )
    with gr.Row():
        mixed_precision = gr.Dropdown(
            label='Mixed precision',
            choices=[
                'no',
                'fp16',
                'bf16',
            ],
            value='fp16',
        )
        save_precision = gr.Dropdown(
            label='Save precision',
            choices=[
                'float',
                'fp16',
                'bf16',
            ],
            value='fp16',
        )
        num_cpu_threads_per_process = gr.Slider(
            minimum=1,
            maximum=os.cpu_count(),
            step=1,
            label='Number of CPU threads per core',
            value=2,
        )
        seed = gr.Textbox(label='Seed', placeholder='(Optional) eg:1234')
        cache_latents = gr.Checkbox(label='Cache latent', value=True)
    with gr.Row():
        learning_rate = gr.Textbox(
            label='Learning rate', value=learning_rate_value
        )
        lr_scheduler = gr.Dropdown(
            label='LR Scheduler',
            choices=[
                'adafactor',
                'constant',
                'constant_with_warmup',
                'cosine',
                'cosine_with_restarts',
                'linear',
                'polynomial',
            ],
            value=lr_scheduler_value,
        )
        lr_warmup = gr.Textbox(
            label='LR warmup (% of steps)', value=lr_warmup_value
        )
        optimizer = gr.Dropdown(
            label='Optimizer',
            choices=[
                'AdamW',
                'AdamW8bit',
                'Adafactor',
                'DAdaptation',
                'Lion',
                'SGDNesterov',
                'SGDNesterov8bit',
            ],
            value='AdamW8bit',
            interactive=True,
        )
    with gr.Row():
        optimizer_args = gr.Textbox(
            label='Optimizer extra arguments',
            placeholder='(Optional) eg: relative_step=True scale_parameter=True warmup_init=True',
        )
    return (
        learning_rate,
        lr_scheduler,
        lr_warmup,
        train_batch_size,
        epoch,
        save_every_n_epochs,
        mixed_precision,
        save_precision,
        num_cpu_threads_per_process,
        seed,
        caption_extension,
        cache_latents,
        optimizer,
        optimizer_args,
    )


def run_cmd_training(**kwargs):
    options = [
        f' --learning_rate="{kwargs.get("learning_rate", "")}"'
        if kwargs.get('learning_rate')
        else '',
        f' --lr_scheduler="{kwargs.get("lr_scheduler", "")}"'
        if kwargs.get('lr_scheduler')
        else '',
        f' --lr_warmup_steps="{kwargs.get("lr_warmup_steps", "")}"'
        if kwargs.get('lr_warmup_steps')
        else '',
        f' --train_batch_size="{kwargs.get("train_batch_size", "")}"'
        if kwargs.get('train_batch_size')
        else '',
        f' --max_train_steps="{kwargs.get("max_train_steps", "")}"'
        if kwargs.get('max_train_steps')
        else '',
        f' --save_every_n_epochs="{kwargs.get("save_every_n_epochs", "")}"'
        if kwargs.get('save_every_n_epochs')
        else '',
        f' --mixed_precision="{kwargs.get("mixed_precision", "")}"'
        if kwargs.get('mixed_precision')
        else '',
        f' --save_precision="{kwargs.get("save_precision", "")}"'
        if kwargs.get('save_precision')
        else '',
        f' --seed="{kwargs.get("seed", "")}"' if kwargs.get('seed') else '',
        f' --caption_extension="{kwargs.get("caption_extension", "")}"'
        if kwargs.get('caption_extension')
        else '',
        ' --cache_latents' if kwargs.get('cache_latents') else '',
        # ' --use_lion_optimizer' if kwargs.get('optimizer') == 'Lion' else '',
        f' --optimizer_type="{kwargs.get("optimizer", "AdamW")}"',
        f' --optimizer_args {kwargs.get("optimizer_args", "")}'
        if not kwargs.get('optimizer_args') == ''
        else '',
    ]
    run_cmd = ''.join(options)
    return run_cmd


# # This function takes a dictionary of keyword arguments and returns a string that can be used to run a command-line training script
# def run_cmd_training(**kwargs):
#     arg_map = {
#         'learning_rate': ' --learning_rate="{}"',
#         'lr_scheduler': ' --lr_scheduler="{}"',
#         'lr_warmup_steps': ' --lr_warmup_steps="{}"',
#         'train_batch_size': ' --train_batch_size="{}"',
#         'max_train_steps': ' --max_train_steps="{}"',
#         'save_every_n_epochs': ' --save_every_n_epochs="{}"',
#         'mixed_precision': ' --mixed_precision="{}"',
#         'save_precision': ' --save_precision="{}"',
#         'seed': ' --seed="{}"',
#         'caption_extension': ' --caption_extension="{}"',
#         'cache_latents': ' --cache_latents',
#         'optimizer': ' --use_lion_optimizer' if kwargs.get('optimizer') == 'Lion' else '',
#     }

#     options = [arg_map[key].format(value) for key, value in kwargs.items() if key in arg_map and value]

#     cmd = ''.join(options)

#     return cmd


def gradio_advanced_training():
    with gr.Row():
        keep_tokens = gr.Slider(
            label='Keep n tokens', value='0', minimum=0, maximum=32, step=1
        )
        clip_skip = gr.Slider(
            label='Clip skip', value='1', minimum=1, maximum=12, step=1
        )
        max_token_length = gr.Dropdown(
            label='Max Token Length',
            choices=[
                '75',
                '150',
                '225',
            ],
            value='75',
        )
        full_fp16 = gr.Checkbox(
            label='Full fp16 training (experimental)', value=False
        )
    with gr.Row():
        gradient_checkpointing = gr.Checkbox(
            label='Gradient checkpointing', value=False
        )
        shuffle_caption = gr.Checkbox(label='Shuffle caption', value=False)
        persistent_data_loader_workers = gr.Checkbox(
            label='Persistent data loader', value=False
        )
        mem_eff_attn = gr.Checkbox(
            label='Memory efficient attention', value=False
        )
    with gr.Row():
        # This use_8bit_adam element should be removed in a future release as it is no longer used
        use_8bit_adam = gr.Checkbox(
            label='Use 8bit adam', value=False, visible=False
        )
        xformers = gr.Checkbox(label='Use xformers', value=True)
        color_aug = gr.Checkbox(label='Color augmentation', value=False)
        flip_aug = gr.Checkbox(label='Flip augmentation', value=False)
    with gr.Row():
        bucket_no_upscale = gr.Checkbox(
            label="Don't upscale bucket resolution", value=True
        )
        bucket_reso_steps = gr.Number(
            label='Bucket resolution steps', value=64
        )
        random_crop = gr.Checkbox(
            label='Random crop instead of center crop', value=False
        )
        noise_offset = gr.Textbox(
            label='Noise offset (0 - 1)', placeholder='(Oprional) eg: 0.1'
        )

    with gr.Row():
        caption_dropout_every_n_epochs = gr.Number(
            label='Dropout caption every n epochs', value=0
        )
        caption_dropout_rate = gr.Slider(
            label='Rate of caption dropout', value=0, minimum=0, maximum=1
        )
    with gr.Row():
        save_state = gr.Checkbox(label='Save training state', value=False)
        resume = gr.Textbox(
            label='Resume from saved training state',
            placeholder='path to "last-state" state folder to resume from',
        )
        resume_button = gr.Button('📂', elem_id='open_folder_small')
        resume_button.click(get_folder_path, outputs=resume)
        max_train_epochs = gr.Textbox(
            label='Max train epoch',
            placeholder='(Optional) Override number of epoch',
        )
        max_data_loader_n_workers = gr.Textbox(
            label='Max num workers for DataLoader',
            placeholder='(Optional) Override number of epoch. Default: 8',
        )
    return (
        use_8bit_adam,
        xformers,
        full_fp16,
        gradient_checkpointing,
        shuffle_caption,
        color_aug,
        flip_aug,
        clip_skip,
        mem_eff_attn,
        save_state,
        resume,
        max_token_length,
        max_train_epochs,
        max_data_loader_n_workers,
        keep_tokens,
        persistent_data_loader_workers,
        bucket_no_upscale,
        random_crop,
        bucket_reso_steps,
        caption_dropout_every_n_epochs,
        caption_dropout_rate,
        noise_offset,
    )


def run_cmd_advanced_training(**kwargs):
    options = [
        f' --max_train_epochs="{kwargs.get("max_train_epochs", "")}"'
        if kwargs.get('max_train_epochs')
        else '',
        f' --max_data_loader_n_workers="{kwargs.get("max_data_loader_n_workers", "")}"'
        if kwargs.get('max_data_loader_n_workers')
        else '',
        f' --max_token_length={kwargs.get("max_token_length", "")}'
        if int(kwargs.get('max_token_length', 75)) > 75
        else '',
        f' --clip_skip={kwargs.get("clip_skip", "")}'
        if int(kwargs.get('clip_skip', 1)) > 1
        else '',
        f' --resume="{kwargs.get("resume", "")}"'
        if kwargs.get('resume')
        else '',
        f' --keep_tokens="{kwargs.get("keep_tokens", "")}"'
        if int(kwargs.get('keep_tokens', 0)) > 0
        else '',
        f' --caption_dropout_every_n_epochs="{kwargs.get("caption_dropout_every_n_epochs", "")}"'
        if int(kwargs.get('caption_dropout_every_n_epochs', 0)) > 0
        else '',
        f' --caption_dropout_rate="{kwargs.get("caption_dropout_rate", "")}"'
        if float(kwargs.get('caption_dropout_rate', 0)) > 0
        else '',
        f' --bucket_reso_steps={int(kwargs.get("bucket_reso_steps", 1))}'
        if int(kwargs.get('bucket_reso_steps', 64)) >= 1
        else '',
        ' --save_state' if kwargs.get('save_state') else '',
        ' --mem_eff_attn' if kwargs.get('mem_eff_attn') else '',
        ' --color_aug' if kwargs.get('color_aug') else '',
        ' --flip_aug' if kwargs.get('flip_aug') else '',
        ' --shuffle_caption' if kwargs.get('shuffle_caption') else '',
        ' --gradient_checkpointing'
        if kwargs.get('gradient_checkpointing')
        else '',
        ' --full_fp16' if kwargs.get('full_fp16') else '',
        ' --xformers' if kwargs.get('xformers') else '',
        ' --use_8bit_adam' if kwargs.get('use_8bit_adam') else '',
        ' --persistent_data_loader_workers'
        if kwargs.get('persistent_data_loader_workers')
        else '',
        ' --bucket_no_upscale' if kwargs.get('bucket_no_upscale') else '',
        ' --random_crop' if kwargs.get('random_crop') else '',
        f' --noise_offset={float(kwargs.get("noise_offset", 0))}'
        if not kwargs.get('noise_offset', '') == ''
        else '',
    ]
    run_cmd = ''.join(options)
    return run_cmd


# def run_cmd_advanced_training(**kwargs):
#     arg_map = {
#         'max_train_epochs': ' --max_train_epochs="{}"',
#         'max_data_loader_n_workers': ' --max_data_loader_n_workers="{}"',
#         'max_token_length': ' --max_token_length={}' if int(kwargs.get('max_token_length', 75)) > 75 else '',
#         'clip_skip': ' --clip_skip={}' if int(kwargs.get('clip_skip', 1)) > 1 else '',
#         'resume': ' --resume="{}"',
#         'keep_tokens': ' --keep_tokens="{}"' if int(kwargs.get('keep_tokens', 0)) > 0 else '',
#         'caption_dropout_every_n_epochs': ' --caption_dropout_every_n_epochs="{}"' if int(kwargs.get('caption_dropout_every_n_epochs', 0)) > 0 else '',
#         'caption_dropout_rate': ' --caption_dropout_rate="{}"' if float(kwargs.get('caption_dropout_rate', 0)) > 0 else '',
#         'bucket_reso_steps': ' --bucket_reso_steps={:d}' if int(kwargs.get('bucket_reso_steps', 64)) >= 1 else '',
#         'save_state': ' --save_state',
#         'mem_eff_attn': ' --mem_eff_attn',
#         'color_aug': ' --color_aug',
#         'flip_aug': ' --flip_aug',
#         'shuffle_caption': ' --shuffle_caption',
#         'gradient_checkpointing': ' --gradient_checkpointing',
#         'full_fp16': ' --full_fp16',
#         'xformers': ' --xformers',
#         'use_8bit_adam': ' --use_8bit_adam',
#         'persistent_data_loader_workers': ' --persistent_data_loader_workers',
#         'bucket_no_upscale': ' --bucket_no_upscale',
#         'random_crop': ' --random_crop',
#     }

#     options = [arg_map[key].format(value) for key, value in kwargs.items() if key in arg_map and value]

#     cmd = ''.join(options)

#     return cmd