"""
Lookup Free Quantization
Proposed in https://arxiv.org/abs/2310.05737

In the simplest setup, each dimension is quantized into {-1, 1}.
An entropy penalty is used to encourage utilization.
"""

from math import log2, ceil
import random
from functools import partial

import torch
from torch import nn, einsum
import torch.nn.functional as F
from torch.nn import Module
from torch.cuda.amp import autocast

from einops import rearrange, reduce, pack, repeat, unpack

# constants

# helper functions


def exists(v):
    return v is not None


def default(*args):
    for arg in args:
        if exists(arg):
            return arg() if callable(arg) else arg
    return None


def pack_one(t, pattern):
    return pack([t], pattern)


def unpack_one(t, ps, pattern):
    return unpack(t, ps, pattern)[0]


# entropy


def log(t, eps=1e-5):
    return t.clamp(min=eps).log()


def entropy(prob):
    return (-prob * log(prob)).sum(dim=-1)


# class


class LFQ(Module):
    def __init__(
        self,
        *,
        dim=None,
        codebook_size=None,
        entropy_loss_weight=0.1,
        commitment_loss_weight=0.25,
        diversity_gamma=3.0,
        straight_through_activation=nn.Identity(),
        num_codebooks=1,
        keep_num_codebooks_dim=None,
        codebook_scale=1.0,  # for residual LFQ, codebook scaled down by 2x at each layer
        quantize_dropout=False,
    ):
        super().__init__()

        # some assert validations

        assert exists(dim) or exists(
            codebook_size
        ), "either dim or codebook_size must be specified for LFQ"
        assert (
            not exists(codebook_size) or log2(codebook_size).is_integer()
        ), f"your codebook size must be a power of 2 for lookup free quantization (suggested {2 ** ceil(log2(codebook_size))})"

        codebook_size = default(codebook_size, lambda: 2**dim)
        codebook_dim = int(log2(codebook_size))

        codebook_dims = codebook_dim * num_codebooks
        dim = default(dim, codebook_dims)

        has_projections = dim != codebook_dims
        self.project_in = (
            nn.Linear(dim, codebook_dims) if has_projections else nn.Identity()
        )
        self.project_out = (
            nn.Linear(codebook_dims, dim) if has_projections else nn.Identity()
        )
        self.has_projections = has_projections

        self.dim = dim
        self.codebook_dim = codebook_dim
        self.num_codebooks = num_codebooks

        keep_num_codebooks_dim = default(keep_num_codebooks_dim, num_codebooks > 1)
        assert not (num_codebooks > 1 and not keep_num_codebooks_dim)
        self.keep_num_codebooks_dim = keep_num_codebooks_dim

        # straight through activation

        self.activation = straight_through_activation

        # entropy aux loss related weights

        self.diversity_gamma = diversity_gamma
        self.entropy_loss_weight = entropy_loss_weight

        # codebook scale

        self.codebook_scale = codebook_scale

        # commitment loss

        self.commitment_loss_weight = commitment_loss_weight

        # for no auxiliary loss, during inference

        self.register_buffer("mask", 2 ** torch.arange(codebook_dim - 1, -1, -1))
        self.register_buffer("zero", torch.tensor(0.0), persistent=False)

        # codes

        all_codes = torch.arange(codebook_size)
        bits = ((all_codes[..., None].int() & self.mask) != 0).float()
        codebook = self.bits_to_codes(bits)

        self.register_buffer("codebook", codebook, persistent=False)

        self.quantize_dropout = quantize_dropout and num_codebooks > 1

    def bits_to_codes(self, bits):
        return bits * self.codebook_scale * 2 - self.codebook_scale

    @property
    def dtype(self):
        return self.codebook.dtype

    def indices_to_codes(self, indices, project_out=True):
        is_img_or_video = indices.ndim >= (3 + int(self.keep_num_codebooks_dim))

        if not self.keep_num_codebooks_dim:
            indices = rearrange(indices, "... -> ... 1")

        # indices to codes, which are bits of either -1 or 1

        bits = ((indices[..., None].int() & self.mask) != 0).to(self.dtype)

        codes = self.bits_to_codes(bits)

        codes = rearrange(codes, "... c d -> ... (c d)")

        # whether to project codes out to original dimensions
        # if the input feature dimensions were not log2(codebook size)

        if project_out:
            codes = self.project_out(codes)

        # rearrange codes back to original shape

        if is_img_or_video:
            codes = rearrange(codes, "b ... d -> b d ...")

        return codes

    def forward(self, x, inv_temperature=100.0, mask=None, transpose=True, **kwargs):
        """
        einstein notation
        b - batch
        n - sequence (or flattened spatial dimensions)
        d - feature dimension, which is also log2(codebook size)
        c - number of codebook dim
        """
        if transpose:
            x = rearrange(x, "b d n -> b n d")
        is_img_or_video = x.ndim >= 4

        # standardize image or video into (batch, seq, dimension)

        if is_img_or_video:
            x = rearrange(x, "b d ... -> b ... d")
            x, ps = pack_one(x, "b * d")

        assert (
            x.shape[-1] == self.dim
        ), f"expected dimension of {self.dim} but received {x.shape[-1]}"

        x = self.project_in(x)

        # split out number of codebooks

        x = rearrange(x, "b n (c d) -> b n c d", c=self.num_codebooks)

        # quantize by eq 3.

        original_input = x

        codebook_value = torch.ones_like(x) * self.codebook_scale
        quantized = torch.where(x > 0, codebook_value, -codebook_value)

        # use straight-through gradients (optionally with custom activation fn) if training

        if self.training:
            x = self.activation(x)
            x = x + (quantized - x).detach()
        else:
            x = quantized

        # calculate indices

        indices = reduce((x > 0).int() * self.mask.int(), "b n c d -> b n c", "sum")

        # entropy aux loss

        if self.training:
            # the same as euclidean distance up to a constant
            distance = -2 * einsum(
                "... i d, j d -> ... i j", original_input, self.codebook
            )

            prob = (-distance * inv_temperature).softmax(dim=-1)

            per_sample_entropy = entropy(prob).mean()

            # account for mask

            if exists(mask):
                prob = prob[mask]

            # distribution over all available tokens in the batch

            avg_prob = reduce(prob, "... c d -> c d", "mean")
            codebook_entropy = entropy(avg_prob).mean()

            # 1. entropy will be nudged to be low for each code, to encourage the network to output confident predictions
            # 2. codebook entropy will be nudged to be high, to encourage all codes to be uniformly used within the batch

            entropy_aux_loss = (
                per_sample_entropy - self.diversity_gamma * codebook_entropy
            )
        else:
            # if not training, just return dummy 0
            entropy_aux_loss = per_sample_entropy = codebook_entropy = self.zero

        # commit loss

        if self.training:
            commit_loss = F.mse_loss(
                original_input, quantized.detach(), reduction="none"
            )

            if exists(mask):
                commit_loss = commit_loss[mask]

            commit_loss = commit_loss.mean()
        else:
            commit_loss = self.zero

        # merge back codebook dim

        x = rearrange(x, "b n c d -> b n (c d)")

        # project out to feature dimension if needed

        x = self.project_out(x)

        # reconstitute image or video dimensions

        if is_img_or_video:
            x = unpack_one(x, ps, "b * d")
            x = rearrange(x, "b ... d -> b d ...")

            indices = unpack_one(indices, ps, "b * c")

        # whether to remove single codebook dim

        if not self.keep_num_codebooks_dim:
            indices = rearrange(indices, "... 1 -> ...")

        if transpose:
            x = rearrange(x, "b ... d -> b d ...")
            original_input = rearrange(original_input, "b ... d -> b d ...")

        # complete aux loss
        aux_loss = (
            entropy_aux_loss * self.entropy_loss_weight
            + commit_loss * self.commitment_loss_weight
        )

        # return ret, LossBreakdown(per_sample_entropy, codebook_entropy, commit_loss)
        return {
            "z": x,
            "codes": indices,
            "latents": original_input,
            "vq/commitment_loss": commit_loss,
            "vq/entropy_loss": entropy_aux_loss,
            "vq/codebook_entropy": codebook_entropy,
            "aux_loss": aux_loss,
        }


