r"""Flux plugin.

This plugin depends on :mod:`diffusers` and :mod:`transformers`. To use it,
install the dependencies in your environment

.. code-block:: console

    pip install diffusers transformers accelerate protobuf sentencepiece

before importing the plugin.

.. code-block:: python

    from azula.plugins import flux

References:
    | FLUX (Black Forest Labs, 2024)
    | https://github.com/black-forest-labs/flux
"""

__all__ = [
    "AutoEncoder",
    "TextEncoder",
    "FluxDenoiser",
    "load_model",
]

import torch
import torch.nn as nn

from collections.abc import Sequence
from einops import rearrange
from functools import cache
from torch import Tensor

from azula.denoise import Denoiser, DiracPosterior
from azula.nn.utils import get_module_dtype, skip_init
from azula.noise import DecaySchedule, Schedule

from ..utils import load_cards, patch_diffusers, patch_mmap_safetensors


class AutoEncoder(nn.Module):
    r"""Creates an auto-encoder wrapper."""

    def __init__(
        self,
        vae: nn.Module,  # AutoencoderKL
        shift: float = 0.0,
        scale: float = 1.0,
    ) -> None:
        super().__init__()

        self.vae = vae
        self.shift = shift
        self.scale = scale

    def encode(self, x: Tensor) -> Tensor:
        r"""Encodes images to latents.

        Arguments:
            x: A batch of images :math:`x`, with shape :math:`(B, 3, H, W)`.
                Pixel values are expected to range between -1 and 1.

        Returns:
            A batch of latents :math:`z \sim q(Z \mid x)`, with shape :math:`(B, H / 16, W / 16, 64)`.
        """

        dtype = get_module_dtype(self.vae)

        q_z_x = self.vae.encode(x.to(dtype)).latent_dist
        z = q_z_x.mean + q_z_x.std * torch.randn_like(q_z_x.mean)
        z = (z - self.shift) * self.scale
        z = rearrange(z, "... C (H h) (W w) -> ... H W (C h w)", h=2, w=2)

        return z.to(x)

    def decode(self, z: Tensor) -> Tensor:
        r"""Decodes latents to images.

        Arguments:
            z: A batch of latents :math:`z`, with shape :math:`(B, H / 16, W / 16, 64)`.

        Returns:
            A batch of images :math:`x = D(z)`, with shape :math:`(B, 3, H, W)`.
        """

        dtype = get_module_dtype(self.vae)

        z = rearrange(z, "... H W (C h w) -> ... C (H h) (W w)", h=2, w=2)
        z = z / self.scale + self.shift
        x = self.vae.decode(z.to(dtype)).sample

        return x.to(z)


class TextEncoder(nn.Module):
    r"""Creates a text encoder."""

    def __init__(
        self,
        clip: nn.Module,
        clip_tokenizer: nn.Module,
        t5: nn.Module,
        t5_tokenizer: nn.Module,
    ) -> None:
        super().__init__()

        self.clip = clip
        self.clip_tokenizer = clip_tokenizer
        self.t5 = t5
        self.t5_tokenizer = t5_tokenizer

    def forward(self, prompt: str | Sequence[str]) -> dict[str, Tensor]:
        r"""
        Arguments:
            prompt: A text prompt or list of text prompts.

        Returns:
            The CLIP and T5 encoded prompt(s).
        """

        if isinstance(prompt, str):
            prompt = [prompt]

        # CLIP
        clip_output = self.clip(
            input_ids=self.clip_tokenizer(
                prompt,
                truncation=True,
                max_length=self.clip_tokenizer.model_max_length,
                padding="max_length",
                return_tensors="pt",
            ).input_ids,
            output_hidden_states=False,
        ).pooler_output

        # T5
        t5_output = self.t5(
            input_ids=self.t5_tokenizer(
                prompt,
                truncation=True,
                max_length=self.t5_tokenizer.model_max_length,
                padding="max_length",
                return_tensors="pt",
            ).input_ids,
            output_hidden_states=False,
        ).last_hidden_state

        return {
            "prompt_clip": clip_output,
            "prompt_t5": t5_output,
        }


