import os
import random
from typing import Optional

import gradio as gr
import numpy as np
import PIL.Image
import spaces
import torch
from diffusers import (AutoencoderKL, DDIMInverseScheduler, DDIMScheduler,
                       StableDiffusionXLPipeline)
from torchvision.transforms import ToTensor 

# pyright: reportPrivateImportUsage=false


DESCRIPTION = f"""
    # 🎨 Noisy-Style 🎨  
    This is an interactive demo of noisy DDIM inversion capabilities on top of Instant-Style styling method 
    
    This method proposed in [Controllability of diffusion models]() *by Eyal Benaroche, Clément Chadebec, Onur Tasar, and Benjamin Aubin* from Jasper Research in the context of Eyal's internship with Ecole Polytechnique.

    A style benchmark : [style-bench](https://gojasper.github.io/style-bench) was also provided to facilitate evaluation of diffusion models for styling purposes.
    """

OPEN_SOURCE_PROMO = f"""
    If you enjoy the space, please also promote *open-source* by giving a ⭐ to our repo [![GitHub Stars](https://img.shields.io/github/stars/gojasper/style-bench?style=social)](https://github.com/gojasper/style-bench)
    """

DISCLAIMER = f"""
    This demo is only for research purpose. Users are solely responsible for any content they create, and it is their obligation to ensure that it adheres to appropriate and ethical standards. """

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

if gr.NO_RELOAD:
    if torch.cuda.is_available():
        vae = AutoencoderKL.from_pretrained(
            "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
        )
        pipe = StableDiffusionXLPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0",
            vae=vae,
            torch_dtype=torch.float16,
            use_safetensors=True,
            variant="fp16",
        )
        pipe.load_ip_adapter(
            "h94/IP-Adapter",
            subfolder="sdxl_models",
            weight_name="ip-adapter_sdxl.safetensors"
        )
        pipe.to(device)

        forward_scheduler = DDIMScheduler.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler"
        )
        invert_scheduler = DDIMInverseScheduler(**forward_scheduler.config)

css = """
    h1 {
        text-align: center;
        display:block;
    }
    p {
        text-align: justify;
        display:block;
    }
"""


def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed


def img_to_latents(x: torch.Tensor, vae: AutoencoderKL):
    x = 2.0 * x - 1.0
    posterior = vae.encode(x).latent_dist
    latents = posterior.mean * 0.18215
    return latents


def invert_image(model, image: np.ndarray, n_steps: int, width:int, height:int):

    model.scheduler = invert_scheduler

    image = PIL.Image.fromarray(image).resize((width, height))
    image_tensor = ToTensor()(image).to(model.device, dtype=torch.float16)
    image_tensor = image_tensor.unsqueeze(0)
    latent = img_to_latents(image_tensor, model.vae)

    print(latent)
    model.set_ip_adapter_scale(0)

    inv_latents = model(
        prompt="",
        negative_prompt="",
        ip_adapter_image=image,
        guidance_scale=1.0,
        output_type="latent",
        return_dict=False,
        num_inference_steps=n_steps,
        latents=latent,
    )[0]

    return inv_latents


@spaces.GPU
def generate(
    prompt: str,
    negative_prompt: str = "",
    prompt_2: str = "",
    negative_prompt_2: str = "",
    use_negative_prompt: bool = False,
    use_prompt_2: bool = False,
    use_negative_prompt_2: bool = False,
    seed: int = 0,
    width: int = 1024,
    height: int = 1024,
    guidance_scale_base: float = 5.0,
    num_inference_steps_base: int = 25,
    style_image_value = None,
    noise_scale: float = 1.5,
) -> PIL.Image.Image:
    torch.manual_seed(seed)

    if not use_negative_prompt:
        negative_prompt = None  # type: ignore
    if not use_prompt_2:
        prompt_2 = None  # type: ignore
    if not use_negative_prompt_2:
        negative_prompt_2 = None  # type: ignore

    # Add scaled noise to the latent
    noise = torch.randn(1, 4, width // 8, height // 8).to(device, dtype=torch.float16)

    # Invert the image and get the latent
    if style_image_value is not None:
        latent = invert_image(pipe, style_image_value, 30, width, height)
        print("Image was inverted")
        print(latent)

        latent = latent + noise_scale * noise
        latent = latent / torch.sqrt(torch.tensor(1 + noise_scale**2).to(device, dtype=torch.float16))

    else:
        latent = noise

    print("Noise added")
    print(latent)

    scale = {
        "up": {"block_0": [0.0, 1.0, 0.0]},
    }
    pipe.set_ip_adapter_scale(scale)

    pipe.scheduler = forward_scheduler
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        ip_adapter_image=style_image_value,
        latents=latent,
        prompt_2=prompt_2,
        negative_prompt_2=negative_prompt_2,
        guidance_scale=guidance_scale_base,
        num_inference_steps=num_inference_steps_base,
        output_type="pil",
    ).images[0]

    return image


