import torch
import torch.nn as nn
from torch.autograd import Function

class VectorQuantization(Function):
    @staticmethod
    def forward(ctx, inputs, codebook):
        with torch.no_grad():
            embedding_size = codebook.size(1)
            inputs_size = inputs.size()
            inputs_flatten = inputs.view(-1, embedding_size)

            codebook_sqr = torch.sum(codebook ** 2, dim=1)
            inputs_sqr = torch.sum(inputs_flatten ** 2, dim=1, keepdim=True)

            # Compute the distances to the codebook
            distances = torch.addmm(codebook_sqr + inputs_sqr,
                inputs_flatten, codebook.t(), alpha=-2.0, beta=1.0)

            _, indices_flatten = torch.min(distances, dim=1)
            indices = indices_flatten.view(*inputs_size[:-1])
            ctx.mark_non_differentiable(indices)

            return indices

    @staticmethod
    def backward(ctx, grad_output):
        raise RuntimeError('Trying to call `.grad()` on graph containing '
            '`VectorQuantization`. The function `VectorQuantization` '
            'is not differentiable. Use `VectorQuantizationStraightThrough` '
            'if you want a straight-through estimator of the gradient.')

class VectorQuantizationStraightThrough(Function):
    @staticmethod
    def forward(ctx, inputs, codebook):
        indices = vq(inputs, codebook)
        indices_flatten = indices.view(-1)
        ctx.save_for_backward(indices_flatten, codebook)
        ctx.mark_non_differentiable(indices_flatten)

        codes_flatten = torch.index_select(codebook, dim=0,
            index=indices_flatten)
        codes = codes_flatten.view_as(inputs)

        return (codes, indices_flatten)

    @staticmethod
    def backward(ctx, grad_output, grad_indices):
        grad_inputs, grad_codebook = None, None

        if ctx.needs_input_grad[0]:
            # Straight-through estimator
            grad_inputs = grad_output.clone()
        if ctx.needs_input_grad[1]:
            # Gradient wrt. the codebook
            indices, codebook = ctx.saved_tensors
            embedding_size = codebook.size(1)

            grad_output_flatten = (grad_output.contiguous()
                                              .view(-1, embedding_size))
            grad_codebook = torch.zeros_like(codebook)
            grad_codebook.index_add_(0, indices, grad_output_flatten)

        return (grad_inputs, grad_codebook)

class MaskStraightThrough(Function):
    @staticmethod
    def forward(ctx, m):
        m_ = (m > 0.5).float()  # 二值化为硬掩码
        ctx.save_for_backward(m, m_)
        return m_

    @staticmethod
    def backward(ctx, grad_output):
        m, m_ = ctx.saved_tensors
        grad_input = grad_output.clone()
        # 梯度裁剪
        max_grad_norm = 1.0  # 设置梯度裁剪的最大范数
        grad_norm = torch.norm(grad_input)
        if grad_norm > max_grad_norm:
            grad_input = grad_input * (max_grad_norm / grad_norm)
        # 确保 m 在 [0, 1] 范围内，以避免数值问题
        m = torch.clamp(m, 0, 1)
        return grad_input * m * (1 - m)

class VectorQuantization_Ts(Function):
    @staticmethod
    def forward(ctx, inputs, quantizer):
        with torch.no_grad():
            quantized, indices = quantizer(inputs)
            indices_flatten = indices.view(-1)
            ctx.mark_non_differentiable(indices_flatten)
            ctx.save_for_backward(indices_flatten)
            return indices

    @staticmethod
    def backward(ctx, grad_output):
        raise RuntimeError('Trying to call `.grad()` on graph containing '
            '`VectorQuantization`. The function `VectorQuantization` '
            'is not differentiable. Use `VectorQuantizationStraightThrough` '
            'if you want a straight-through estimator of the gradient.')

class VectorQuantizationStraightThrough_Ts(Function):
    @staticmethod
    def forward(ctx, inputs, quantizer):
        quantized, indices = quantizer(inputs)
        indices_flatten = indices.view(-1)
        ctx.save_for_backward(indices_flatten, quantized)
        ctx.mark_non_differentiable(indices_flatten)
        return (quantized, indices_flatten)

    @staticmethod
    def backward(ctx, grad_output, grad_indices):
        grad_inputs, grad_codebook = None, None

        if ctx.needs_input_grad[0]:
            # Straight-through estimator
            grad_inputs = grad_output.clone()
        if ctx.needs_input_grad[1]:
            # Gradient wrt. the codebook
            indices, codebook = ctx.saved_tensors
            embedding_size = codebook.size(2)  # Update based on the new quantizer

            grad_output_flatten = (grad_output.contiguous()
                                              .view(-1, embedding_size))
            grad_codebook = torch.zeros_like(codebook)
            grad_codebook.index_add_(0, indices, grad_output_flatten)

        return (grad_inputs, grad_codebook)


ms = MaskStraightThrough.apply
vq = VectorQuantization.apply
vq_st = VectorQuantizationStraightThrough.apply
vq_t = VectorQuantization_Ts.apply
vq_stt = VectorQuantizationStraightThrough_Ts.apply
__all__ = [vq, vq_st]
