import os
from contextlib import nullcontext

os.environ['TOKENIZERS_PARALLELISM'] = 'False'

import numpy as np
from packaging.version import parse as V
import torch
import transformers
from transformers import AutoModel, AutoTokenizer, T5Tokenizer, T5EncoderModel

from typing import *

from .engine import RepEngineBase

transformers.logging.set_verbosity(transformers.logging.ERROR)


AVAILABLE_MODELS = {
    'esm2_t48_15B_UR50D': 5120,
    'esm2_t36_3B_UR50D': 2560,
    'esm2_t33_650M_UR50D': 1280,
    'esm1b_t33_650M_UR50S': 1280,
    'esm2_t30_150M_UR50D': 640,
    'esm2_t12_35M_UR50D': 480,
    'esm2_t6_8M_UR50D': 320,
    'ESMplusplus_small': 960,
    'ESMplusplus_large': 1152,
    'prot_t5_xxl_uniref50': 1024,
    'prot_t5_xl_half_uniref50-enc': 1024,
    'prot_bert': 1024,
    'ProstT5': 1024,
    'ankh-base': 768,
    'ankh-large': 1536,
    'MoLFormer-XL-both-10pct': 768,
    'ChemBERTa-77M-MLM': 384,
    'ChemBERTa-100M-MLM': 768,
    'PeptideCLM-23M-all': 768,
    'PeptideMTR_lg': 1024,
}

SYNONYMS = {
    'prot-t5-xl': 'prot_t5_xl_half_uniref50-enc',
    'prot-t5-xxl': 'prot_t5_xxl_uniref50',
    'protbert': 'prot_bert',
    'prost-t5': 'ProstT5',
    'esm2-15b': 'esm2_t48_15B_UR50D',
    'esm2-3b': 'esm2_t36_3B_UR50D',
    'esm2-650m': 'esm2_t33_650M_UR50D',
    'esm1b': 'esm1b_t33_650M_UR50S',
    'esm2-150m': 'esm2_t30_150M_UR50D',
    'esm2-35m': 'esm2_t12_35M_UR50D',
    'esm2-8m': 'esm2_t6_8M_UR50D',
    'esmc-300m': 'ESMplusplus_small',
    'esmc-600m': 'ESMplusplus_large',
    'ankh-base': 'ankh-base',
    'ankh-large': 'ankh-large',
    'molformer-xl': 'MoLFormer-XL-both-10pct',
    'chemberta-2': 'ChemBERTa-77M-MLM',
    'chemberta-3': 'ChemBERTa-100M-MLM',
    'peptideclm': 'PeptideCLM-23M-all',
    'peptidemtr': "PeptideMTR_lg",
}


