import torch
import re 
from transformers import T5Tokenizer, T5EncoderModel
from transformers import AutoTokenizer, AutoModel


class EmbeddingExtractor:
    def __init__(self, device):
        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')

        self.prot_tokenizer = T5Tokenizer.from_pretrained('Rostlab/prot_t5_xl_half_uniref50-enc', do_lower_case=False)
        self.prot_model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_half_uniref50-enc").to(device)
        
        self.mol_tokenizer = AutoTokenizer.from_pretrained("ibm/MoLFormer-XL-both-10pct", deterministic_eval=True, trust_remote_code=True)
        self.mol_model = AutoModel.from_pretrained("ibm/MoLFormer-XL-both-10pct", trust_remote_code=True).to(self.device)

    def get_protein_embedding(self, sequence):
        tokens = self.prot_tokenizer(sequence, return_tensors='pt', padding=True, truncation=True, max_length=3200).to(self.device)
        with torch.no_grad():
            embedding = self.prot_model(**tokens).last_hidden_state.mean(dim=1)
        return embedding
    
    def get_molecule_embedding(self, smiles):
        tokens = self.mol_tokenizer(smiles, return_tensors='pt', padding=True, truncation=True, max_length=278).to(self.device)
        with torch.no_grad():
            embedding = self.mol_model(**tokens).last_hidden_state.mean(dim=1)
        return embedding

    def get_combined_embedding(self, sequence, smiles):
        prot_embedding = self.get_protein_embedding(sequence)
        mol_embedding = self.get_molecule_embedding(smiles)
        return prot_embedding, mol_embedding

def preprocess_function(df):
    df['seq'] = df['seq'].apply(lambda x: " ".join(re.sub(r"[UZOB]", "X", x)))
    return df