#!/usr/bin/env python

from __future__ import annotations

import os

import gradio as gr

from constants import UploadTarget
from inference import InferencePipeline
from trainer import Trainer


def create_training_demo(
    trainer: Trainer, pipe: InferencePipeline | None = None, disable_run_button: bool = False
) -> gr.Blocks:
    def read_log() -> str:
        with open(trainer.log_file) as f:
            lines = f.readlines()
        return "".join(lines[-10:])

    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                with gr.Box():
                    gr.Markdown("Training Data")
                    training_video = gr.File(label="Training video")
                    training_prompt = gr.Textbox(label="Training prompt", max_lines=1, placeholder="A man is surfing")
                    gr.Markdown(
                        """
                        - Upload a video and write a `Training Prompt` that describes the video.
                        """
                    )

            with gr.Column():
                with gr.Box():
                    gr.Markdown("Training Parameters")
                    with gr.Row():
                        base_model = gr.Text(label="Base Model", value="CompVis/stable-diffusion-v1-4", max_lines=1)
                        resolution = gr.Dropdown(
                            choices=["512", "768"], value="512", label="Resolution", visible=False
                        )

                    hf_token = gr.Text(
                        label="Hugging Face Write Token", type="password", visible=os.getenv("HF_TOKEN") is None
                    )
                    with gr.Accordion(label="Advanced options", open=False):
                        num_training_steps = gr.Number(label="Number of Training Steps", value=300, precision=0)
                        learning_rate = gr.Number(label="Learning Rate", value=0.000035)
                        gradient_accumulation = gr.Number(
                            label="Number of Gradient Accumulation", value=1, precision=0
                        )
                        seed = gr.Slider(label="Seed", minimum=0, maximum=100000, step=1, randomize=True, value=0)
                        fp16 = gr.Checkbox(label="FP16", value=True)
                        use_8bit_adam = gr.Checkbox(label="Use 8bit Adam", value=False)
                        checkpointing_steps = gr.Number(label="Checkpointing Steps", value=1000, precision=0)
                        validation_epochs = gr.Number(label="Validation Epochs", value=100, precision=0)
                    gr.Markdown(
                        """
                        - The base model must be a Stable Diffusion model compatible with [diffusers](https://github.com/huggingface/diffusers) library.
                        - Expected time to train a model for 300 steps: ~20 minutes with T4
                        - You can check the training status by pressing the "Open logs" button if you are running this on your Space.
                        """
                    )

        with gr.Row():
            with gr.Column():
                gr.Markdown("Output Model")
                output_model_name = gr.Text(label="Name of your model", placeholder="The surfer man", max_lines=1)
                validation_prompt = gr.Text(
                    label="Validation Prompt", placeholder="prompt to test the model, e.g: a dog is surfing"
                )
            with gr.Column():
                gr.Markdown("Upload Settings")
                with gr.Row():
                    upload_to_hub = gr.Checkbox(label="Upload model to Hub", value=True)
                    use_private_repo = gr.Checkbox(label="Private", value=True)
                    delete_existing_repo = gr.Checkbox(label="Delete existing repo of the same name", value=False)
                    upload_to = gr.Radio(
                        label="Upload to",
                        choices=[_.value for _ in UploadTarget],
                        value=UploadTarget.MODEL_LIBRARY.value,
                    )

        pause_space_after_training = gr.Checkbox(
            label="Pause this Space after training",
            value=False,
            interactive=bool(os.getenv("SPACE_ID")),
            visible=False,
        )
        run_button = gr.Button("Start Training", interactive=not disable_run_button)

        with gr.Box():
            gr.Text(label="Log", value=read_log, lines=10, max_lines=10, every=1)

        if pipe is not None:
            run_button.click(fn=pipe.clear)
        run_button.click(
            fn=trainer.run,
            inputs=[
                training_video,
                training_prompt,
                output_model_name,
                delete_existing_repo,
                validation_prompt,
                base_model,
                resolution,
                num_training_steps,
                learning_rate,
                gradient_accumulation,
                seed,
                fp16,
                use_8bit_adam,
                checkpointing_steps,
                validation_epochs,
                upload_to_hub,
                use_private_repo,
                delete_existing_repo,
                upload_to,
                pause_space_after_training,
                hf_token,
            ],
        )
    return demo


if __name__ == "__main__":
    trainer = Trainer()
    demo = create_training_demo(trainer)
    demo.queue(api_open=False, max_size=1).launch()