import torch
import torch.nn as nn
import math

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

ALPHA_INIT = 1.0
BETA_INIT = 1.0

DTYPE = torch.float64

class PGELU(nn.Module):
    '''
    P-GELU: A Novel Activation Function to Optimize Whisper for Darija Speech Translation
    https://ieeexplore.ieee.org/document/11016691
    Formula: f(x) = x * (1 + tan(alpha * x + beta * x^3))
    '''
    def __init__(self, alpha_init=1.0, beta_init=1.0):
        super(PGELU, self).__init__()
        self.alpha = nn.Parameter(torch.tensor(alpha_init, dtype=DTYPE))
        self.beta = nn.Parameter(torch.tensor(beta_init, dtype=DTYPE))
        self.pi_half = math.pi / 2.0
        self.eps = 1e-7 # Epsilon for double

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        inner = self.alpha * x + self.beta * x.pow(3)
        inner_clamped = torch.clamp(inner, -self.pi_half + self.eps, self.pi_half - self.eps)
        return x * (1.0 + torch.tan(inner_clamped))

class Model(nn.Module):
    def __init__(self, alpha_init=1.0, beta_init=1.0):
        super(Model, self).__init__()
        self.act = PGELU(alpha_init, beta_init)
    
    def forward(self, x):
        return self.act(x)

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

def get_init_inputs():
    return [ALPHA_INIT, BETA_INIT]