import random

import gradio as gr
import numpy as np
import torch
import torchvision.transforms as transforms
from torchmetrics.functional.image import structural_similarity_index_measure as ssim
from transformers import CLIPModel, CLIPProcessor

from tools import synth

device = "cuda" if torch.cuda.is_available() else "cpu"
model_path = "runwayml/stable-diffusion-v1-5"
clip_model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(device)
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")

if torch.cuda.is_available():
    torch.cuda.max_memory_allocated(device=device)
    pipe = synth.pipe_img(
        model_path=model_path,
        device=device,
        use_torchcompile=False,
    )
else:
    pipe = synth.pipe_img(
        model_path=model_path,
        device=device,
        apply_optimization=False,
    )

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024


def infer(
    input_image,
    prompt1,
    prompt2,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    interpolation_step,
    num_inference_steps,
    num_interpolation_steps,
    sample_mid_interpolation,
    remove_n_middle,
):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Input Validation
    try:
        assert num_interpolation_steps % 2 == 0
    except AssertionError:
        raise ValueError("num_interpolation_steps must be an even number")
    try:
        assert sample_mid_interpolation % 2 == 0
    except AssertionError:
        raise ValueError("sample_mid_interpolation must be an even number")
    try:
        assert remove_n_middle % 2 == 0
    except AssertionError:
        raise ValueError("remove_n_middle must be an even number")
    try:
        assert num_interpolation_steps >= sample_mid_interpolation
    except AssertionError:
        raise ValueError(
            "num_interpolation_steps must be greater than or equal to sample_mid_interpolation"
        )
    try:
        assert num_interpolation_steps >= 2 and sample_mid_interpolation >= 2
    except AssertionError:
        raise ValueError(
            "num_interpolation_steps and sample_mid_interpolation must be greater than or equal to 2"
        )
    try:
        assert sample_mid_interpolation - remove_n_middle >= 2
    except AssertionError:
        raise ValueError(
            "sample_mid_interpolation must be greater than or equal to remove_n_middle + 2"
        )

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    prompts = [prompt1, prompt2]
    generator = torch.Generator().manual_seed(seed)

    interpolated_prompt_embeds, prompt_metadata = synth.interpolatePrompts(
        prompts,
        pipe,
        num_interpolation_steps,
        sample_mid_interpolation,
        remove_n_middle=remove_n_middle,
        device=device,
    )
    negative_prompts = [negative_prompt, negative_prompt]
    if negative_prompts != ["", ""]:
        interpolated_negative_prompts_embeds, _ = synth.interpolatePrompts(
            negative_prompts,
            pipe,
            num_interpolation_steps,
            sample_mid_interpolation,
            remove_n_middle=remove_n_middle,
            device=device,
        )
    else:
        interpolated_negative_prompts_embeds, _ = [None] * len(
            interpolated_prompt_embeds
        ), None

    latents = torch.randn(
        (1, pipe.unet.config.in_channels, height // 8, width // 8),
        generator=generator,
    ).to(device)
    embed_pairs = zip(interpolated_prompt_embeds, interpolated_negative_prompts_embeds)
    embed_pairs_list = list(embed_pairs)
    print(len(embed_pairs_list))
    # offset step by -1
    prompt_embeds, negative_prompt_embeds = embed_pairs_list[interpolation_step - 1]
    preprocess_input = transforms.Compose(
        [transforms.ToTensor(), transforms.Resize((512, 512))]
    )
    input_img_tensor = preprocess_input(input_image).unsqueeze(0)
    if negative_prompt_embeds is not None:
        npe = negative_prompt_embeds[None, ...]
    else:
        npe = None
    image = pipe(
        height=height,
        width=width,
        num_images_per_prompt=1,
        prompt_embeds=prompt_embeds[None, ...],
        negative_prompt_embeds=npe,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=generator,
        latents=latents,
        image=input_img_tensor,
    ).images[0]
    pred_image = transforms.ToTensor()(image).unsqueeze(0)
    ssim_score = ssim(pred_image, input_img_tensor).item()
    real_inputs = clip_processor(
        text=prompts, padding=True, images=input_image, return_tensors="pt"
    ).to(device)
    real_output = clip_model(**real_inputs)
    synth_inputs = clip_processor(
        text=prompts, padding=True, images=image, return_tensors="pt"
    ).to(device)
    synth_output = clip_model(**synth_inputs)
    cos_sim = torch.nn.CosineSimilarity(dim=1)
    cosine_sim = (
        cos_sim(real_output.image_embeds, synth_output.image_embeds)
        .detach()
        .cpu()
        .numpy()
        .squeeze()
        * 100
    )

    return image, seed, round(ssim_score, 4), round(cosine_sim, 2)


examples1 = [
    "A photo of a chain saw, chainsaw",
    "A photo of a Shih-Tzu, a type of dog",
]
examples2 = [
    "A photo of a golf ball",
    "A photo of a beagle, a type of dog",
]


def update_steps(total_steps, interpolation_step):
    if interpolation_step > total_steps:
        return gr.update(maximum=total_steps // 2, value=total_steps)
    return gr.update(maximum=total_steps // 2)


def update_sampling_steps(total_steps, sample_steps):
    # if sample_steps > total_steps:
    #     return gr.update(value=total_steps)
    return gr.update(value=total_steps)


def update_format(image_format):
    return gr.update(format=image_format)


if torch.cuda.is_available():
    power_device = "GPU"
else:
    power_device = "CPU"

with gr.Blocks(title="Generative Date Augmentation Demo") as demo:

    gr.Markdown(
        """
    # Data Augmentation with Image-to-Image Diffusion Models via Prompt Interpolation.
    Main GitHub Repo: [Generative Data Augmentation](https://github.com/zhulinchng/generative-data-augmentation) | Image Classification Demo: [Generative Augmented Classifiers](https://huggingface.co/spaces/czl/generative-augmented-classifiers).
    """
    )
    with gr.Row():
        with gr.Column():

            input_image = gr.Image(type="pil", label="Image to Augment")

            with gr.Row():
                prompt1 = gr.Text(
                    label="Prompt for the image to synthesize. (Actual class)",
                    show_label=True,
                    max_lines=1,
                    placeholder="Enter your first prompt",
                    container=False,
                )
            with gr.Row():
                prompt2 = gr.Text(
                    label="Prompt to augment against. (Confusing class)",
                    show_label=True,
                    max_lines=1,
                    placeholder="Enter your second prompt",
                    container=False,
                )
            with gr.Row():
                gr.Examples(
                    examples=examples1, inputs=[prompt1], label="Example for Prompt 1"
                )
                gr.Examples(
                    examples=examples2, inputs=[prompt2], label="Example for Prompt 2"
                )

            with gr.Row():
                interpolation_step = gr.Slider(
                    label="Specific Interpolation Step",
                    minimum=1,
                    maximum=8,
                    step=1,
                    value=8,
                )
                num_interpolation_steps = gr.Slider(
                    label="Total interpolation steps",
                    minimum=2,
                    maximum=32,
                    step=2,
                    value=16,
                )
                num_interpolation_steps.change(
                    fn=update_steps,
                    inputs=[num_interpolation_steps, interpolation_step],
                    outputs=[interpolation_step],
                )
                run_button = gr.Button("Run", scale=0)
            with gr.Accordion("Advanced Settings", open=True):
                negative_prompt = gr.Text(
                    label="Negative prompt",
                    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)
                gr.Markdown("Negative Prompt: ")
                with gr.Row():
                    negative_prompt = gr.Text(
                        label="Negative Prompt",
                        show_label=True,
                        max_lines=1,
                        value="blurry image, disfigured, deformed, distorted, cartoon, drawings",
                        container=False,
                    )
                with gr.Row():
                    width = gr.Slider(
                        label="Width",
                        minimum=256,
                        maximum=MAX_IMAGE_SIZE,
                        step=32,
                        value=512,
                    )
                    height = gr.Slider(
                        label="Height",
                        minimum=256,
                        maximum=MAX_IMAGE_SIZE,
                        step=32,
                        value=512,
                    )
                with gr.Row():
                    guidance_scale = gr.Slider(
                        label="Guidance scale",
                        minimum=0.0,
                        maximum=10.0,
                        step=0.1,
                        value=8.0,
                    )
                    num_inference_steps = gr.Slider(
                        label="Number of inference steps",
                        minimum=1,
                        maximum=80,
                        step=1,
                        value=25,
                    )
                with gr.Row():
                    sample_mid_interpolation = gr.Slider(
                        label="Number of sampling steps in the middle of interpolation",
                        minimum=2,
                        maximum=80,
                        step=2,
                        value=16,
                    )
                    num_interpolation_steps.change(
                        fn=update_sampling_steps,
                        inputs=[num_interpolation_steps, sample_mid_interpolation],
                        outputs=[sample_mid_interpolation],
                    )
                with gr.Row():
                    remove_n_middle = gr.Slider(
                        label="Number of middle steps to remove from interpolation",
                        minimum=0,
                        maximum=80,
                        step=2,
                        value=0,
                    )
                with gr.Row():
                    image_type = gr.Radio(
                        choices=[
                            "webp",
                            "png",
                            "jpeg",
                        ],
                        label="Download Image Format",
                        value="jpeg",
                    )
        with gr.Column():
            result = gr.Image(label="Result", show_label=False, format="jpeg")
            image_type.change(
                fn=update_format,
                inputs=[image_type],
                outputs=[result],
            )
            gr.Markdown(
                """
                Metadata:
                """
            )
            with gr.Row():
                show_seed = gr.Label(label="Seed:", value="Randomized seed")
                ssim_score = gr.Label(
                    label="SSIM Score:", value="Generate to see score"
                )
                cos_sim = gr.Label(label="CLIP Score:", value="Generate to see score")
            if power_device == "GPU":
                gr.Markdown(
                    f"""
Currently running on {power_device}.
                    """
                )
            else:
                gr.Markdown(
                    f"""
Currently running on {power_device}.
Note: Running on CPU will take longer (approx. 6 minutes with default settings).
                    """
                )
            gr.Markdown(
                """
This demo is created as part of the 'Investigating the Effectiveness of Generative Diffusion Models in Synthesizing Images for Data Augmentation in Image Classification' dissertation.

The user can augment an image by interpolating between two prompts, and specify the number of interpolation steps and the specific step to generate the image.
            """
            )
        run_button.click(
            fn=infer,
            inputs=[
                input_image,
                prompt1,
                prompt2,
                negative_prompt,
                seed,
                randomize_seed,
                width,
                height,
                guidance_scale,
                interpolation_step,
                num_inference_steps,
                num_interpolation_steps,
                sample_mid_interpolation,
                remove_n_middle,
            ],
            outputs=[result, show_seed, ssim_score, cos_sim],
        )

demo.queue().launch(show_error=True)

"""
    input_image,
    prompt1,
    prompt2,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    interpolation_step,
    num_inference_steps,
    num_interpolation_steps,
    sample_mid_interpolation,
    remove_n_middle,
    """