import torch
import torch.nn as nn

class Trigger(nn.Module):

    def __init__(self, size: int = 32, epsilon: float=16/255, transparency: float = 1.) -> None:
        super().__init__()

        self.size = size
        self.epsilon = epsilon
        self.mask = nn.Parameter(torch.rand(size, size,device=torch.device('cuda')),requires_grad=True)
        self.transparency = transparency
        self.trigger = nn.Parameter(torch.rand(3, size, size,device=torch.device('cuda')) * 4 - 2,requires_grad=True)
        # self.trigger = nn.Parameter(torch.rand(3, size, size,device=torch.device('cuda')),requires_grad=True)


    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x = torch.min(torch.max(tri(x), x - epsilon), x + epsilon)
        if len(x.shape) == 4:
            return self.transparency * self.mask * self.trigger.repeat(len(x), 1, 1, 1) + (1 - self.mask * self.transparency) * x
            # return torch.clamp(self.transparency * torch.clamp(self.mask, 0, 1) * torch.clamp(self.trigger.repeat(len(x), 1, 1, 1), -self.epsilon,
            #                                                                               self.epsilon) + (
            #                            1 - torch.clamp(self.mask, 0, 1) * self.transparency) * x, -1, 1)
        else:
            return self.transparency * self.mask * self.trigger + (1 - self.mask * self.transparency) * x
            # return torch.clamp(self.transparency * torch.clamp(self.mask, 0, 1) * torch.clamp(self.trigger, -self.epsilon, self.epsilon) + (1 - torch.clamp(self.mask, 0, 1) * self.transparency) * x, -1, 1)
    
class UAP(nn.Module):

    def __init__(self, size: int = 32) -> None:
        super().__init__()

        self.size = size
        self.perturbation = nn.Parameter(torch.zeros(3, size, size,device=torch.device('cuda')),requires_grad=True)


    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.perturbation