from transformers import Pipeline
import torch
import torch.nn.functional as F
import numpy as np
from transformers import Pipeline
import torch
import torch.nn.functional as F
import re

class SmilesDiffusionPipe(Pipeline):
    def __init__(self, model, voc, **kwargs):
        super().__init__(model=model, **kwargs)
        self.voc = voc

    def restore_halogen(self, smiles):
        smiles = re.sub(r'(?<!\[)L(?!\])', 'Cl', smiles)
        smiles = re.sub(r'(?<!\[)R(?!\])', 'Br', smiles)
        return smiles
    
    def replace_halogen(self, smiles):
        return smiles.replace('Cl', 'L').replace('Br', 'R')

    def tokenize_smiles(self, smiles):
        regex = r'(\[[^\[\]]{1,6}\])'
        smiles = self.replace_halogen(smiles)
        char_list = re.split(regex, smiles)
        tokenized = ['[GO]']
        for char in char_list:
            if char.startswith('['):
                tokenized.append(char)
            else:
                chars = [unit for unit in char]
                [tokenized.append(unit) for unit in chars]
        tokenized.append("[EOS]")
        return tokenized

    def encode(self, char_list):
        smiles_matrix = np.zeros(len(char_list), dtype=np.float32)
    
        unk_id = self.voc.vocab.get("[UNK]")
        if unk_id is None:
            raise ValueError("[UNK] token not found in vocabulary.")
    
        for i, char in enumerate(char_list):
            smiles_matrix[i] = self.voc.vocab.get(char, unk_id)
    
        return smiles_matrix
        
    def decode(self, matrix):
        """"Takes an array of indices (can be list of tensors or ints) and returns the corresponding SMILES"""
        matrix = [int(i) for i in matrix] # This handles tensors from PyTorch
        # Using tokenizer's decode method to convert token IDs to text
        smiles = self.voc.decode(matrix)
        smiles = self.restore_halogen(smiles)
        return smiles

    def _sanitize_parameters(self, **kwargs):
        preprocess_kwargs = {}
        postprocess_kwargs = {}
        if "steps" in kwargs:
            preprocess_kwargs["steps"] = kwargs["steps"]
        if "k" in kwargs:
            preprocess_kwargs["k"] = kwargs["k"]
        return preprocess_kwargs, {}, postprocess_kwargs

    def preprocess(self, smiles_str, steps=30, k=1):
        tokenized = self.tokenize_smiles(smiles_str)
        encoded = self.encode(tokenized)
        return (encoded, steps, k)

    def _forward(self, inputs):
        encoded, steps, k = inputs
        return self.unmask_partial_smiles(encoded, self.model, self.voc, steps, k)

    def postprocess(self, result):
        return result

    def multinomial_sample(self, probs):
        return torch.multinomial(probs, 1)


    def unmask_smiles_argmax(self, input_ids, model, voc):
        """
        Unmasks all [MASK] tokens in a single forward pass using argmax decoding.
        """
        model.eval()
        with torch.no_grad():
            sequences = torch.tensor(
                input_ids, dtype=torch.long, device=self.device
            ).unsqueeze(0)
    
            mask_token = voc.vocab["[MASK]"]
            pad_token = voc.vocab.get("[PAD]", None)
    
            # diffusion time = fraction masked
            if pad_token is not None:
                valid = sequences != pad_token
                frac_masked = (sequences == mask_token).sum() / valid.sum()
            else:
                frac_masked = (sequences == mask_token).sum() / sequences.size(1)
    
            t = frac_masked.unsqueeze(0).to(self.device)
    
            # forward
            logits = model(sequences, t=t)
            probs = F.softmax(logits, dim=-1)
    
            # argmax over vocab
            pred_ids = torch.argmax(probs, dim=-1)
    
            # replace all masks at once
            mask_positions = sequences == mask_token
            sequences[mask_positions] = pred_ids[mask_positions]
    
            # decode
            decoded = [
                t for t in sequences[0].tolist()
                if t not in (pad_token, voc.vocab.get("<s>"), voc.vocab.get("</s>"))
            ]
    
            if len(decoded) > 2:
                decoded = decoded[1:-1]
    
            return self.decode(decoded).replace(" ", "")