class FluxDenoiser(Denoiser):
    r"""Creates a Flux denoiser.

    Arguments:
        backbone: A time conditional network.
        schedule: A noise schedule. If :py:`None`, use
            :class:`azula.noise.DecaySchedule` instead.
    """

    def __init__(
        self,
        backbone: nn.Module,
        schedule: Schedule | None = None,
    ) -> None:
        super().__init__()

        self.backbone = backbone

        if schedule is None:
            self.schedule = DecaySchedule()
        else:
            self.schedule = schedule

    @staticmethod
    @cache
    def coordinates(
        H: int,
        W: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> Tensor:
        return torch.stack(
            torch.meshgrid(
                torch.zeros(1, dtype=dtype, device=device),
                torch.arange(H, dtype=dtype, device=device),
                torch.arange(W, dtype=dtype, device=device),
                indexing="ij",
            ),
            dim=-1,
        ).reshape(-1, 3)

    def forward(
        self,
        z_t: Tensor,
        t: Tensor,
        prompt_clip: Tensor,
        prompt_t5: Tensor,
        guidance: float | Tensor = 4.0,
        **kwargs,
    ) -> DiracPosterior:
        r"""
        Arguments:
            z_t: A noisy tensor :math:`z_t`, with shape :math:`(B, H, W, 64)`.
            t: The time :math:`t`, with shape :math:`()` or :math:`(B)`.
            prompt_clip: The CLIP-encoded text prompt :math:`y`, with shape :math:`(B, F)`.
            prompt_t5: The T5-encoded text prompt :math:`y`, with shape :math:`(B, L, D)`.
            guidance: The guidance strength :math:`\omega \in \mathbb{R}`.
            kwargs: Optional keyword arguments.

        Returns:
            The Dirac delta :math:`\delta(Z - \mu_\phi(z_t \mid y))`.
        """

        alpha_t, sigma_t = self.schedule(t)

        while alpha_t.ndim < z_t.ndim:
            alpha_t, sigma_t = alpha_t[..., None], sigma_t[..., None]

        c_in = 1 / (alpha_t + sigma_t)
        c_out = -sigma_t / (alpha_t + sigma_t)
        c_skip = 1 / (alpha_t + sigma_t)
        c_time = (sigma_t / (alpha_t + sigma_t)).flatten()

        B, H, W, C = z_t.shape
        _, L, D = prompt_t5.shape

        dtype = get_module_dtype(self.backbone)

        img_ids = self.coordinates(H, W, dtype=dtype, device=z_t.device)
        txt_ids = torch.zeros((L, 3), dtype=dtype, device=z_t.device)

        if guidance is not None:
            guidance = torch.as_tensor(guidance, dtype=dtype, device=z_t.device).expand(B)

        output = (
            self.backbone(
                timestep=c_time.to(dtype).expand(B),
                hidden_states=(c_in * z_t).to(dtype).reshape(B, H * W, C),
                encoder_hidden_states=prompt_t5.to(dtype).expand(B, L, D),
                pooled_projections=prompt_clip.to(dtype),
                img_ids=img_ids,
                txt_ids=txt_ids,
                guidance=guidance,
                **kwargs,
            )
            .sample.reshape_as(z_t)
            .to(z_t)
        )

        mean = c_skip * z_t + c_out * output

        return DiracPosterior(mean=mean)


def load_model(
    name: str = "flux_1_dev",
    **kwargs,
) -> tuple[Denoiser, AutoEncoder, TextEncoder]:
    r"""Loads a pre-trained Flux latent denoiser.

    Arguments:
        name: The pre-trained model name.
        kwargs: Keyword arguments passed to :func:`diffusers.FluxPipeline.from_pretrained`.

    Returns:
        A pre-trained latent denoiser and the corresponding auto-encoder and text encoder.
    """

    from diffusers import FluxPipeline

    card = load_cards(__name__)[name]

    with skip_init(), patch_diffusers():
        pipe = FluxPipeline.from_pretrained(
            card.repo,
            variant=card.variant,
            torch_dtype=card.dtype_map,
            **kwargs,
        )

    patch_mmap_safetensors(pipe.transformer, pipe.vae, pipe.text_encoder, pipe.text_encoder_2)

    denoiser = FluxDenoiser(backbone=pipe.transformer)

    autoencoder = AutoEncoder(
        vae=pipe.vae,
        shift=pipe.vae.config.shift_factor,
        scale=pipe.vae.config.scaling_factor,
    )

    textencoder = TextEncoder(
        clip=pipe.text_encoder,
        clip_tokenizer=pipe.tokenizer,
        t5=pipe.text_encoder_2,
        t5_tokenizer=pipe.tokenizer_2,
    )

    return denoiser.eval(), autoencoder.eval(), textencoder.eval()
