import gradio as gr

from absl import flags
from absl import app
from ml_collections import config_flags
import os

import spaces #[uncomment to use ZeroGPU]
import torch


import io
import random
import tempfile

import numpy as np
import torch
import torch.nn.functional as F
from torchvision.utils import save_image
from torchvision.transforms import ToPILImage
from huggingface_hub import hf_hub_download

from absl import logging
import ml_collections

from diffusion.flow_matching import ODEEulerFlowMatchingSolver
import utils
import libs.autoencoder
from libs.clip import FrozenCLIPEmbedder
from configs import t2i_512px_clip_dimr, t2i_256px_clip_dimr


def unpreprocess(x: torch.Tensor) -> torch.Tensor:
    x = 0.5 * (x + 1.0)
    x.clamp_(0.0, 1.0)
    return x
    
def cosine_similarity_torch(latent1: torch.Tensor, latent2: torch.Tensor) -> torch.Tensor:
    latent1_flat = latent1.view(-1)
    latent2_flat = latent2.view(-1)
    cosine_similarity = F.cosine_similarity(
        latent1_flat.unsqueeze(0), latent2_flat.unsqueeze(0), dim=1
    )
    return cosine_similarity

def kl_divergence(latent1: torch.Tensor, latent2: torch.Tensor) -> torch.Tensor:
    latent1_prob = F.softmax(latent1, dim=-1)
    latent2_prob = F.softmax(latent2, dim=-1)
    latent1_log_prob = torch.log(latent1_prob)
    kl_div = F.kl_div(latent1_log_prob, latent2_prob, reduction="batchmean")
    return kl_div

def batch_decode(_z: torch.Tensor, decode, batch_size: int = 10) -> torch.Tensor:
    num_samples = _z.size(0)
    decoded_batches = []

    for i in range(0, num_samples, batch_size):
        batch = _z[i : i + batch_size]
        decoded_batch = decode(batch)
        decoded_batches.append(decoded_batch)

    return torch.cat(decoded_batches, dim=0)

def get_caption(llm: str, text_model, prompt_dict: dict, batch_size: int):
    if batch_size == 3:
        # Only addition or only subtraction mode.
        assert len(prompt_dict) == 2, "Expected 2 prompts for batch_size 3."
        batch_prompts = list(prompt_dict.values()) + [" "]
    elif batch_size == 4:
        # Addition and subtraction mode.
        assert len(prompt_dict) == 3, "Expected 3 prompts for batch_size 4."
        batch_prompts = list(prompt_dict.values()) + [" "]
    elif batch_size >= 5:
        # Linear interpolation mode.
        assert len(prompt_dict) == 2, "Expected 2 prompts for linear interpolation."
        batch_prompts = [prompt_dict["prompt_1"]] + [" "] * (batch_size - 2) + [prompt_dict["prompt_2"]]
    else:
        raise ValueError(f"Unsupported batch_size: {batch_size}")

    if llm == "clip":
        latent, latent_and_others = text_model.encode(batch_prompts)
        context = latent_and_others["token_embedding"].detach()
    elif llm == "t5":
        latent, latent_and_others = text_model.get_text_embeddings(batch_prompts)
        context = (latent_and_others["token_embedding"] * 10.0).detach()
    else:
        raise NotImplementedError(f"Language model {llm} not supported.")

    token_mask = latent_and_others["token_mask"].detach()
    tokens = latent_and_others["tokens"].detach()
    captions = batch_prompts

    return context, token_mask, tokens, captions

# Load configuration and initialize models.
# config_dict = t2i_512px_clip_dimr.get_config()
config_dict = t2i_256px_clip_dimr.get_config()
config = ml_collections.ConfigDict(config_dict)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info(f"Using device: {device}")

# Freeze configuration.
config = ml_collections.FrozenConfigDict(config)

torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024  # Currently not used.

# Load the main diffusion model.
repo_id = "QHL067/CrossFlow"
# filename = "pretrained_models/t2i_512px_clip_dimr.pth"
filename = "pretrained_models/t2i_256px_clip_dimr.pth"
checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
nnet = utils.get_nnet(**config.nnet)
nnet = nnet.to(device)
state_dict = torch.load(checkpoint_path, map_location=device)
nnet.load_state_dict(state_dict)
nnet.eval()

# Initialize text model.
llm = "clip"
clip = FrozenCLIPEmbedder()
clip.eval()
clip.to(device)

# Load autoencoder.
autoencoder = libs.autoencoder.get_model(**config.autoencoder)
autoencoder.to(device)


@torch.cuda.amp.autocast()
def encode(_batch: torch.Tensor) -> torch.Tensor:
    """Encode a batch of images using the autoencoder."""
    return autoencoder.encode(_batch)


@torch.cuda.amp.autocast()
def decode(_batch: torch.Tensor) -> torch.Tensor:
    """Decode a batch of latent vectors using the autoencoder."""
    return autoencoder.decode(_batch)


