import os

import gradio as gr
import numpy as np
import torch
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
from PIL import Image

from video_diffusion.inpaint_zoom.utils.zoom_in_utils import dummy, image_grid, shrink_and_paste_on_blank, write_video

os.environ["CUDA_VISIBLE_DEVICES"] = "0"


stable_paint_model_list = ["stabilityai/stable-diffusion-2-inpainting", "runwayml/stable-diffusion-inpainting"]

stable_paint_prompt_list = [
    "children running in the forest , sunny, bright, by studio ghibli painting, superior quality, masterpiece,  traditional Japanese colors, by Grzegorz Rutkowski, concept art",
    "A beautiful landscape of a mountain range with a lake in the foreground",
]

stable_paint_negative_prompt_list = [
    "lurry, bad art, blurred, text, watermark",
]


class StableDiffusionZoomIn:
    def __init__(self):
        self.pipe = None

    def load_model(self, model_id):
        if self.pipe is None:
            self.pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="fp16")
            self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
            self.pipe = self.pipe.to("cuda")
            self.pipe.safety_checker = dummy
            self.pipe.enable_attention_slicing()
            self.pipe.enable_xformers_memory_efficient_attention()
            self.g_cuda = torch.Generator(device="cuda")

        return self.pipe

    def generate_video(
        self,
        model_id,
        prompt,
        negative_prompt,
        guidance_scale,
        num_inference_steps,
    ):
        pipe = self.load_model(model_id)

        num_init_images = 2
        seed = 42
        height = 512
        width = height

        current_image = Image.new(mode="RGBA", size=(height, width))
        mask_image = np.array(current_image)[:, :, 3]
        mask_image = Image.fromarray(255 - mask_image).convert("RGB")
        current_image = current_image.convert("RGB")

        init_images = pipe(
            prompt=[prompt] * num_init_images,
            negative_prompt=[negative_prompt] * num_init_images,
            image=current_image,
            guidance_scale=guidance_scale,
            height=height,
            width=width,
            generator=self.g_cuda.manual_seed(seed),
            mask_image=mask_image,
            num_inference_steps=num_inference_steps,
        )[0]

        image_grid(init_images, rows=1, cols=num_init_images)

        init_image_selected = 1  # @param
        if num_init_images == 1:
            init_image_selected = 0
        else:
            init_image_selected = init_image_selected - 1

        num_outpainting_steps = 20  # @param
        mask_width = 128  # @param
        num_interpol_frames = 30  # @param

        current_image = init_images[init_image_selected]
        all_frames = []
        all_frames.append(current_image)

        for i in range(num_outpainting_steps):
            print("Generating image: " + str(i + 1) + " / " + str(num_outpainting_steps))

            prev_image_fix = current_image

            prev_image = shrink_and_paste_on_blank(current_image, mask_width)

            current_image = prev_image

            # create mask (black image with white mask_width width edges)
            mask_image = np.array(current_image)[:, :, 3]
            mask_image = Image.fromarray(255 - mask_image).convert("RGB")

            # inpainting step
            current_image = current_image.convert("RGB")
            images = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                image=current_image,
                guidance_scale=guidance_scale,
                height=height,
                width=width,
                # this can make the whole thing deterministic but the output less exciting
                # generator = g_cuda.manual_seed(seed),
                mask_image=mask_image,
                num_inference_steps=num_inference_steps,
            )[0]
            current_image = images[0]
            current_image.paste(prev_image, mask=prev_image)

            # interpolation steps bewteen 2 inpainted images (=sequential zoom and crop)
            for j in range(num_interpol_frames - 1):
                interpol_image = current_image
                interpol_width = round(
                    (1 - (1 - 2 * mask_width / height) ** (1 - (j + 1) / num_interpol_frames)) * height / 2
                )
                interpol_image = interpol_image.crop(
                    (interpol_width, interpol_width, width - interpol_width, height - interpol_width)
                )

                interpol_image = interpol_image.resize((height, width))

                # paste the higher resolution previous image in the middle to avoid drop in quality caused by zooming
                interpol_width2 = round((1 - (height - 2 * mask_width) / (height - 2 * interpol_width)) / 2 * height)
                prev_image_fix_crop = shrink_and_paste_on_blank(prev_image_fix, interpol_width2)
                interpol_image.paste(prev_image_fix_crop, mask=prev_image_fix_crop)

                all_frames.append(interpol_image)

            all_frames.append(current_image)

        video_file_name = "infinite_zoom_out"
        fps = 30
        save_path = video_file_name + ".mp4"
        write_video(save_path, all_frames, fps)
        return save_path

    def app():
        with gr.Blocks():
            with gr.Row():
                with gr.Column():
                    text2image_in_model_path = gr.Dropdown(
                        choices=stable_paint_model_list, value=stable_paint_model_list[0], label="Text-Image Model Id"
                    )

                    text2image_in_prompt = gr.Textbox(lines=2, value=stable_paint_prompt_list[0], label="Prompt")

                    text2image_in_negative_prompt = gr.Textbox(
                        lines=1, value=stable_paint_negative_prompt_list[0], label="Negative Prompt"
                    )

                    with gr.Row():
                        with gr.Column():
                            text2image_in_guidance_scale = gr.Slider(
                                minimum=0.1, maximum=15, step=0.1, value=7.5, label="Guidance Scale"
                            )

                            text2image_in_num_inference_step = gr.Slider(
                                minimum=1, maximum=100, step=1, value=50, label="Num Inference Step"
                            )

                    text2image_in_predict = gr.Button(value="Generator")

                with gr.Column():
                    output_image = gr.Video(label="Output")

            text2image_in_predict.click(
                fn=StableDiffusionZoomIn().generate_video,
                inputs=[
                    text2image_in_model_path,
                    text2image_in_prompt,
                    text2image_in_negative_prompt,
                    text2image_in_guidance_scale,
                    text2image_in_num_inference_step,
                ],
                outputs=output_image,
            )