import torch
import torch.nn as nn
import torch.nn.functional as F

BATCH_SIZE = 4096
HIDDEN_DIM = 4096
SHAPE = (BATCH_SIZE, HIDDEN_DIM)

GAMMA_VAL = 1.25
DELTA_VAL = 0.85

class Pserf(nn.Module):
    '''
    "ErfAct and Pserf: Non-monotonic Smooth Trainable Activation Functions" (AAAI 2022)
    https://doi.org/10.1609/aaai.v36i6.20557
    Formula: f(x) = x * erf(gamma * log(1 + exp(delta * x)))
    '''
    def __init__(self, gamma=1.25, delta=0.85):
        super(Pserf, self).__init__()
        self.gamma = gamma
        self.delta = delta

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        sp = F.softplus(self.delta * x)
        erf_val = torch.erf(self.gamma * sp)
        return x * erf_val

class Model(nn.Module):
    def __init__(self, gamma=1.25, delta=0.85):
        super(Model, self).__init__()
        self.act = Pserf(gamma, delta)
    
    def forward(self, x):
        return self.act(x)

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

def get_init_inputs():
    return [GAMMA_VAL, DELTA_VAL]