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):
        """Restores halogens back to their proper form after decoding."""
        return smiles.replace('L', 'Cl').replace('R', 'Br')
    
    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 = []
        for char in char_list:
            if char.startswith('['):
                tokenized.append(char)
            else:
                tokenized.extend(list(char))
        return tokenized

    def encode(self, char_list):
        smiles_matrix = np.zeros(len(char_list), dtype=np.float32)
        for i, char in enumerate(char_list):
            if char not in self.voc.vocab:
                raise ValueError(f"Token '{char}' not in vocabulary.")
            smiles_matrix[i] = self.voc.vocab[char]
        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)  # Restore halogens after decoding
        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 topk_sample(self, probs, k=5):
        topk_probs, topk_indices = torch.topk(probs, k)
        topk_probs = topk_probs / topk_probs.sum(dim=-1, keepdim=True)
        sampled = torch.multinomial(topk_probs, 1)
        return topk_indices.gather(-1, sampled)

    def unmask_partial_smiles(self, input_ids, model, voc, steps=30, k=1):
        model.eval()
        with torch.no_grad():
            sequences = torch.tensor(input_ids, dtype=torch.long, device=self.device).unsqueeze(0)

            total_tokens = sequences.size(1)
            original_num_masked = (sequences == voc.vocab['[MASK]']).sum().item()
            print(f"Original number of masked tokens: {original_num_masked}")

            if total_tokens == 0:
                raise ValueError("Encoded SMILES has no tokens.")
            if original_num_masked == 0:
                print("No masked tokens found. Returning original SMILES.")
                return self.decode(input_ids)

            filled_tokens_confidence = {}

            for step in range(steps):
                t = torch.tensor([(sequences == voc.vocab['[MASK]']).sum().item() / total_tokens], device=self.device)
                logits = model(sequences, t=t)
                probs = F.softmax(logits, dim=-1)

                mask_indices = (sequences == voc.vocab['[MASK]']).nonzero(as_tuple=True)[1]
                if len(mask_indices) == 0:
                    print("All tokens filled.")
                    break

                confidences = []
                filled_indices = []

                for idx in mask_indices:
                    token_probs = probs[0, idx]
                    topk_token = self.topk_sample(token_probs, k)
                    confidence = token_probs[topk_token].item()
                    sequences[0, idx] = topk_token
                    confidences.append(confidence)
                    filled_indices.append(idx)
                    filled_tokens_confidence[idx] = confidence

                avg_conf = sum(confidences) / len(confidences)

                for idx, conf in zip(filled_indices, confidences):
                    if conf < avg_conf:
                        sequences[0, idx] = voc.vocab['<MASK>']
                        filled_tokens_confidence[idx] = None

            sorted_filled_tokens = sorted(
                ((idx, conf) for idx, conf in filled_tokens_confidence.items() if conf is not None),
                key=lambda x: x[0]
            )

            for idx, confidence in sorted_filled_tokens:
                print(f"Final confidence for token at index {idx}: {confidence:.3f}")

            decoded = self.decode([
                token for token in sequences[0].tolist()
                if token not in (voc.vocab.get('<s>'), voc.vocab.get('</s>'))
            ]).replace(" ", "")
            return decoded