import torch
import torch.nn as nn

BATCH_SIZE = 32
CHANNELS = 64
HEIGHT = 56
WIDTH = 56
SHAPE = (BATCH_SIZE, CHANNELS, HEIGHT, WIDTH)
BETA_SHAPE = (BATCH_SIZE, CHANNELS, 1, 1)

DTYPE = torch.float64

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

    def forward(self, x, beta):
        # p1, p2: (1, C, 1, 1) broadcast to (N, C, H, W)
        # beta:   (N, C, 1, 1) broadcast to (N, C, H, W)
        # x:      (N, C, H, W)
        d1 = self.p1 * x
        d2 = self.p2 * x
        diff = d1 - d2
        return diff * torch.sigmoid(beta * diff) + d2

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

def get_inputs():
    x = torch.randn(SHAPE, dtype=DTYPE)
    # In a real model, this comes from Conv/BN layers
    beta = torch.randn(BETA_SHAPE, dtype=DTYPE)
    return [x.contiguous(), beta.contiguous()]

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