from __future__ import annotations

import random
from typing import Callable, List, Optional, Tuple, Union

import torch
import torch.nn as nn
from torch import Tensor


def round_ste(z: Tensor) -> Tensor:
    zhat = z.round()
    return z + (zhat - z).detach()


def round_ste_with_dropout(z: Tensor, drop_quant_p: float) -> Tensor:
    if drop_quant_p <= 0.0:
        return round_ste(z)
    zhat = z.round()
    b = z.shape[0]
    mask = torch.bernoulli(torch.full((b,), drop_quant_p, device=z.device)).view(
        b, *([1] * (z.ndim - 1))
    )
    return z + ((1 - mask) * (zhat - z)).detach()


class FSQ(nn.Module):
    """Finite Scalar Quantizer with index mapping compatible with OAT-style token IDs."""

    def __init__(
        self,
        levels: List[int],
        drop_quant_p: float = 0.0,
        corrupt_tokens_p: float = 0.0,
        min_corrupt_tokens_p: Optional[float] = None,
        apply_corrupt_tokens_p: float = 0.2,
        packed_call: bool = True,
    ):
        super().__init__()
        if len(levels) == 0:
            raise ValueError("levels must be non-empty")

        lv = torch.tensor(levels, dtype=torch.int32)
        basis = torch.cumprod(torch.tensor([1] + levels[:-1], dtype=torch.int32), dim=0)

        self.register_buffer("_levels", lv, persistent=False)
        self.register_buffer("_basis", basis, persistent=False)
        self.dim = len(levels)
        self.codebook_size = int(lv.prod().item())

        codebook = self.indices_to_embedding(torch.arange(self.codebook_size))
        self.register_buffer("implicit_codebook", codebook, persistent=False)

        self.drop_quant_p = drop_quant_p
        self.corrupt_tokens_p = corrupt_tokens_p
        self.min_corrupt_tokens_p = (
            min_corrupt_tokens_p
            if min_corrupt_tokens_p is not None
            else corrupt_tokens_p
        )
        self.apply_corrupt_tokens_p = apply_corrupt_tokens_p
        self.packed_call = packed_call

    def __repr__(self) -> str:
        return (
            "FSQ(\n"
            f"  levels={self._levels.tolist()!r},\n"
            f"  codebook_size={self.codebook_size!r},\n"
            f"  drop_quant_p={self.drop_quant_p!r},\n"
            ")"
        )

    def bound(self, z: Tensor, eps: float = 1e-3) -> Tensor:
        half_l = (self._levels - 1) * (1 + eps) / 2
        offset = torch.where(self._levels % 2 == 0, 0.5, 0.0)
        shift = (offset / half_l).atanh()
        return (z + shift).tanh() * half_l - offset

    def quantize(self, z: Tensor) -> Tensor:
        bounded = self.bound(z)
        drop_p = self.drop_quant_p if self.training else 0.0
        quant = round_ste_with_dropout(bounded, drop_p)
        half_width = (self._levels // 2).to(device=quant.device, dtype=quant.dtype)
        return quant / half_width

    def _scale_and_shift(self, zhat_normalized: Tensor) -> Tensor:
        half_width = (self._levels // 2).to(
            device=zhat_normalized.device, dtype=zhat_normalized.dtype
        )
        return (zhat_normalized * half_width) + half_width

    def _scale_and_shift_inverse(self, zhat: Tensor) -> Tensor:
        half_width = (self._levels // 2).to(device=zhat.device, dtype=zhat.dtype)
        return (zhat - half_width) / half_width

    def codes_to_indices(self, zhat: Tensor) -> Tensor:
        if zhat.shape[-1] != self.dim:
            raise ValueError(f"Expected last dim={self.dim}, got {zhat.shape[-1]}")
        zhat = self._scale_and_shift(zhat)
        basis = self._basis.to(device=zhat.device, dtype=zhat.dtype)
        return (zhat * basis).sum(dim=-1).to(torch.int32)

    def indices_to_embedding(self, indices: Tensor) -> Tensor:
        idx = indices.to(torch.int32).unsqueeze(-1)
        basis = self._basis.to(device=indices.device)
        levels = self._levels.to(device=indices.device)
        codes_non_centered = (idx // basis) % levels
        return self._scale_and_shift_inverse(codes_non_centered.float())

    def corrupt_quant(self, quant: Tensor) -> Tensor:
        shape = quant.shape[:-1]
        random_indices = torch.randint(
            0, self.codebook_size, size=shape, device=quant.device
        )
        random_quant = self.implicit_codebook.to(
            device=quant.device, dtype=quant.dtype
        )[random_indices]
        ratio = random.uniform(self.min_corrupt_tokens_p, self.corrupt_tokens_p)
        mask = (torch.rand(shape, device=quant.device) < ratio).unsqueeze(-1)
        return torch.where(mask, random_quant, quant)

    def _forward_tensor(self, z: Tensor) -> Tuple[Tensor, Tensor]:
        if z.shape[-1] != self.dim:
            raise ValueError(f"Expected latent dim={self.dim}, got {z.shape[-1]}")
        quant = self.quantize(z.float())
        if (
            self.training
            and self.corrupt_tokens_p > 0.0
            and random.random() < self.apply_corrupt_tokens_p
        ):
            quant = self.corrupt_quant(quant)
        tokens = self.codes_to_indices(quant).long()
        return quant, tokens

    def _packed_call(
        self, fn: Callable[[Tensor], Union[Tensor, Tuple[Tensor, ...]]], x
    ):
        if isinstance(x, torch.Tensor):
            return fn(x)

        x_list = x
        x_seqlens = [xi.shape[1] for xi in x_list]
        x_packed = torch.cat(x_list, dim=1)
        results_packed = fn(x_packed)

        results = []
        if isinstance(results_packed, tuple):
            for res_i in results_packed:
                if (
                    isinstance(res_i, torch.Tensor)
                    and res_i.ndim > 1
                    and res_i.shape[1] == sum(x_seqlens)
                ):
                    results.append(list(torch.split(res_i, x_seqlens, dim=1)))
                else:
                    results.append(res_i)
            return tuple(results)
        return list(torch.split(results_packed, x_seqlens, dim=1))

    def forward(
        self, latents: Union[Tensor, List[Tensor]]
    ) -> Tuple[Union[Tensor, List[Tensor]], Union[Tensor, List[Tensor]]]:
        if self.packed_call:
            return self._packed_call(self._forward_tensor, latents)
        if isinstance(latents, list):
            quant: List[Tensor] = []
            tokens: List[Tensor] = []
            for z_i in latents:
                q_i, t_i = self._forward_tensor(z_i)
                quant.append(q_i)
                tokens.append(t_i)
            return quant, tokens

        return self._forward_tensor(latents)
