import gradio as gr
import shutil
import zipfile
import tensorflow as tf
import pandas as pd
import pathlib
import PIL.Image
import os
import subprocess

def pad_image(image: PIL.Image.Image) -> PIL.Image.Image:
    w, h = image.size
    if w == h:
        return image
    elif w > h:
        new_image = PIL.Image.new(image.mode, (w, w), (0, 0, 0))
        new_image.paste(image, (0, (w - h) // 2))
        return new_image
    else:
        new_image = PIL.Image.new(image.mode, (h, h), (0, 0, 0))
        new_image.paste(image, ((h - w) // 2, 0))
        return new_image


class ModelTrainer:
    def __init__(self):
        self.training_pictures = []
        self.training_model = None

    def unzip_file(self, zip_file_path):
        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
            extracted_path = zip_file_path.replace('.zip', '')
            zip_ref.extractall(extracted_path)
            file_names = zip_ref.namelist()
            for file_name in file_names:
                if file_name.endswith(('.jpeg', '.jpg', '.png')):
                    self.training_pictures.append(f'{extracted_path}/{file_name}')

    def train(self, pretrained_model_name_or_path: str, instance_images: list | None):
        output_model_name = 'a-xyz-model'
        resolution = 512
        repo_dir = pathlib.Path(__file__).parent
        subdirs = ['train-instance', 'train-class', 'experiments']
        dir_paths = []

        for subdir in subdirs:
            dir_path = repo_dir / subdir / output_model_name
            dir_paths.append(dir_path)
            shutil.rmtree(dir_path, ignore_errors=True)
            os.makedirs(dir_path, exist_ok=True)

        instance_data_dir, class_data_dir, output_dir = dir_paths

        for i, temp_path in enumerate(instance_images):
            image = PIL.Image.open(temp_path.name)
            image = pad_image(image)
            image = image.resize((resolution, resolution))
            image = image.convert('RGB')
            out_path = instance_data_dir / f'{i:03d}.jpg'
            image.save(out_path, format='JPEG', quality=100)

        command = [
            'python', '-u',
            'train_dreambooth_cloneofsimo_lora.py',
            '--pretrained_model_name_or_path', pretrained_model_name_or_path,
            '--instance_data_dir', instance_data_dir,
            '--class_data_dir', class_data_dir,
            '--resolution', '768',
            '--output_dir', output_dir,
            '--instance_prompt', 'a photo of a pwsm dog',
            '--with_prior_preservation',
            '--class_prompt', 'a dog',
            '--prior_loss_weight', '1.0',
            '--num_class_images', '100',
            '--learning_rate', '0.0004',
            '--train_batch_size', '1',
            '--sample_batch_size', '1',
            '--max_train_steps', '400',
            '--gradient_accumulation_steps', '1',
            '--gradient_checkpointing',
            '--train_text_encoder',
            '--learning_rate_text', '5e-6',
            '--save_steps', '100',
            '--seed', '1337',
            '--lr_scheduler', 'constant',
            '--lr_warmup_steps', '0'
        ]

        result = subprocess.run(command)
        return result

    def generate_picture(self, row):
        num_of_training_steps, learning_rate, checkpoint_steps, abc = row
        return f'Picture generated for num_of_training_steps: {num_of_training_steps}, learning_rate: {learning_rate}, checkpoint_steps: {checkpoint_steps}'

    def generate_pictures(self, csv_input):
        csv = pd.read_csv(csv_input.name)
        result = []
        for index, row in csv.iterrows():
            result.append(self.generate_picture(row))
        return "\n".join(str(item) for item in result)

loader = ModelTrainer()

with gr.Blocks() as demo:
    with gr.Box():
        instance_images = gr.Files(label='Instance images')
        pretrained_model_name_or_path = gr.inputs.Textbox(lines=1, label='pretrained_model_name_or_path', default='stabilityai/stable-diffusion-2-1')
        output_message = gr.Markdown()
        train_button = gr.Button('Train')
        train_button.click(fn=loader.train, inputs=[pretrained_model_name_or_path, instance_images], outputs=[output_message])
    with gr.Box():
        csv_input = gr.inputs.File(label='CSV File')
        output_message2 = gr.Markdown()
        generate_button = gr.Button('Generate Pictures from CSV')
        generate_button.click(fn=loader.generate_pictures, inputs=[csv_input], outputs=[output_message2])

demo.launch()