import torch


class Masker:
    def __init__(self, tokenizer) -> None:
        self.tokenizer = tokenizer
        self.mask_token_id = self.tokenizer.mask_token_id

    def random_mask(self, input_ids, mask_prob=0.15):
        device = input_ids.device
        mask = (torch.rand(input_ids.shape) < mask_prob).to(device)
        mask = mask & (torch.logical_not(self.tokenizer.is_special_token(input_ids)))
        masked_input_ids = input_ids.clone()
        masked_input_ids[mask] = self.mask_token_id
        return masked_input_ids, mask

    def mask_ptm_tokens(
        self,
        input_ids,
    ):
        device = input_ids.device
        is_ptm_mask = self.tokenizer.is_ptm_token(input_ids).to(device)
        is_ptm_mask = is_ptm_mask & (
            torch.logical_not(self.tokenizer.is_special_token(input_ids))
        )
        masked_input_ids = input_ids.clone()
        masked_input_ids[is_ptm_mask] = self.mask_token_id
        return masked_input_ids, is_ptm_mask

    def random_and_ptm_mask(self, input_ids, mask_prob=0.15):
        device = input_ids.device
        mask = (torch.rand(input_ids.shape) < mask_prob).to(device)
        mask = mask & (torch.logical_not(self.tokenizer.is_special_token(input_ids)))
        is_ptm_mask = self.tokenizer.is_ptm_token(input_ids).to(device)
        is_ptm_mask = is_ptm_mask & (
            torch.logical_not(self.tokenizer.is_special_token(input_ids))
        )
        mask = mask | is_ptm_mask
        masked_input_ids = input_ids.clone()
        masked_input_ids[mask] = self.mask_token_id
        return masked_input_ids, mask

    def random_or_random_and_ptm_mask(
        self, input_ids, ranom_mask_prob=0.15, alternate_prob=0.2
    ):
        """
        alternate between [(1) random mask] and [(2) random mask & ptm mask] by probability alternate_prob
        """
        p = torch.rand(1).item()
        if p < alternate_prob:
            return self.random_mask(input_ids, ranom_mask_prob)
        else:
            return self.random_and_ptm_mask(input_ids, ranom_mask_prob)