class RepEngineLM(RepEngineBase):
    """
    Class `RepEngineLM` is a subclass of `RepEngineBase` designed to compute molecular representations 
    using pre-trained language models (LMs) such as T5, ESM, or ChemBERTa. This engine generates vector-based 
    embeddings for input sequences, typically protein or peptide sequences, by leveraging transformer-based models.

    Attributes:
        :type engine: str
        :param engine: The name of the engine. Default is `'lm'`, indicating a language model-based representation.

        :type device: str
        :param device: The device on which the model runs, either `'cuda'` for GPU or `'cpu'`.

        :type model: object
        :param model: The pre-trained model used for generating representations. The model is loaded from a repository 
                        based on the `model` parameter.

        :type name: str
        :param name: The name of the model engine combined with the model type.

        :type dimension: int
        :param dimension: The dimensionality of the output representation, corresponding to the model's embedding size.

        :type model_name: str
        :param model_name: The specific model name used for generating representations.

        :type tokenizer: object
        :param tokenizer: The tokenizer associated with the model, used for converting sequences into tokenized input.

        :type lab: str
        :param lab: The laboratory or organization associated with the model (e.g., 'Rostlab', 'facebook', etc.).
    """
    engine = 'lm'

    def __init__(self, model: str, average_pooling: Optional[bool] = True,
                 cls_token: Optional[bool] = False, fp16: bool = True):
        """
        Initializes the `RepEngineLM` with the specified model and pooling options. The model is loaded based on 
        the given `model` name and its associated tokenizer.

        :type model: str
          :param model: The pre-trained model to use for generating representations (e.g., 'esm2_t48_15B_UR50D').

        :type average_pooling: Optional[bool]
          :param average_pooling: If `True`, the embeddings are averaged across all tokens. Default is `True`.

        :type cls_token: Optional[bool]
          :param cls_token: If `True`, only the representation of the [CLS] token is used. Default is `False`.

        :rtype: None
        """
        super().__init__(model, average_pooling=average_pooling,
                         cls_token=cls_token)
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.model = None
        self.name = f'{self.engine}-{model}'
        self.fp16 = fp16
        self._load_model(model)

    def move_to_device(self, device: str):
        """
        Moves the model to the specified device (e.g., 'cuda' or 'cpu').

        :type device: str
          :param device: The target device to move the model to.

        :rtype: None
        """
        self.device = device
        self.model.to(self.device)

    def dim(self) -> int:
        """
        Returns the dimensionality of the output representation generated by the model.

        :rtype: int
          :return: The dimensionality (embedding size) of the model's output.
        """
        return self.dimension

    def max_len(self) -> int:
        """
        Returns the maximum allowed sequence length for the model. Some models have a specific maximum sequence length.

        :rtype: int
          :return: The maximum sequence length for the model.
        """
        if self.lab == 'facebook':
            return 1022
        elif self.lab.lower() == 'evolutionaryscale':
            return 2046
        elif self.lab.lower() == 'deepchem':
            return 512
        else:
            return 2046

    def get_num_params(self, human_readable: bool = False) -> int:
        """
        Returns the total number of parameters in the model.

        :rtype: int
          :return: The number of parameters in the model.
        """
        if not human_readable:
            return sum(p.numel() for p in self.model.parameters())
        else:
            return f"{sum(p.numel() for p in self.model.parameters())/1e6:,.3f}M"

    def _load_model(self, model: str):
        """
        Loads the specified pre-trained model and tokenizer based on the provided model name. 
        The model is selected from the available models in the `AVAILABLE_MODELS` dictionary.

        :type model: str
          :param model: The model name or synonym to load (e.g., 'esm2_t48_15B_UR50D').

        :raises NotImplementedError: If the specified model is not found in `AVAILABLE_MODELS` or `SYNONYMS`.
        :rtype: None
        """
        if model not in AVAILABLE_MODELS and SYNONYMS[model.lower()] not in AVAILABLE_MODELS:
            raise NotImplementedError(
                f"Model: {model} not implemented.",
                f"Available models: {', '.join(AVAILABLE_MODELS)}"
            )
        if model not in AVAILABLE_MODELS:
            model = SYNONYMS[model.lower()]
        if model.lower().startswith('pro'):
            self.lab = 'Rostlab'
        elif 'plusplus' in model.lower():
            self.lab = 'Synthyra'
        elif 'esmc' in model.lower():
            self.lab = 'EvolutionaryScale'
        elif 'esm' in model.lower():
            self.lab = 'facebook'
        elif 'lobster' in model.lower():
            self.lab = 'asalam91'
        elif 'ankh' in model.lower():
            self.lab = 'ElnaggarLab'
        elif 'molformer' in model.lower():
            if V(transformers.__version__) >= V('5.0.0'):
                print("Warning: Molformer does not support transformers>=5.0.0, we recommend using an earlier version, e.g., transformers==4.41.2")
            self.lab = 'ibm-research'
        elif 'chemberta' in model.lower():
            self.lab = 'DeepChem'
        elif (('clm' in model.lower() or 'mtr' in model.lower()) and
              'peptide' in model.lower()):
            self.lab = 'aaronfeller'
        if 't5' in model.lower():
            self.tokenizer = T5Tokenizer.from_pretrained(f'Rostlab/{model}',
                                                         do_lower_case=False)
            self.model = T5EncoderModel.from_pretrained(f"Rostlab/{model}")
        elif 'feller' in self.lab.lower():
            import os
            import urllib
            import urllib.request as request
            try:
                from .utils.peptideclm_tokenizer import SMILES_SPE_Tokenizer
            except ImportError:
                raise ImportError("This function requires smilespe. Please install: `pip install smilespe`")
            if os.getenv('HF_HOME') is None:
                hf_home = os.path.abspath(f"{os.getenv('HOME')}/.cache/huggingface/hub/")
            else:
                hf_home = os.path.abspath(os.getenv('HF_HOME'))
            path = os.path.join(hf_home, 'peptideclm_tokenizer')
            vocab = os.path.join(path, 'new_vocab.txt')
            splits = os.path.join(path, 'new_splits.txt')

            if not (os.path.exists(vocab) and os.path.exists(splits)):
                os.makedirs(path, exist_ok=True)
                try:
                    url1 = 'https://raw.githubusercontent.com/AaronFeller/PeptideCLM/refs/heads/master/tokenizer/new_vocab.txt'
                    url2 = 'https://raw.githubusercontent.com/AaronFeller/PeptideCLM/refs/heads/master/tokenizer/new_splits.txt'
                    request.urlretrieve(url1, vocab)
                    request.urlretrieve(url2, splits)
                except urllib.error.URLError:
                    raise RuntimeError("Tokenizer could not be downloaded. Please try again later and if the problem persists,",
                                       "raise an issue in on the AutoPeptideML github so that the issue can be",
                                       "investigated: https://github.com/IBM/AutoPeptideML/issues")
            self.tokenizer = SMILES_SPE_Tokenizer(vocab_file=vocab,
                                                  spe_file=splits)
            self.model = AutoModel.from_pretrained(f'{self.lab}/{model}',
                                                   trust_remote_code=True)
        else:
            self.model = AutoModel.from_pretrained(f'{self.lab}/{model}',
                                                   trust_remote_code=True)
            if 'plusplus' in model.lower():
                self.tokenizer = self.model.tokenizer
            else:
                self.tokenizer = AutoTokenizer.from_pretrained(
                    f'{self.lab}/{model}', trust_remote_code=True,
                )

        self.dimension = AVAILABLE_MODELS[model]
        self.model_name = model
        self.model.to(self.device)

    def _preprocess_batch(self, sequences: List[str]) -> List[List[str]]:
        """
        Preprocesses a batch of input sequences by adjusting formatting, truncating, and applying special tokens 
        based on the model type.

        :type sequences: List[str]
          :param sequences: A list of input sequences (e.g., protein sequences in FASTA format).

        :rtype: List[List[str]]
          :return: A list of preprocessed sequences.
        """
        if self.lab == 'Rostlab':
            sequences = [' '.join([char for char in seq]) for seq in sequences]
        if self.model_name == 'ProstT5':
            sequences = ["<AA2fold> " + seq for seq in sequences]
        sequences = [seq[:self.max_len()] for seq in sequences]
        return sequences

    def _rep_batch(
        self, batch: List[str],
    ) -> List[np.ndarray]:
        """
        Generates representations for a batch of sequences using the loaded pre-trained model. The representations 
        are extracted from the model's output and returned based on the specified pooling strategy.

        :type batch: List[str]
          :param batch: A list of input sequences (e.g., protein sequences in FASTA format).

        :rtype: List[np.ndarray]
          :return: A list of numpy arrays representing the embeddings of each input sequence.
        """
        batch = [b if len(b) <= self.max_len() else b[:self.max_len()]
                 for b in batch]
        inputs = self.tokenizer(batch, add_special_tokens=True,
                                truncation=True,
                                padding="longest", return_tensors="pt")
        inputs = inputs.to(self.device)
        mps_autocast = int(torch.__version__.split('.')[1]) >= 6
        autocast = self.fp16 and (self.device == 'cuda' or
                                  (self.device == 'mps' and mps_autocast) or
                                  self.device == 'cpu')
        if autocast:
            autocast = torch.autocast(
                device_type=self.device,
                dtype=torch.bfloat16
            )
        else:
            autocast = nullcontext()

        with torch.no_grad():
            with autocast:
                if self.lab == 'ElnaggarLab':
                    embd_rpr = self.model(
                        input_ids=inputs['input_ids'],
                        attention_mask=inputs['attention_mask'],
                        decoder_input_ids=inputs['input_ids']
                    ).last_hidden_state
                elif 'peptidemtr' in self.model_name.lower():
                    embd_rpr = self.model(**inputs)['last_layer']
                else:
                    embd_rpr = self.model(**inputs).last_hidden_state
        output = []
        for idx in range(len(batch)):
            if self.lab == 'facebook' or self.lab == 'EvolutionaryScale':
                initial = 1
                final = len(batch[idx]) + 1
            elif self.lab == 'RostLab':
                initial = 0
                final = len(batch[idx].replace(' ', ''))
            else:
                initial = 0
                final = len(batch[idx])

            if self.average_pooling:
                output.append(embd_rpr[idx, initial:final].mean(0).float().detach().cpu().numpy())
            elif self.cls_token:
                output.append(embd_rpr[idx, 0].float().detach().cpu().numpy())
            else:
                output.append(embd_rpr[idx, initial:final].detach().cpu().numpy())

            if autocast:
                output[-1] = output[-1].astype(np.float16)
        return output
