import torch
import torch.nn as nn

BATCH_SIZE = 64
CHANNELS = 128
HEIGHT = 56
WIDTH = 56
SHAPE = (BATCH_SIZE, CHANNELS, HEIGHT, WIDTH)

DTYPE = torch.float64

class AconC(nn.Module):
    """ ACON activation (activate or not).
    # AconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is a learnable parameter
    # according to "Activate or Not: Learning Customized Activation" <https://arxiv.org/pdf/2009.04759.pdf>.
    """
    def __init__(self, p1, p2, beta):
        super().__init__()
        self.p1 = nn.Parameter(p1)
        self.p2 = nn.Parameter(p2)
        self.beta = nn.Parameter(beta)

    def forward(self, x):
        return (self.p1 * x - self.p2 * x) * torch.sigmoid(self.beta * (self.p1 * x - self.p2 * x)) + self.p2 * x

class Model(nn.Module):
    def __init__(self, p1, p2, beta):
        super(Model, self).__init__()
        self.act = AconC(p1, p2, beta)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.act(x)

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

def get_init_inputs():
    p1 = torch.randn(1, CHANNELS, 1, 1, dtype=DTYPE)
    p2 = torch.randn(1, CHANNELS, 1, 1, dtype=DTYPE)
    beta = torch.ones(1, CHANNELS, 1, 1, dtype=DTYPE)
    return [p1, p2, beta]