def round_up_multiple(num, mult):
    return ceil(num / mult) * mult


# main class


class ResidualLFQ(Module):
    """Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf"""

    def __init__(
        self,
        *,
        dim,
        num_quantizers,
        codebook_size,
        quantize_dropout=False,
        quantize_dropout_cutoff_index=0,
        quantize_dropout_multiple_of=1,
        **kwargs,
    ):
        super().__init__()
        codebook_dim = int(log2(codebook_size))

        requires_projection = codebook_dim != dim
        self.project_in = (
            nn.Linear(dim, codebook_dim) if requires_projection else nn.Identity()
        )
        self.project_out = (
            nn.Linear(codebook_dim, dim) if requires_projection else nn.Identity()
        )
        self.has_projections = requires_projection

        self.num_quantizers = num_quantizers

        self.layers = nn.ModuleList([])

        for ind in range(num_quantizers):
            codebook_scale = 2**-ind

            lfq = LFQ(dim=codebook_dim, codebook_scale=codebook_scale, **kwargs)

            self.layers.append(lfq)

        assert all([not lfq.has_projections for lfq in self.layers])

        self.quantize_dropout = quantize_dropout and num_quantizers > 1

        assert quantize_dropout_cutoff_index >= 0

        self.quantize_dropout_cutoff_index = quantize_dropout_cutoff_index
        self.quantize_dropout_multiple_of = quantize_dropout_multiple_of  # encodec paper proposes structured dropout, believe this was set to 4

    @property
    def codebooks(self):
        codebooks = [layer.codebook for layer in self.layers]
        codebooks = torch.stack(codebooks, dim=0)
        return codebooks

    def get_codes_from_indices(self, indices):
        batch, quantize_dim = indices.shape[0], indices.shape[-1]

        # may also receive indices in the shape of 'b h w q' (accept_image_fmap)

        indices, ps = pack([indices], "b * q")

        # because of quantize dropout, one can pass in indices that are coarse
        # and the network should be able to reconstruct

        if quantize_dim < self.num_quantizers:
            assert (
                self.quantize_dropout > 0.0
            ), "quantize dropout must be greater than 0 if you wish to reconstruct from a signal with less fine quantizations"
            indices = F.pad(indices, (0, self.num_quantizers - quantize_dim), value=-1)

        # get ready for gathering

        codebooks = repeat(self.codebooks, "q c d -> q b c d", b=batch)
        gather_indices = repeat(indices, "b n q -> q b n d", d=codebooks.shape[-1])

        # take care of quantizer dropout

        mask = gather_indices == -1.0
        gather_indices = gather_indices.masked_fill(
            mask, 0
        )  # have it fetch a dummy code to be masked out later

        all_codes = codebooks.gather(2, gather_indices)  # gather all codes

        # mask out any codes that were dropout-ed

        all_codes = all_codes.masked_fill(mask, 0.0)

        # if (accept_image_fmap = True) then return shape (quantize, batch, height, width, dimension)

        (all_codes,) = unpack(all_codes, ps, "q b * d")

        return all_codes

    def get_output_from_indices(self, indices):
        codes = self.get_codes_from_indices(indices)
        codes_summed = reduce(codes, "q ... -> ...", "sum")
        return self.project_out(codes_summed)

    def forward(
        self,
        x,
        mask=None,
        rand_quantize_dropout_fixed_seed=None,
        **kwargs,
    ):
        num_quant, quant_dropout_multiple_of, device = (
            self.num_quantizers,
            self.quantize_dropout_multiple_of,
            x.device,
        )

        x = rearrange(x, "b d ... -> b ... d")
        x = self.project_in(x)
        x = torch.tanh(x)

        quantized_out = 0.0
        residual = x

        all_aux_losses = []
        all_commitment_losses = []
        all_entropy_losses = []
        all_indices = []

        should_quantize_dropout = self.training and self.quantize_dropout

        # sample a layer index at which to dropout further residual quantization
        # also prepare null indices and loss

        if should_quantize_dropout:
            rand = (
                random.Random(rand_quantize_dropout_fixed_seed)
                if exists(rand_quantize_dropout_fixed_seed)
                else random
            )

            rand_quantize_dropout_index = rand.randrange(
                self.quantize_dropout_cutoff_index, num_quant
            )

            if quant_dropout_multiple_of != 1:
                rand_quantize_dropout_index = (
                    round_up_multiple(
                        rand_quantize_dropout_index + 1, quant_dropout_multiple_of
                    )
                    - 1
                )

            null_indices = torch.full(
                x.shape[:2], -1.0, device=device, dtype=torch.long
            )

        # go through the layers

        with autocast(enabled=False):
            for quantizer_index, layer in enumerate(self.layers):
                if (
                    should_quantize_dropout
                    and quantizer_index > rand_quantize_dropout_index
                ):
                    all_indices.append(null_indices)
                    continue

                assert isinstance(layer, LFQ)
                layer_out = layer.forward(residual, mask=mask, transpose=False)
                quantized, indices = (layer_out["z"], layer_out["codes"])

                residual = residual - quantized.detach()
                quantized_out = quantized_out + quantized

                all_indices.append(indices)
                all_commitment_losses.append(layer_out["vq/commitment_loss"])
                all_entropy_losses.append(layer_out["vq/entropy_loss"])
                all_aux_losses.append(layer_out["aux_loss"])

        # project out, if needed

        quantized_out = self.project_out(quantized_out)

        # stack all losses and indices
        all_indices = torch.stack(all_indices, dim=-1)
        all_commitment_losses = torch.stack(all_commitment_losses, dim=-1)
        all_aux_losses = torch.stack(all_aux_losses, dim=-1)
        all_entropy_losses = torch.stack(all_entropy_losses, dim=-1)

        quantized_out = rearrange(quantized_out, "b ... d -> b d ...")
        x = rearrange(x, "b ... d -> b d ...")

        return {
            "z": quantized_out,
            "codes": all_indices,
            "latents": x,
            "vq/commitment_loss": all_commitment_losses.mean(),
            "vq/entropy_loss": all_entropy_losses.mean(),
            "aux_loss": all_aux_losses.mean(),
        }