@spaces.GPU #[uncomment to use ZeroGPU]
def infer(
    prompt1,
    prompt2,
    seed,
    randomize_seed,
    guidance_scale,
    num_inference_steps,
    num_of_interpolation,
    save_gpu_memory=True,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    torch.manual_seed(seed)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(seed)

    # Only support interpolation in this implementation.
    prompt_dict = {"prompt_1": prompt1, "prompt_2": prompt2}
    for key, value in prompt_dict.items():
        assert value is not None, f"{key} must not be None."
    assert num_of_interpolation >= 5, "For linear interpolation, please sample at least five images."

    # Get text embeddings and tokens.
    _context, _token_mask, _token, _caption = get_caption(
        llm, clip, prompt_dict=prompt_dict, batch_size=num_of_interpolation
    )

    with torch.no_grad():
        _z_gaussian = torch.randn(num_of_interpolation, *config.z_shape, device=device)
        _z_x0, _mu, _log_var = nnet(
            _context, text_encoder=True, shape=_z_gaussian.shape, mask=_token_mask
        )
        _z_init = _z_x0.reshape(_z_gaussian.shape)

        # Prepare the initial latent representations based on the number of interpolations.
        if num_of_interpolation == 3:
            # Addition or subtraction mode.
            if config.prompt_a is not None:
                assert config.prompt_s is None, "Only one of prompt_a or prompt_s should be provided."
                z_init_temp = _z_init[0] + _z_init[1]
            elif config.prompt_s is not None:
                assert config.prompt_a is None, "Only one of prompt_a or prompt_s should be provided."
                z_init_temp = _z_init[0] - _z_init[1]
            else:
                raise NotImplementedError("Either prompt_a or prompt_s must be provided for 3-sample mode.")
            mean = z_init_temp.mean()
            std = z_init_temp.std()
            _z_init[2] = (z_init_temp - mean) / std

        elif num_of_interpolation == 4:
            z_init_temp = _z_init[0] + _z_init[1] - _z_init[2]
            mean = z_init_temp.mean()
            std = z_init_temp.std()
            _z_init[3] = (z_init_temp - mean) / std

        elif num_of_interpolation >= 5:
            tensor_a = _z_init[0]
            tensor_b = _z_init[-1]
            num_interpolations = num_of_interpolation - 2
            interpolations = [
                tensor_a + (tensor_b - tensor_a) * (i / (num_interpolations + 1))
                for i in range(1, num_interpolations + 1)
            ]
            _z_init = torch.stack([tensor_a] + interpolations + [tensor_b], dim=0)

        else:
            raise ValueError("Unsupported number of interpolations.")

        assert guidance_scale > 1, "Guidance scale must be greater than 1."

        has_null_indicator = hasattr(config.nnet.model_args, "cfg_indicator")
        ode_solver = ODEEulerFlowMatchingSolver(
            nnet,
            bdv_model_fn=None,
            step_size_type="step_in_dsigma",
            guidance_scale=guidance_scale,
        )
        _z, _ = ode_solver.sample(
            x_T=_z_init,
            batch_size=num_of_interpolation,
            sample_steps=num_inference_steps,
            unconditional_guidance_scale=guidance_scale,
            has_null_indicator=has_null_indicator,
        )

        if save_gpu_memory:
            image_unprocessed = batch_decode(_z, decode)
        else:
            image_unprocessed = decode(_z)

        samples = unpreprocess(image_unprocessed).contiguous()


    to_pil = ToPILImage()
    pil_images = [to_pil(img) for img in samples]

    first_image = pil_images[0]
    last_image = pil_images[-1]

    gif_buffer = io.BytesIO()
    pil_images[0].save(gif_buffer, format="GIF", save_all=True, append_images=pil_images[1:], duration=10, loop=0)
    gif_buffer.seek(0)
    gif_bytes = gif_buffer.read()

    # Save the GIF bytes to a temporary file and get its path
    temp_gif = tempfile.NamedTemporaryFile(delete=False, suffix=".gif")
    temp_gif.write(gif_bytes)
    temp_gif.close()
    gif_path = temp_gif.name

    return first_image, last_image, gif_path, seed
    # return first_image, last_image, seed


# examples = [
#     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
#     "An astronaut riding a green horse",
#     "A delicious ceviche cheesecake slice",
# ]

examples = [
    ["A dog cooking dinner in the kitchen", "An orange cat wearing sunglasses on a ship"],
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""
with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# CrossFlow")
        gr.Markdown("CrossFlow directly transforms text representations into images for text-to-image generation, enabling interpolation in the input text latent space.")

        with gr.Row():
            prompt1 = gr.Text(
                label="Prompt_1",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt for the first image",
                container=False,
            )
        
        with gr.Row():
            prompt2 = gr.Text(
                label="Prompt_2",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt for the second image",
                container=False,
            )

        with gr.Row():
            run_button = gr.Button("Run", scale=0, variant="primary")

        # Create separate outputs for the first image, last image, and the animated GIF
        first_image_output = gr.Image(label="Image if the first prompt", show_label=True)
        last_image_output = gr.Image(label="Image if the second prompt", show_label=True)
        gif_output = gr.Image(label="Linear interpolation", show_label=True)

        with gr.Accordion("Advanced Settings", open=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():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=7.0,  # Replace with defaults that work for your model
                )
            with gr.Row():
                num_inference_steps = gr.Slider(
                    label="Number of inference steps - 50 inference steps are recommended; but you can reduce to 20 if the demo fails.",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=50,  # Replace with defaults that work for your model
                )
            with gr.Row():
                num_of_interpolation = gr.Slider(
                    label="Number of images for interpolation - More images yield smoother transitions but require more resources and may fail.",
                    minimum=5,
                    maximum=50,
                    step=1,
                    value=5,  # Replace with defaults that work for your model
                )

        gr.Examples(examples=examples, inputs=[prompt1, prompt2])
    gr.on(
        triggers=[run_button.click, prompt1.submit, prompt2.submit],
        fn=infer,
        inputs=[
            prompt1,
            prompt2,
            seed,
            randomize_seed,
            guidance_scale,
            num_inference_steps,
            num_of_interpolation,
        ],
        outputs=[first_image_output, last_image_output, gif_output, seed],
        # outputs=[first_image_output, last_image_output, seed],
    )


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