from abc import ABC, abstractmethod
import logging

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
import torch
import transformers
from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForCausalLM,
    AutoTokenizer,
    BloomForCausalLM,
    GPTNeoXForCausalLM,
    LlamaForCausalLM,
    LlamaTokenizer,
)
from transformers.generation.utils import GenerationConfig

from ._registry import register_model

__all__ = ['tnl3', 'tnl', 'llama', 'baichuan', 'baichuan2', 'bloom', 'mamba', 'LLModel', 'HuggingFaceModel']

# detect transformers version
logging.info(f'transformers version: {transformers.__version__}')


class LLModel(ABC):

    @abstractmethod
    def get_tokenizer(self):
        pass

    @abstractmethod
    def get_model(self):
        pass

    @abstractmethod
    def chat(self):
        pass


class HuggingFaceModel(LLModel):

    def __init__(self, tokenizer, model):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        model.eval()
        model.to(self.device)

        self.tok = tokenizer
        self.model = model

    def get_tokenizer(self) -> AutoTokenizer:
        return self.tok

    def get_model(self) -> AutoModelForCausalLM:
        return self.model

    def chat(self, messages, config):
        inputs = self.tok(messages, return_tensors='pt').to(self.device)
        outputs = self.model.generate(inputs.input_ids, do_sample=True, **config)
        try:
            response = self.tok.batch_decode(outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0]
        except Exception as e:
            logging.error(e)
            response = ''

        return response


class HuggingFaceModelWrap(HuggingFaceModel):

    def __init__(self, repo_or_path, tok_configs, model_configs):
        tok: AutoTokenizer = AutoTokenizer.from_pretrained(repo_or_path,
                                                           trust_remote_code=True,
                                                           add_bos_token=False,
                                                           padding_side='left',
                                                           **tok_configs)
        model: AutoModelForCausalLM = AutoModelForCausalLM.from_pretrained(repo_or_path,
                                                                           trust_remote_code=True,
                                                                           **model_configs)

        model_vocab_size = model.get_input_embeddings().weight.size(0)
        tokenzier_vocab_size = len(tok)
        logging.info(f'Vocab of the base model: {model_vocab_size}')
        logging.info(f'Vocab of the tokenizer: {tokenzier_vocab_size}')

        super().__init__(tok, model)


@register_model
def tnl3(repo_or_path) -> HuggingFaceModel:
    try:
        import tiktoken
    except Exception as e:
        logging.error('Try to pip install tiktoken!')
        raise e

    tok_configs = {'use_fast': True}
    model_configs = {'torch_dtype': torch.bfloat16}

    return HuggingFaceModelWrap(repo_or_path, tok_configs, model_configs)


@register_model
def tnl(repo_or_path) -> HuggingFaceModel:
    tok_configs = {'use_fast': True}
    model_configs = {'torch_dtype': torch.bfloat16}

    model = HuggingFaceModelWrap(repo_or_path, tok_configs, model_configs)

    model.tok.pad_token_id = 0
    model.model.config.pad_token_id = 0

    return model


@register_model
def llama(repo_or_path) -> HuggingFaceModel:

    tok = LlamaTokenizer.from_pretrained(repo_or_path,
                                         padding_side='left',
                                         add_bos_token=False,
                                         use_fast=True,
                                         legacy=False)
    if 'llama2' or 'LLaMA2' or 'llama-2' or 'LLaMA-2' in repo_or_path:
        config = AutoConfig.from_pretrained(repo_or_path, pretraining_tp=0)
    else:
        config = AutoConfig.from_pretrained(repo_or_path)

    model = LlamaForCausalLM.from_pretrained(repo_or_path, config=config, torch_dtype=torch.bfloat16)

    model_vocab_size = model.get_input_embeddings().weight.size(0)
    tokenzier_vocab_size = len(tok)
    logging.info(f'Vocab of the base model: {model_vocab_size}')
    logging.info(f'Vocab of the tokenizer: {tokenzier_vocab_size}')

    tok.pad_token = tok.eos_token
    model.config.end_token_id = tok.eos_token_id
    model.config.pad_token_id = model.config.eos_token_id


@register_model
def baichuan(repo_or_path) -> HuggingFaceModel:
    model_configs = {'torch_dtype': torch.float16}
    model = HuggingFaceModelWrap(repo_or_path, {}, model_configs)

    model.tok.pad_token_id = 0
    model.model.config.pad_token_id = 0

    return model


@register_model
def baichuan2(repo_or_path) -> HuggingFaceModel:
    tok_configs = {'use_fast': False}
    model_configs = {'torch_dtype': torch.bfloat16}
    model = HuggingFaceModelWrap(repo_or_path, tok_configs, model_configs)

    model.model.generation_config = GenerationConfig.from_pretrained(repo_or_path)

    return model


@register_model
def bloom(repo_or_path) -> HuggingFaceModel:
    tok = AutoTokenizer.from_pretrained(repo_or_path, trust_remote_code=True, add_bos_token=False, padding_side='left')

    model = BloomForCausalLM.from_pretrained(repo_or_path)

    model_vocab_size = model.get_input_embeddings().weight.size(0)
    tokenzier_vocab_size = len(tok)
    logging.info(f'Vocab of the base model: {model_vocab_size}')
    logging.info(f'Vocab of the tokenizer: {tokenzier_vocab_size}')

    return HuggingFaceModel(tok, model)


class MambaModel(LLModel):

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

        tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
        model = MambaLMHeadModel.from_pretrained(repo_or_path, device=device, dtype=torch.float16)
        model.eval()

        self.tok = tokenizer
        self.model = model

    def get_tokenizer(self) -> AutoTokenizer:
        return self.tok

    def get_model(self) -> AutoModelForCausalLM:
        return self.model

    def chat(self, messages, config={}):
        input_ids = self.tok(messages, return_tensors='pt').to(self.device).input_ids
        max_length = input_ids.shape[1] + config['max_new_tokens']
        del config['max_new_tokens']

        outputs = self.model.generate(input_ids, max_length=max_length, cg=True, enable_timing=False, **config)
        try:
            response = self.tok.batch_decode(outputs[:, input_ids.shape[1]:], skip_special_tokens=True)[0]
        except Exception as e:
            logging.error(e)
            response = ''

        return response


@register_model
def mamba(repo_or_path) -> LLModel:
    return MambaModel(repo_or_path)