class GroupedLFQ(Module):
    def __init__(
        self,
        *,
        dim,
        num_quantizers,
        codebook_size,
        quantize_dropout=False,
        quantize_dropout_cutoff_index=0,
        quantize_dropout_multiple_of=1,
        **kwargs,
    ):
        super().__init__()
        codebook_dim = int(log2(codebook_size))

        requires_projection = codebook_dim != dim
        self.project_in = (
            nn.Linear(dim, codebook_dim * num_quantizers)
            if requires_projection
            else nn.Identity()
        )
        self.project_out = (
            nn.Linear(codebook_dim * num_quantizers, dim)
            if requires_projection
            else nn.Identity()
        )
        self.has_projections = requires_projection

        self.num_quantizers = num_quantizers

        self.layers = nn.ModuleList([])

        for ind in range(num_quantizers):
            lfq = LFQ(dim=codebook_dim, codebook_size=codebook_size, **kwargs)

            self.layers.append(lfq)

        self.quantize_dropout = quantize_dropout and num_quantizers > 1

        assert quantize_dropout_cutoff_index >= 0

        self.quantize_dropout_cutoff_index = quantize_dropout_cutoff_index
        self.quantize_dropout_multiple_of = quantize_dropout_multiple_of  # encodec paper proposes structured dropout, believe this was set to 4

    @property
    def codebooks(self):
        codebooks = [layer.codebook for layer in self.layers]
        codebooks = torch.stack(codebooks, dim=0)
        return codebooks

    def get_codes_from_indices(self, indices):
        batch, quantize_dim = indices.shape[0], indices.shape[-1]

        # may also receive indices in the shape of 'b h w q' (accept_image_fmap)

        indices, ps = pack([indices], "b * q")

        # because of quantize dropout, one can pass in indices that are coarse
        # and the network should be able to reconstruct

        if quantize_dim < self.num_quantizers:
            assert (
                self.quantize_dropout > 0.0
            ), "quantize dropout must be greater than 0 if you wish to reconstruct from a signal with less fine quantizations"
            indices = F.pad(indices, (0, self.num_quantizers - quantize_dim), value=-1)

        # get ready for gathering

        codebooks = repeat(self.codebooks, "q c d -> q b c d", b=batch)
        gather_indices = repeat(indices, "b n q -> q b n d", d=codebooks.shape[-1])

        # take care of quantizer dropout

        mask = gather_indices == -1.0
        gather_indices = gather_indices.masked_fill(
            mask, 0
        )  # have it fetch a dummy code to be masked out later

        all_codes = codebooks.gather(2, gather_indices)  # gather all codes

        # mask out any codes that were dropout-ed

        all_codes = all_codes.masked_fill(mask, 0.0)

        # if (accept_image_fmap = True) then return shape (quantize, batch, height, width, dimension)

        (all_codes,) = unpack(all_codes, ps, "q b * d")

        return all_codes

    def get_output_from_indices(self, indices):
        outputs = []
        for layer in self.layers:
            output = layer.indices_to_codes(indices)
            outputs.append(output)
        outputs = torch.stack(outputs, dim=-2)
        return outputs

    def forward(
        self,
        x,
        mask=None,
        rand_quantize_dropout_fixed_seed=None,
        **kwargs,
    ):
        num_quant, quant_dropout_multiple_of, device = (
            self.num_quantizers,
            self.quantize_dropout_multiple_of,
            x.device,
        )

        x = rearrange(x, "b d ... -> b ... d")
        x = self.project_in(x)
        x = torch.tanh(x)
        # split x into num_quantizers
        x = rearrange(x, "b n (c d) -> b n c d", c=self.num_quantizers)

        quantized_out = torch.zeros_like(x)

        all_aux_losses = []
        all_commitment_losses = []
        all_entropy_losses = []
        all_codebook_entropy_losses = []
        all_indices = []

        should_quantize_dropout = self.training and self.quantize_dropout

        # sample a layer index at which to dropout further residual quantization
        # also prepare null indices and loss

        if should_quantize_dropout:
            rand = (
                random.Random(rand_quantize_dropout_fixed_seed)
                if exists(rand_quantize_dropout_fixed_seed)
                else random
            )

            rand_quantize_dropout_index = rand.randrange(
                self.quantize_dropout_cutoff_index, num_quant
            )

            if quant_dropout_multiple_of != 1:
                rand_quantize_dropout_index = (
                    round_up_multiple(
                        rand_quantize_dropout_index + 1, quant_dropout_multiple_of
                    )
                    - 1
                )

            null_indices = torch.full(
                x.shape[:2], -1.0, device=device, dtype=torch.long
            )

        # go through the layers

        with autocast(enabled=False):
            for quantizer_index, layer in enumerate(self.layers):
                if (
                    should_quantize_dropout
                    and quantizer_index > rand_quantize_dropout_index
                ):
                    all_indices.append(null_indices)
                    continue

                assert isinstance(layer, LFQ)
                x_i = x[:, :, quantizer_index, :]
                layer_out = layer.forward(x_i, mask=mask, transpose=False)
                quantized, indices = (layer_out["z"], layer_out["codes"])

                quantized_out[:, :, quantizer_index, :] = quantized

                all_indices.append(indices)
                all_commitment_losses.append(layer_out["vq/commitment_loss"])
                all_entropy_losses.append(layer_out["vq/entropy_loss"])
                all_codebook_entropy_losses.append(layer_out["vq/codebook_entropy"])
                all_aux_losses.append(layer_out["aux_loss"])

        # project out, if needed
        quantized_out = rearrange(quantized_out, "b n c d -> b n (c d)")
        quantized_out = self.project_out(quantized_out)

        # stack all losses and indices
        all_indices = torch.stack(all_indices, dim=-1)
        all_commitment_losses = torch.stack(all_commitment_losses, dim=-1)
        all_aux_losses = torch.stack(all_aux_losses, dim=-1)
        all_entropy_losses = torch.stack(all_entropy_losses, dim=-1)
        all_codebook_entropy_losses = torch.stack(all_codebook_entropy_losses, dim=-1)

        quantized_out = rearrange(quantized_out, "b ... d -> b d ...")
        x = rearrange(x, "b ... d -> b d ...")

        return {
            "z": quantized_out,
            "codes": all_indices,
            "latents": x,
            "vq/commitment_loss": all_commitment_losses.mean(),
            "vq/entropy_loss": all_entropy_losses.mean(),
            "vq/codebook_entropy": all_codebook_entropy_losses.mean(),
            "aux_loss": all_aux_losses.mean(),
        }