examples = [
    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
    "An astronaut riding a green horse",
]

with gr.Blocks(css=css) as demo:
    gr.Markdown(DESCRIPTION)

    gr.Markdown(OPEN_SOURCE_PROMO)

    with gr.Row():

        with gr.Blocks():

            with gr.Column():

                style_image = gr.Image()

                noise_scale = gr.Slider(
                    label="Noise Scale",
                    minimum=0,
                    maximum=5,
                    step=0.1,
                    value=1.5,
                )

        with gr.Blocks():

            with gr.Column():
                with gr.Row():
                    prompt = gr.Text(
                        label="Prompt",
                        show_label=False,
                        max_lines=1,
                        placeholder="Enter your prompt",
                        container=False,
                    )
                    run_button = gr.Button("Run", scale=0)

                result = gr.Image(label="Result", show_label=False)

    with gr.Accordion("Advanced options", open=False):
        with gr.Row():
            use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
            use_prompt_2 = gr.Checkbox(label="Use prompt 2", value=False)
            use_negative_prompt_2 = gr.Checkbox(
                label="Use negative prompt 2", value=False
            )
        negative_prompt = gr.Text(
            label="Negative prompt",
            max_lines=1,
            placeholder="Enter a negative prompt",
            visible=False,
        )
        prompt_2 = gr.Text(
            label="Prompt 2",
            max_lines=1,
            placeholder="Enter your prompt",
            visible=False,
        )
        negative_prompt_2 = gr.Text(
            label="Negative prompt 2",
            max_lines=1,
            placeholder="Enter a negative prompt",
            visible=False,
        )

        seed = gr.Slider(
            label="Seed",
            minimum=0,
            maximum=MAX_SEED,
            step=1,
            value=0,
        )
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        with gr.Row():
            width = gr.Slider(
                label="Width",
                minimum=256,
                maximum=MAX_IMAGE_SIZE,
                step=32,
                value=1024,
            )
            height = gr.Slider(
                label="Height",
                minimum=256,
                maximum=MAX_IMAGE_SIZE,
                step=32,
                value=1024,
            )

        with gr.Row():
            guidance_scale_base = gr.Slider(
                label="Guidance scale for base",
                minimum=1,
                maximum=20,
                step=0.1,
                value=5.0,
            )
            num_inference_steps_base = gr.Slider(
                label="Number of inference steps for base",
                minimum=10,
                maximum=100,
                step=1,
                value=25,
            )
        with gr.Row(visible=False) as refiner_params:
            guidance_scale_refiner = gr.Slider(
                label="Guidance scale for refiner",
                minimum=1,
                maximum=20,
                step=0.1,
                value=5.0,
            )
            num_inference_steps_refiner = gr.Slider(
                label="Number of inference steps for refiner",
                minimum=10,
                maximum=100,
                step=1,
                value=25,
            )

    gr.Examples(
        examples=examples,
        inputs=prompt,
        outputs=result,
        fn=generate,
    )

    gr.Markdown("## Disclaimer")
    gr.Markdown(DISCLAIMER)

    use_negative_prompt.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt,
        outputs=negative_prompt,
        queue=False,
        api_name=False,
    )
    use_prompt_2.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_prompt_2,
        outputs=prompt_2,
        queue=False,
        api_name=False,
    )
    use_negative_prompt_2.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt_2,
        outputs=negative_prompt_2,
        queue=False,
        api_name=False,
    )

    gr.on(
        triggers=[
            prompt.submit,
            negative_prompt.submit,
            prompt_2.submit,
            negative_prompt_2.submit,
            run_button.click,
        ],
        fn=randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate,
        inputs=[
            prompt,
            negative_prompt,
            prompt_2,
            negative_prompt_2,
            use_negative_prompt,
            use_prompt_2,
            use_negative_prompt_2,
            seed,
            width,
            height,
            guidance_scale_base,
            num_inference_steps_base,
            style_image,
            noise_scale,
        ],
        outputs=result,
        api_name="run",
    )

if __name__ == "__main__":
    demo.launch()