from transformers import PretrainedConfig, PreTrainedModel
import metl.models as models
from metl.main import *

class METLConfig(PretrainedConfig):
    IDENT_UUID_MAP = IDENT_UUID_MAP
    UUID_URL_MAP = UUID_URL_MAP
    model_type = "METL"

    def __init__(
            self,
            id:str = None,
            **kwargs,
    ):
        self.id = id
        super().__init__(**kwargs)

class METLModel(PreTrainedModel):
    config_class = METLConfig
    def __init__(self, config:METLConfig):
        super().__init__(config)
        self.model = None
        self.encoder = None
        self.config = config
        
    def forward(self, X, pdb_fn=None):
        if pdb_fn:
            return self.model(X, pdb_fn=pdb_fn)
        return self.model(X)
    
    def load_from_uuid(self, id):
        if id:
            id = id.lower()
            assert id in self.config.UUID_URL_MAP, "ID given does not reference a valid METL model in the IDENT_UUID_MAP"
            self.config.id = id

        self.model, self.encoder = get_from_uuid(self.config.UUID_URL_MAP[self.config.id])

    def load_from_indent(self, id):
        if id:
            id = id.lower()
            assert id in self.config.IDENT_UUID_MAP, "ID given does not reference a valid METL model in the IDENT_UUID_MAP"
            self.config.id = id

        self.model, self.encoder = get_from_ident(self.config.IDENT_UUID_MAP[self.config.id])