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

# --- Hyperparameters ---
N, C, H, W = 16, 16, 64, 64
BETA = 1.0
THRESHOLD = 20.0


class Softplus(nn.Module):

    def __init__(self, beta=BETA, threshold=THRESHOLD):
        super().__init__()
        self.beta = beta
        self.threshold = threshold

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        scaled_input = input * self.beta

        mask = (scaled_input > self.threshold)

        stable_output = (1.0 / self.beta) * torch.log1p(torch.exp(scaled_input))

        linear_output = input

        return torch.where(mask, linear_output, stable_output)


class Model(nn.Module):
    def __init__(self, beta=BETA, threshold=THRESHOLD):
        super().__init__()
        self.op = Softplus(beta=beta, threshold=threshold)

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return self.op(input)


def get_inputs():
    torch.manual_seed(42)

    x = torch.randn(N, C, H, W, dtype=torch.float32) * (THRESHOLD / BETA) / 5.0

    x[0, 0, 0, 0] = THRESHOLD / BETA + 1.0
    return [x]


def get_init_inputs():
    return [BETA, THRESHOLD]