# Load model directly
import torch.nn as nn
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM, EsmConfig

class ESMEncoder(nn.Module):
    def __init__(self, protein_tower, 
                 protein_max_len = 1024, 
                 delay_load=False,
                 select_layer = -1,
                 use_smiles = False,
                 smiles_type = None) -> None:
        super().__init__()
        self.protein_tower_name = protein_tower
        self.protein_max_len = protein_max_len
        self.select_layer = select_layer

        # if use_smiles:
        #     self.smiles_tokenizer = AutoTokenizer.from_pretrained(smiles_type)
        #     self.smiles_model = AutoModelForMaskedLM.from_pretrained(smiles_type)

        if not delay_load:
            self.load_model()
        else:
            self.cfg_only = EsmConfig.from_pretrained(self.protein_tower_name)

    def load_model(self):
        self.protein_tower_tokenizer = AutoTokenizer.from_pretrained(self.protein_tower_name)
        self.protein_tower = AutoModelForMaskedLM.from_pretrained(self.protein_tower_name)
        self.protein_tower.requires_grad_(False)
        # print(self.protein_tower.config)
        self.is_loaded = True

    def feature_select(self, protein_forward_outs):
        # image_features = image_forward_outs.hidden_states[self.select_layer]
        # if self.select_feature == 'patch':
        #     image_features = image_features[:, 1:]
        # elif self.select_feature == 'cls_patch':
        #     image_features = image_features
        # else:
        #     raise ValueError(f'Unexpected select feature: {self.select_feature}')
        return protein_forward_outs.hidden_states[self.select_layer]

    @torch.no_grad()
    def forward(self, proteins, SMILES=None):
        # if type(proteins) is list:
        #     protein_features = []
        #     for protein in proteins:
        #         tokens = self.protein_tower_tokenizer(protein, return_tensors="pt", max_length=self.protein_max_len)
        #         protein_forward_out = self.vision_tower(tokens.to(device=self.device, dtype=self.dtype), output_hidden_states=True)
        #         protein_feature = self.feature_select(protein_forward_out).to(protein.dtype)
        #         protein_features.append(protein_feature)
        # else:clear
        tokens = self.protein_tower_tokenizer(proteins, return_tensors="pt", max_length=self.protein_max_len, padding=True, truncation=True)
        # print(tokens)
        tokens = tokens.to(device=self.device)
        protein_forward_outs = self.protein_tower(**tokens, output_hidden_states=True)
        protein_features = self.feature_select(protein_forward_outs)
        protein_features = protein_features.to(self.dtype)
        # print(protein_features.shape)
        return protein_features

    @property
    def dummy_feature(self):
        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)

    @property
    def dtype(self):
        return self.protein_tower.dtype

    @property
    def device(self):
        return self.protein_tower.device

    @property
    def config(self):
        if self.is_loaded:
            return self.protein_tower.config
        else:
            return self.cfg_only

    @property
    def hidden_size(self):
        return self.config.hidden_size

    @property
    def num_patches(self):
        return self.config.max_position_embeddings