import torch
import torch.nn as nn

BATCH_SIZE = 4096
DIM = 4096
SHAPE = (BATCH_SIZE, DIM)
SCALE = 1
MODE = 'PN'

class PairNorm(nn.Module):
    '''
    PairNorm
    https://openreview.net/pdf?id=rkecl1rtwB
    '''
    def __init__(self, mode='PN', scale=1):
        """
            mode:
              'None' : No normalization 
              'PN'   : Original version
              'PN-SI'  : Scale-Individually version
              'PN-SCS' : Scale-and-Center-Simultaneously version
        """
        assert mode in ['None', 'PN',  'PN-SI', 'PN-SCS']
        super(PairNorm, self).__init__()
        self.mode = mode
        self.scale = scale
                
    def forward(self, x):
        if self.mode == 'None':
            return x
        
        col_mean = x.mean(dim=0)      
        if self.mode == 'PN':
            x = x - col_mean
            rownorm_mean = (1e-6 + x.pow(2).sum(dim=1).mean()).sqrt() 
            x = self.scale * x / rownorm_mean

        if self.mode == 'PN-SI':
            x = x - col_mean
            rownorm_individual = (1e-6 + x.pow(2).sum(dim=1, keepdim=True)).sqrt()
            x = self.scale * x / rownorm_individual

        if self.mode == 'PN-SCS':
            rownorm_individual = (1e-6 + x.pow(2).sum(dim=1, keepdim=True)).sqrt()
            x = self.scale * x / rownorm_individual - col_mean

        return x

class Model(nn.Module):
    def __init__(self, mode, scale):
        super(Model, self).__init__()
        self.norm = PairNorm(mode=mode, scale=scale)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.norm(x)

def get_inputs():
    x = torch.randn(SHAPE, dtype=torch.float32)
    return [x.contiguous()]

def get_init_inputs():
    return [MODE, SCALE]