import torch
import torch.nn as nn

class quant(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, T):
        ctx.save_for_backward(input)
        ctx.T = T
        return torch.round(torch.clamp(input, min=0, max=T))
        
    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        grad_input[input > ctx.T] = 0
        return grad_input, None



class MultiSpike(torch.nn.Module):
    def __init__(self, dim: int, T=4):
        super().__init__()
        self.T = T
        self.spike = quant()
        self.momentum = 0.1
        self.eps = 1e-5
        self.register_buffer("running_stats", torch.zeros(dim))
    
    def __repr__(self):
        return f"MultiSpike(T={self.T})"    
    
    def forward(self, x, iiter=0):
        #v7
        # print('iiter:{}'.format(iiter))
        if self.training:
            Stats = x.max(dim=0).values.max(dim=0).values
            # Stats = x.abs().mean(dim=[0,1])
            with torch.no_grad():
                self.running_stats = self.momentum * Stats + (1-self.momentum) * self.running_stats
        else:
            Stats = self.running_stats

        scale = self.T / (Stats[None, None, :] + self.eps)

        return self.spike.apply(scale* x, self.T) / scale