#note: I think layer 0 is the input embedding.
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

import torch
import huggingface_hub
import transformers
from transformers import AutoConfig, AutoModel, AutoTokenizer
from transformers import pipeline
try:
    from transformers.utils import logging
except ImportError:
    print("Warning: Unable to importing transformers.utils logging")

import numpy as np
import nltk
try:
    nltk.data.find('tokenizers/punkt/PY3/english.pickle')
except:
    nltk.download('punkt')
#try:
#    nltk.data.find('tokenizers/punkt_tab/english/')
#except:
#    nltk.download('punkt_tab')

from nltk.tokenize import sent_tokenize

import os, sys

from collections import Counter
from tqdm import tqdm
from transformer_embs import transformer_embeddings

def set_hg_gated_access(access_token):
    """
    Local save of the access token for gated models on hg.
    
    Parameters
    ----------
    access_token : str
        Steps to get the access_token:
        1. Log in to your Hugging Face account.
        2. Click on your profile picture in the top right corner.
        3. Select ‘Settings’ from the dropdown menu.
        4. In the settings, you’ll find an option to generate a new token.
        Or, visit URL: https://huggingface.co/settings/tokens
    """
    huggingface_hub.login(access_token)
    print("Successfully login to Huggingface!")
    
def del_hg_gated_access():
    """
    Remove the access_token saved locally.

    """
    huggingface_hub.logout()
    print("Successfully logout to Huggingface!")

def set_logging_level(logging_level):
    """
    Set the logging level

    Parameters
    ----------
    logging_level : str
        set logging level, options: critical, error, warning, info, debug
    """
    logging_level = logging_level.lower()
    # default level is warning, which is in between "error" and "info"
    if logging_level in ['warn', 'warning']:
        logging.set_verbosity_warning()
    elif logging_level == "critical":
        logging.set_verbosity(50)
    elif logging_level == "error":
        logging.set_verbosity_error()
    elif logging_level == "info":
        logging.set_verbosity_info()
    elif logging_level == "debug":
        logging.set_verbosity_debug()
    else:
        print("Warning: Logging level {l} is not an option.".format(l=logging_level))
        print("\tUse one of: critical, error, warning, info, debug")

def set_tokenizer_parallelism(tokenizer_parallelism):
    if tokenizer_parallelism:
        os.environ["TOKENIZERS_PARALLELISM"] = "true"
    else:
        os.environ["TOKENIZERS_PARALLELISM"] = "false"

def get_device(device):
    """
    Get device and device number

    Parameters
    ----------
    device : str
        name of device: 'cpu', 'gpu', 'cuda', 'mps', or of the form 'gpu:k', 'cuda:k', or 'mps:0'
        where k is a specific device number

    Returns
    -------
    device : str
        final selected device name
    device_num : int
        device number, -1 for CPU
    """
    device = device.lower()
    if not device.startswith('cpu') and not device.startswith('gpu') and not device.startswith('cuda') and not device.startswith('mps'):
        print("device must be 'cpu', 'gpu', 'cuda', 'mps', or of the form 'gpu:k', 'cuda:k', or 'mps:0'")
        print("\twhere k is an integer value for the device")
        print("Trying CPUs")
        device = 'cpu'
    
    device_num = -1
    if device != 'cpu':
        attached = False
        
        if hasattr(torch.backends, "mps"):
            mps_available = torch.backends.mps.is_available()
        else:
            mps_available = False
       # print(f"MPS for Mac available: {mps_available}")
        if torch.cuda.is_available():
            if device == 'gpu' or device == 'cuda': 
                # assign to first gpu device number
                device = 'cuda'
                device_num = list(range(torch.cuda.device_count()))[0]
                attached = True
            elif 'gpu:' in device or 'cuda:' in device:
                try:
                    device_num = int(device.split(":")[-1])
                    device = 'cuda:' + str(device_num)
                    attached = True
                except:
                    attached = False
                    print(f"Device number {str(device_num)} does not exist! Use 'device = gpus' to see available gpu numbers.")
            elif 'gpus' in device:
                device = 'cuda'
                device_num = list(range(torch.cuda.device_count()))
                device = [device + ':' + str(num1) for num1 in device_num]
                attached = True
                print(f"Running on {str(len(device))} GPUs!")
                print(f"Available gpus to set: \n {device}")
        elif "mps" in device:
            if not torch.backends.mps.is_available():
                if not torch.backends.mps.is_built():
                    print("MPS not available because the current PyTorch install was not built with MPS enabled.")
                else:
                    print("MPS not available because the current MacOS version is not 12.3+ and/or you do not have an MPS-enabled device on this machine.")
            else:
                device_num = 0 # list(range(torch.cuda.device_count()))[0]
                device = 'mps:' + str(device_num)
                attached = True
               # print("Using Metal Performance Shaders (MPS) backend for GPU training acceleration!")
        else:
            attached = False
        if not attached:
            print("Unable to use MPS (Mac M1+), CUDA (GPU), using CPU")
            device = "cpu"
            device_num = -1

    return device, device_num

def get_model(model, tokenizer_only=False, config_only=False, hg_gated=False, hg_token="", trust_remote_code=False):
    """
    Get model and tokenizer from model string

    Parameters
    ----------
    model : str
        shortcut name for Hugging Face pretained model
        Full list https://huggingface.co/transformers/pretrained_models.html
    hg_gated : bool
        Set to True if the model is gated
    hg_token: str
        The token to access the gated model got in huggingface website
    
    Returns
    -------
    config
    tokenizer
    model
    """
    if "megatron-bert" in model:
        try:
            from transformers import BertTokenizer, MegatronBertForMaskedLM
        except:
            print("WARNING: You must install transformers>4.10 to use MegatronBertForMaskedLM")
            print("\tPlease try another model.")
            sys.exit()

        config = AutoConfig.from_pretrained(model, output_hidden_states=True)
        if not config_only:
            if "megatron-bert-cased" in model:
                tokenizer = BertTokenizer.from_pretrained('nvidia/megatron-bert-cased-345m')
            else:
                tokenizer = BertTokenizer.from_pretrained('nvidia/megatron-bert-uncased-345m')
            transformer_model = MegatronBertForMaskedLM.from_pretrained(model, config=config)
    elif "bigscience/bloom" in model:
        try:
            from transformers import BloomTokenizerFast, BloomModel, BloomConfig
        except:
            print("WARNING: You must install transformers>4.21.0 to use BloomModel")
            print("\tPlease try another model.")
            sys.exit()
        
        config = BloomConfig()
        if not config_only:
            tokenizer = BloomTokenizerFast.from_pretrained(model)
            transformer_model = BloomModel.from_pretrained(model, config=config)
    else:
        if hg_gated:
            set_hg_gated_access(access_token=hg_token)
        else: 
            pass
        config = AutoConfig.from_pretrained(model, output_hidden_states=True)
        if not config_only:
            tokenizer = AutoTokenizer.from_pretrained(model)
            transformer_model = AutoModel.from_pretrained(model, config=config, trust_remote_code=trust_remote_code)
            
    if config_only:
        return config
    elif tokenizer_only:
        # Do not know how to fix this. Some decoder-only files do not have pad_token.
        if tokenizer.pad_token is None:
            print("The language model entered might have issues since the model does not provide the padding_token.")
            print("Consider use BERT-like models instead if meeting errors.")
        #    tokenizer.pad_token = tokenizer.eos_token
        #    tokenizer.pad_token_id = tokenizer.eos_token_id
        return tokenizer
    else:
        if tokenizer.pad_token is None:
            print("The language model entered might have issues since the model does not provide the padding_token.")
            print("Consider use BERT-like models instead if meeting errors.")    
        #    tokenizer.pad_token = tokenizer.eos_token
        #    tokenizer.pad_token_id = tokenizer.eos_token_id        
        return config, tokenizer, transformer_model

def get_number_of_hidden_layers(model, logging_level = "error", hg_gated = False, hg_token = "", trust_remote_code = False):
    """
    Return the number of hidden layers for a given model.
    Returns -1 if the model's config doesn't have the num_hidden_layers parameter
    """
    set_logging_level(logging_level)
    config = get_model(model, config_only=True, hg_gated=hg_gated, hg_token=hg_token, trust_remote_code=trust_remote_code)
    number_of_hidden_layers = -1
    try:
        number_of_hidden_layers = config.num_hidden_layers
    except:
        print("Warning: Unable to get number of hidden layers")
        print("         num_hidden_layers is not a parameter of transformer_model.config")
        pass
    return number_of_hidden_layers

def hgTransformerGetPipeline(text_strings,
                            task = '',
                            model = '',
                            device = 'cpu',
                            tokenizer_parallelism = False,
                            logging_level = 'warning',
                            force_return_results = False,
                            hg_gated = False,
                            hg_token = "",
                            trust_remote_code = False,
                            set_seed = None,
                            **kwargs):
    """
    Simple interface getting Huggingface Pipeline
    https://huggingface.co/docs/transformers/main_classes/pipelines
    
    Parameters
    ----------
    text_strings : list
        list of strings, each is embedded separately
    task : str
        String representing task
        Task descriptions https://huggingface.co/docs/transformers/v4.20.1/en/task_summary
    model : str
        shortcut name for Hugging Face pretained model
        Full list https://huggingface.co/transformers/pretrained_models.html
    device : str
        name of device: 'cpu', 'gpu', or 'gpu:k' where k is a specific device number
    tokenizer_parallelism :  bool
        something
    logging_level : str
        set logging level, options: critical, error, warning, info, debug
    force_return_results : bool
        return results if they are not properly formatted for the task
    hg_gated : bool
        Set to True if the accessed model is gated
    hg_token: str
        The token needed to access the gated model, gen in huggingface website 
    trust_remote_code : bool
        use a model with custom code on the Huggingface Hub
    set_seed : int
        integer value for manually setting seed
    kwargs : dict
        pipeline task specific arguments
    
    Returns
    -------
    sentiment_scores : list
        list of dictionaries with sentiment scores and labels
    """

    if isinstance(set_seed, int):
        torch.manual_seed(set_seed)
    set_logging_level(logging_level)
    set_tokenizer_parallelism(tokenizer_parallelism)
    device, device_num = get_device(device)
    
    # check and adjust input types
    if not isinstance(text_strings, list):
        text_strings = [text_strings]
    
    if model:
        config, tokenizer, transformer_model = get_model(model,hg_gated=hg_gated, hg_token=hg_token, trust_remote_code=trust_remote_code)
        if device_num >= 0:
            task_pipeline = pipeline(task, model=model, tokenizer=tokenizer, device=device_num)
        else:
            task_pipeline = pipeline(task, model=model, tokenizer=tokenizer)
    else:
        if device_num >= 0:
            task_pipeline = pipeline(task, device=device_num)
        else:
            task_pipeline = pipeline(task)

    if transformers.__version__ >= "4.20" and "return_all_scores" in kwargs:
        return_all_scores = kwargs["return_all_scores"]
        if return_all_scores == True:
            kwargs["top_k"] = None
        else:
            kwargs["top_k"] = 1
        del kwargs["return_all_scores"]
    
    task_scores = []
    if task in ['question-answering', 'zero-shot-classification']:
        task_scores = task_pipeline(**kwargs)
    else:
        task_scores = task_pipeline(text_strings, **kwargs)
    
    return task_scores


# Convert floats to integers or propagate None
def _as_integer(x):
    if isinstance(x, int) or isinstance(x, np.integer):
        return x
    elif isinstance(x, float) or isinstance(x, np.floating):
        return int(x)
    else:
        return None

def hgTransformerGetTextGeneration(text_strings,
                            model = '',
                            device = 'cpu',
                            tokenizer_parallelism = False,
                            max_length = None,
                            max_new_tokens = 20,
                            min_length = 0,
                            min_new_tokens = None,
                            logging_level = 'warning',
                            force_return_results = False,
                            set_seed = None,
                            return_tensors = False,
                            #return_text = True,
                            return_full_text = True,
                            clean_up_tokenization_spaces = False,
                            prefix = '', 
                            handle_long_generation = None):
    # Prepare kwargs
    if max_new_tokens is not None and max_new_tokens <= 0:
        print(f"Warning: `max_new_tokens` must be greater than 0, but is {max_new_tokens}")
        print( "         Using default value…")
        max_new_tokens = None
    generation_kwargs = {
        'max_length': _as_integer(max_length),
        'min_length': _as_integer(min_length),
        'min_new_tokens': _as_integer(min_new_tokens)
    }
    # `max_new_tokens` should not be explicitly None
    max_new_tokens = _as_integer(max_new_tokens)
    if max_new_tokens is not None:
        generation_kwargs['max_new_tokens'] = max_new_tokens
    
    if return_tensors:
        if return_full_text:
            print("Warning: you set return_tensors and return_text (or return_full_text)")
            print("         Returning tensors only, as you cannot return both tensors and text.")
            print("         Please set return_tensors = FALSE if you need the generated text.")
        generated_texts = hgTransformerGetPipeline(text_strings = text_strings,
                            task = 'text-generation',
                            model = model,
                            device = device,
                            tokenizer_parallelism = tokenizer_parallelism,
                            logging_level = logging_level,
                            force_return_results = force_return_results,
                            set_seed = set_seed,
                            return_tensors = return_tensors, 
                            clean_up_tokenization_spaces = clean_up_tokenization_spaces, 
                            prefix = prefix,
                            handle_long_generation = handle_long_generation,
                            **generation_kwargs)
    else:
        generated_texts = hgTransformerGetPipeline(text_strings = text_strings,
                            task = 'text-generation',
                            model = model,
                            device = device,
                            tokenizer_parallelism = tokenizer_parallelism,
                            logging_level = logging_level,
                            force_return_results = force_return_results,
                            set_seed = set_seed,
                            #return_tensors = return_tensors, 
                            #return_text = return_text, 
                            return_full_text = return_full_text, 
                            clean_up_tokenization_spaces = clean_up_tokenization_spaces, 
                            prefix = prefix,
                            handle_long_generation = handle_long_generation,
                            **generation_kwargs)
    return generated_texts

def hgTransformerGetNER(text_strings,
                            model = '',
                            device = 'cpu',
                            tokenizer_parallelism = False,
                            logging_level = 'warning',
                            force_return_results = False,
                            set_seed = None):
    ner_scores = hgTransformerGetPipeline(text_strings = text_strings,
                            task = 'ner',
                            model = model,
                            device = device,
                            tokenizer_parallelism = tokenizer_parallelism,
                            logging_level = logging_level,
                            force_return_results = force_return_results,
                            set_seed = set_seed)
    return ner_scores

def hgTransformerGetZeroShot(sequences,
                            candidate_labels,
                            model = '',
                            device = 'cpu',
                            tokenizer_parallelism = False,
                            logging_level = 'warning',
                            force_return_results = False,
                            set_seed = None,
                            hypothesis_template = "This example is {}.",
                            multi_label = False):
    classifier_output = hgTransformerGetPipeline(text_strings = [],
                            task = 'zero-shot-classification',
                            model = model,
                            device = device,
                            tokenizer_parallelism = tokenizer_parallelism,
                            logging_level = logging_level,
                            force_return_results = force_return_results,
                            set_seed = set_seed,
                            sequences = sequences,
                            candidate_labels = candidate_labels,
                            hypothesis_template = hypothesis_template,
                            multi_label = multi_label)
    return classifier_output

def hgTransformerGetSentiment(text_strings,
                            model = '',
                            device = 'cpu',
                            tokenizer_parallelism = False,
                            logging_level = 'warning',
                            force_return_results = False,
                            set_seed = None,
                            return_all_scores = False,
                            function_to_apply = None):
    sentiment_scores = hgTransformerGetPipeline(text_strings = text_strings,
                            task = 'sentiment-analysis',
                            model = model,
                            device = device,
                            tokenizer_parallelism = tokenizer_parallelism,
                            logging_level = logging_level,
                            force_return_results = force_return_results,
                            set_seed = set_seed,
                            return_all_scores = return_all_scores,
                            function_to_apply = function_to_apply)
    return sentiment_scores

def hgTransformerGetSummarization(text_strings,
                            model = '',
                            device = 'cpu',
                            tokenizer_parallelism = False,
                            logging_level = 'warning',
                            force_return_results = False,
                            set_seed = None,
                            return_text = True,
                            return_tensors = False,
                            clean_up_tokenization_spaces = False, 
                            min_length = 10,
                            max_length = 200):
    summarizations = hgTransformerGetPipeline(text_strings = text_strings,
                            task = 'summarization',
                            model = model,
                            device = device,
                            tokenizer_parallelism = tokenizer_parallelism,
                            logging_level = logging_level,
                            force_return_results = force_return_results,
                            set_seed = set_seed,
                            return_text = return_text, 
                            return_tensors = return_tensors, 
                            clean_up_tokenization_spaces = clean_up_tokenization_spaces, 
                            min_length = min_length,
                            max_length = max_length)
    return summarizations

def hgTransformerGetQA(question,
                        context,
                        model = '',
                        device = 'cpu',
                        tokenizer_parallelism = False,
                        logging_level = 'warning',
                        force_return_results = False,
                        set_seed = None,
                        top_k = 1,
                        doc_stride = 128,
                        max_answer_len = 15,
                        max_seq_len = 384,
                        max_question_len = 64,
                        handle_impossible_answer = False):
    qas = hgTransformerGetPipeline(text_strings = [],
                            task = 'question-answering',
                            model = model,
                            device = device,
                            tokenizer_parallelism = tokenizer_parallelism,
                            logging_level = logging_level,
                            force_return_results = force_return_results,
                            set_seed = set_seed,
                            question = question, 
                            context = context, 
                            top_k = top_k, 
                            doc_stride = doc_stride, 
                            max_answer_len = max_answer_len, 
                            max_seq_len = max_seq_len, 
                            max_question_len = max_question_len, 
                            handle_impossible_answer = handle_impossible_answer)
    return qas

def hgTransformerGetTranslation(text_strings,
                            model = '',
                            device = 'cpu',
                            tokenizer_parallelism = False,
                            logging_level = 'warning',
                            force_return_results = False,
                            set_seed = None,
                            source_lang = '',
                            target_lang = '',
                            return_tensors = False,
                            return_text = True,
                            clean_up_tokenization_spaces = False, 
                            max_length = ''):
    task = 'translation'
    if source_lang and target_lang:
        task = "translation_{s}_to_{t}".format(s=source_lang, t=target_lang)
    translations = hgTransformerGetPipeline(text_strings = text_strings,
                            task = task,
                            model = model,
                            device = device,
                            tokenizer_parallelism = tokenizer_parallelism,
                            logging_level = logging_level,
                            force_return_results = force_return_results,
                            set_seed = set_seed,
                            src_lang = source_lang,
                            tgt_lang = target_lang,
                            return_tensors = return_tensors,
                            return_text = return_text,
                            clean_up_tokenization_spaces = clean_up_tokenization_spaces, 
                            max_length = max_length)
    return translations

def hgTransformerGetEmbedding(text_strings,
                              model = 'bert-large-uncased',
                              layers = 'all',
                              return_tokens = True,
                              max_token_to_sentence = 4,
                              device = 'cpu',
                              tokenizer_parallelism = False,
                              model_max_length = None,
                              hg_gated = False,
                              hg_token = "",
                              trust_remote_code = False,
                              logging_level = 'warning',
                              sentence_tokenize = True):
    """
    Simple Python method for embedding text with pretained Hugging Face models

    Parameters
    ----------
    text_strings : list
        list of strings, each is embedded separately
    model : str
        shortcut name for Hugging Face pretained model
        Full list https://huggingface.co/transformers/pretrained_models.html
    layers : str or list
        'all' or an integer list of layers to keep
    return_tokens : boolean
        return tokenized version of text_strings
    max_token_to_sentence : int
        maximum number of tokens in a string to handle before switching to embedding text
        sentence by sentence
    device : str
        name of device: 'cpu', 'gpu', or 'gpu:k' where k is a specific device number
    tokenizer_parallelism :  bool
        something
    model_max_length : int
        maximum length of the tokenized text
    hg_gated : bool
        Whether the accessed model is gated
    hg_token: str
        The token generated in huggingface website
    trust_remote_code : bool
        use a model with custom code on the Huggingface Hub
    logging_level : str
        set logging level, options: critical, error, warning, info, debug
    sentence_tokenize : bool
        tokenize long documents into sentences before embedding

    Returns
    -------
    all_embs : list
        embeddings for each item in text_strings
    all_toks : list, optional
        tokenized version of text_strings
    """
    #print("I am in hgTransformerGetEmbedding function now!!!!")
    #print(f"!!!!hg_gated: {hg_gated} !!!")
    #print(f"!!!!hg_token: {hg_token} !!!")
                                  
    set_logging_level(logging_level)
    set_tokenizer_parallelism(tokenizer_parallelism)
    device, device_num = get_device(device)

    config, tokenizer, transformer_model = get_model(model, hg_gated=hg_gated, hg_token=hg_token, trust_remote_code=trust_remote_code)

    if device != 'cpu':
        transformer_model.to(device)

    max_tokens = tokenizer.max_len_sentences_pair

    # check and adjust input types
    if not isinstance(text_strings, list):
        text_strings = [text_strings]

    if layers != 'all':
        if not isinstance(layers, list):
            layers = [layers]
        layers = [int(i) for i in layers]

    all_embs = []
    all_toks = []

    for text_string in text_strings:
        # if length of text_string is > max_token_to_sentence*4
        # embedd each sentence separately
        if len(text_string) > max_token_to_sentence*4 and sentence_tokenize:
            sentence_batch = [s for s in sent_tokenize(text_string)]
            if model_max_length is None:
                batch = tokenizer(sentence_batch, padding=True, truncation=True, add_special_tokens=True)
            else:
                batch = tokenizer(sentence_batch, padding=True, truncation=True, add_special_tokens=True, max_length=model_max_length)
            input_ids = torch.tensor(batch["input_ids"])
            attention_mask = torch.tensor(batch['attention_mask'])
            if device != 'cpu':
                input_ids = input_ids.to(device)
                attention_mask = attention_mask.to(device)

            if return_tokens:
                tokens = []
                for ids in input_ids:
                    tokens.extend([token for token in tokenizer.convert_ids_to_tokens(ids) if token != '[PAD]' and token != '<pad>'])
                all_toks.append(tokens)

            with torch.no_grad():
                hidden_states = transformer_model(input_ids,attention_mask=attention_mask)[-1]
                if layers != 'all':
                    hidden_states = [hidden_states[l] for l in layers]
                hidden_states = [h.tolist() for h in hidden_states]

            sent_embedding = []

            for l in range(len(hidden_states)): # iterate over layers
                layer_embedding = []
                for m in range(len(hidden_states[l])): # iterate over sentences
                    layer_embedding.extend([tok for ii, tok in enumerate(hidden_states[l][m]) if attention_mask[m][ii]>0])
                sent_embedding.append(layer_embedding)

            all_embs.append([[l] for l in sent_embedding])
        else:
            input_ids = tokenizer.encode(text_string, add_special_tokens=True)
            if return_tokens:
                tokens = tokenizer.convert_ids_to_tokens(input_ids)

            if device != 'cpu':
                input_ids = torch.tensor([input_ids]).to(device)
            else:
                input_ids = torch.tensor([input_ids])

            with torch.no_grad():
                hidden_states = transformer_model(input_ids)[-1]
                if layers != 'all':
                    hidden_states = [hidden_states[l] for l in layers]
                hidden_states = [h.tolist() for h in hidden_states]
                all_embs.append(hidden_states)
                if return_tokens:
                    all_toks.append(tokens)

    if hg_gated:
        del_hg_gated_access()                              
    if return_tokens:
        return all_embs, all_toks
    else:
        return all_embs

def hgDLATKTransformerGetEmbedding(text_strings,
                                   text_ids,
                                   group_ids,
                                   model = 'bert-large-uncased',
                                   layers = 'all',
                                   device = 'cpu',
                                   tokenizer_parallelism = False,
                                   model_max_length = None,
                                   hg_gated = False,
                                   hg_token = "",
                                   trust_remote_code = False,
                                   logging_level = 'warning',
                                    batch_size = 1,
                                    aggregations= ['mean'],
                                    aggregate_tokens_to_message = None,
                                    return_tokens = True, 
                                    debug = False
                                   ):
    """
    Simple Python method for embedding text with pretrained Hugging Face models
    using DLATK's hypercontextualized embeddings.

    Parameters
    ----------
    text_strings : list or str
        List of strings, each embedded separately. If a single string is
        passed, it will be wrapped into a list.
    text_ids : list
        List of unique identifiers for each text_string.
        If empty, will be set to range(len(text_strings)).
    group_ids : list
        List of unique identifiers for each group of text_strings.
        If empty, will be set to range(len(text_strings)).
    model : str
        Shortcut name for Hugging Face pretrained model.
    layers : str or list or int
        'all' or an integer / integer list of layers to keep.
    device : str
        Name of device: 'cpu', 'gpu', 'cuda', 'cuda:0', 'mps', etc.
    tokenizer_parallelism : bool
        Enable parallelism for tokenizer.
    model_max_length : int
        Maximum length of the tokenized text (currently unused here but kept
        for API compatibility).
    hg_gated : bool
        Whether the accessed model is gated.
    hg_token : str
        Hugging Face token for gated models.
    trust_remote_code : bool
        Use a model with custom code on the Hugging Face Hub.
    logging_level : str
        Logging level: 'critical', 'error', 'warning', 'info', 'debug'.
    batch_size : int
        Batch size for generating embeddings.
    aggregations : list
        Aggregations for message -> group (correl field) embeddings.
    aggregate_tokens_to_message : None or list or str
        Controls token -> message aggregation:
        - None      : do NOT aggregate tokens to messages; cf_embeddings = [].
        - 'mean' or ['mean'] : aggregate tokens to messages using "mean".
        - list of str: passed to DLATK as wordAggregations.
    return_tokens : bool
        If True, return token embeddings and tokens along with cf embeddings.
    debug : bool
        If True, print detailed debug info.

    Returns
    -------
    cf_ids_all : list
        list of unique identifiers for each group of text_strings (only returned if return_tokens is False)
    cf_embeddings : list
        embeddings for each group of text_strings (only returned if return_tokens is False)
    msg_ids_all : list
        list of unique identifiers for each message (only returned if return_tokens is True)
    token_embeddings_all_grouped : list, optional
        embeddings for each token in text_strings (only returned if return_tokens is True)
    tokens_all_grouped : list, optional
        tokenized version of text_strings (only returned if return_tokens is True)
    """ 
    def getMessagesForCorrelFieldGroups(cfGrp, text_strings, group_ids, text_ids):
        """Return a list of [cfId, msgId, msg] for the given cfGrp."""
        filtered_rows = [
            (group_id, text_ids[i], text_strings[i])
            for i, group_id in enumerate(group_ids)
            if group_id in cfGrp
        ]
        return filtered_rows

    # -------------------------------------------------------------------------
    # Setup logging / device / model
    # -------------------------------------------------------------------------
    set_logging_level(logging_level)
    set_tokenizer_parallelism(tokenizer_parallelism)
    device, device_num = get_device(device)

    config, tokenizer, transformer_model = get_model(
        model,
        hg_gated=hg_gated,
        hg_token=hg_token,
        trust_remote_code=trust_remote_code
    )

    if device != 'cpu':
        transformer_model.to(device)

    # -------------------------------------------------------------------------
    # Normalize inputs
    # -------------------------------------------------------------------------
    # If a single string is given, wrap into a list
    if isinstance(text_strings, str):
        text_strings = [text_strings]

    # Layers: 'all' or int or list of int
    if layers == 'all':
        layers = list(range(transformer_model.config.num_hidden_layers))
    elif isinstance(layers, int):
        layers = [layers]
    layers = [int(i) for i in layers]

    # text_ids: default 0..n-1, must be unique and same length as texts
    if text_ids == []:
        text_ids = list(range(len(text_strings)))
    assert len(text_ids) == len(text_strings), \
        "Length of text_ids must be equal to length of text_strings"
    assert len(set(text_ids)) == len(text_ids), "text_ids must be unique"

    # group_ids: default 0..n-1, same length as texts
    if group_ids == []:
        group_ids = list(range(len(text_strings)))
    assert len(group_ids) == len(text_strings), \
        "Length of group_ids must be equal to length of text_strings"

    # Interpret token->message aggregation
    if aggregate_tokens_to_message is None:
        do_aggregate = False
        # DLATK's transformer_embeddings still needs a wordAggregations value,
        # but we will *not* call message_aggregate.
        word_aggregations = ['mean']
    else:
        do_aggregate = True
        if isinstance(aggregate_tokens_to_message, str):
            word_aggregations = [aggregate_tokens_to_message]
        else:
            try:
                word_aggregations = list(aggregate_tokens_to_message)
            except TypeError:
                word_aggregations = [str(aggregate_tokens_to_message)]

    if debug:
        print("=== hgDLATKTransformerGetEmbedding DEBUG ===")
        print(f"n_texts = {len(text_strings)}")
        print(f"text_ids = {list(text_ids)}")
        print(f"group_ids = {list(group_ids)}")
        print(f"layers = {layers}")
        print(f"do_aggregate = {do_aggregate}")
        print(f"word_aggregations (token->msg) = {word_aggregations}")
        print(f"aggregations (msg->group) = {aggregations}")
        print(f"return_tokens = {return_tokens}")
        print(f"aggregate_tokens_to_message = {aggregate_tokens_to_message}")
        print("------------------------------------------")

    # -------------------------------------------------------------------------
    # Create embedding generator
    # -------------------------------------------------------------------------
    embedding_generator = transformer_embeddings(
        modelObj=transformer_model,
        tokenizerObj=tokenizer,
        layersToKeep=layers,
        aggregations=aggregations,           # message -> group
        layerAggregations=['concatenate'],
        wordAggregations=word_aggregations,  # token -> message
        maxTokensPerSeg=255,
        batchSize=batch_size,
        noContext=False
    )

    from tqdm import tqdm
    import numpy as np

    # =========================================================================
    # PATH A: TOKEN-ONLY MODE (your main use case)
    # aggregate_tokens_to_message is None AND return_tokens is True
    # -> we process each text independently, avoid DLATK multi-message logic.
    # =========================================================================
    if (aggregate_tokens_to_message is None) and return_tokens:
        if debug:
            print(">>> TOKEN-ONLY MODE: per-text processing, no msg/group aggregation.")

        cf_embeddings = []  # explicit: no message/group embeddings

        msg_ids_all = []
        token_embeddings_all = []
        tokens_all = []

        # iterate over each text independently
        for idx, (cfId, msgId, msg) in enumerate(zip(group_ids, text_ids, text_strings)):
            if debug:
                print(f"\n--- Processing text {idx} (cfId={cfId}, msgId={msgId}) ---")

            # One cfId, one message
            groupedMessageRows = [[cfId, [[msgId, msg]]]]

            tokenIdsDict, (cfId_seq, msgId_seq) = embedding_generator.prepare_messages(
                groupedMessageRows,
                sent_tok_onthefly=True,
                noContext=False
            )

            n_tok_segments = len(tokenIdsDict["input_ids"])
            if debug:
                print(f"prepare_messages -> n_tok_segments = {n_tok_segments}")

            if n_tok_segments == 0:
                if debug:
                    print("No segments produced for this text; skipping.")
                continue

            encSelectedLayers = embedding_generator.generate_transformer_embeddings(
                tokenIdsDict
            )
            if encSelectedLayers is None:
                if debug:
                    print("generate_transformer_embeddings returned None; skipping.")
                continue

            # decode tokens for each segment
            decoded_tokens = [
                tokenizer.convert_ids_to_tokens(input_ids)
                for input_ids in tokenIdsDict["input_ids"]
            ]

            arr = encSelectedLayers[0]
            n_emb_segments = arr.shape[0]

            if debug:
                print(f"decoded_tokens segments = {len(decoded_tokens)}")
                print(f"encSelectedLayers[0] segments = {n_emb_segments}")

            # --- critical alignment: keep only the common number of segments ---
            if n_tok_segments != n_emb_segments:
                n_common = min(n_tok_segments, n_emb_segments)
                if debug:
                    print(
                        f"*** Segment mismatch for msgId={msgId}: "
                        f"n_tok_segments={n_tok_segments}, "
                        f"n_emb_segments={n_emb_segments}. "
                        f"Keeping first {n_common} segments."
                    )
                decoded_tokens = decoded_tokens[:n_common]
                arr = arr[:n_common, :, :]
            else:
                n_common = n_tok_segments

            token_embeddings = []
            for seg_i in range(n_common):
                n_tok_seg = len(decoded_tokens[seg_i])
                seg_emb = arr[seg_i, :n_tok_seg, :].tolist()
                token_embeddings.append(seg_emb)

            # extend global lists – all three get n_common entries
            tokens_all.extend(decoded_tokens)
            token_embeddings_all.extend(token_embeddings)
            msg_ids_all.extend([msgId] * n_common)

            if debug:
                print(f"Accumulated so far: len(msg_ids_all)={len(msg_ids_all)}, "
                      f"len(token_embeddings_all)={len(token_embeddings_all)}, "
                      f"len(tokens_all)={len(tokens_all)}")

        # no tokens at all?
        if len(msg_ids_all) == 0:
            return cf_embeddings, [], [], []

        # defensive check (should now be aligned)
        if not (len(msg_ids_all) == len(token_embeddings_all) == len(tokens_all)):
            min_len = min(len(msg_ids_all), len(token_embeddings_all), len(tokens_all))
            print(
                "Warning: residual length mismatch in token-only mode "
                f"(msg_ids_all={len(msg_ids_all)}, "
                f"token_embeddings_all={len(token_embeddings_all)}, "
                f"tokens_all={len(tokens_all)}). "
                f"Truncating to {min_len}."
            )
            msg_ids_all = msg_ids_all[:min_len]
            token_embeddings_all = token_embeddings_all[:min_len]
            tokens_all = tokens_all[:min_len]

        # group segments by msg_id
        token_embeddings_all_grouped = {}
        tokens_all_grouped = {}
        for i in range(len(msg_ids_all)):
            mid = msg_ids_all[i]
            if mid not in token_embeddings_all_grouped:
                token_embeddings_all_grouped[mid] = [token_embeddings_all[i]]
                tokens_all_grouped[mid] = [tokens_all[i]]
            else:
                token_embeddings_all_grouped[mid].append(token_embeddings_all[i])
                tokens_all_grouped[mid].append(tokens_all[i])

        msg_ids_all_grouped = sorted(token_embeddings_all_grouped.keys())
        token_embeddings_all_grouped = [
            token_embeddings_all_grouped[mid] for mid in msg_ids_all_grouped
        ]
        tokens_all_grouped = [
            tokens_all_grouped[mid] for mid in msg_ids_all_grouped
        ]

        if debug:
            print("\n=== FINAL GROUPED (TOKEN-ONLY) ===")
            print(f"msg_ids_all_grouped = {msg_ids_all_grouped}")
            print(f"#token_embeddings_all_grouped = {len(token_embeddings_all_grouped)}")
            print(f"#tokens_all_grouped           = {len(tokens_all_grouped)}")

        return cf_embeddings, msg_ids_all_grouped, token_embeddings_all_grouped, tokens_all_grouped

    # =========================================================================
    # PATH B: USE DLATK'S BATCHING + MESSAGE/GROUP AGGREGATION
    # (covers cases where aggregate_tokens_to_message is not None, and/or
    #  return_tokens=False)
    # =========================================================================

    from collections import Counter
    num_rows_per_group = Counter(group_ids)

    msgsInBatch = 0
    cfGroups = [[]]
    for cfId, cfNumMsgs in num_rows_per_group.items():
        if msgsInBatch + cfNumMsgs > batch_size and len(cfGroups[-1]) >= 1:
            cfGroups.append([])
            msgsInBatch = 0
        cfGroups[-1].append(cfId)
        msgsInBatch += cfNumMsgs

    num_cfs = 0
    cfGroups = [cfGrp for cfGrp in cfGroups if cfGrp]
    
    msg_embeddings = []
    cf_embeddings = []
    token_embeddings_all = []
    tokens_all = []
    msg_ids_all = []
    cf_ids_all = []
    for cfGrp in tqdm(cfGroups):
        # mIdSeen = set() #currently seen message ids
        # mIdList = [] #only for keepMsgFeats

        cfId_msgId_map = {}
        groupedMessageRows = {}
        messageRows = getMessagesForCorrelFieldGroups(
            cfGrp=cfGrp,
            text_strings=text_strings,
            group_ids=group_ids,
            text_ids=text_ids
        )

        if debug:
            msg_pairs = [(cfId, msgId) for (cfId, msgId, _msg) in messageRows]
            print(f"messageRows (cfId, msgId) pairs: {msg_pairs}")
            print(f"Number of messageRows in this batch: {len(messageRows)}")

        for cfId, msgId, msg in messageRows:
            if cfId not in cfId_msgId_map:
                cfId_msgId_map[cfId] = set()
                groupedMessageRows[cfId] = []
            if msgId not in cfId_msgId_map[cfId]:
                cfId_msgId_map[cfId].add(msgId)
                groupedMessageRows[cfId].append([msgId, msg])
        groupedMessageRows = [[cfId, groupedMessageRows[cfId]] for cfId in cfId_msgId_map]

        tokenIdsDict, (cfId_seq, msgId_seq) = embedding_generator.prepare_messages(groupedMessageRows, sent_tok_onthefly=True, noContext=False)
        if len(tokenIdsDict["input_ids"]) == 0:
            continue

        encSelectedLayers = embedding_generator.generate_transformer_embeddings(
            tokenIdsDict
        )
        if encSelectedLayers is None:
            if debug:
                print("generate_transformer_embeddings returned None; skipping this batch.")
            continue

        msg_reps, msgIds_new, cfIds_new = embedding_generator.message_aggregate(encSelectedLayers, msgId_seq, cfId_seq)

        if all([len(cfId_msgId_map[cfId]) == 1 for cfId in cfId_msgId_map]):
            cf_reps = msg_reps
        else:
            msg_reps_dict = dict(zip([i for i in msgIds_new], [msg_reps[i] for i in range(len(msgIds_new))]))
            cf_reps, cfIds_new = embedding_generator.correl_field_aggregate(msg_reps_dict, cfId_msgId_map)
            cf_reps = [[cf_reps[i][j].rep for j in range(len(aggregations))] for i in range(len(cf_reps))]    

        msg_embeddings.extend(msg_reps)
        cf_embeddings.extend(cf_reps)
        msg_ids_all.extend(msgIds_new)
        cf_ids_all.extend(cfIds_new)

        if return_tokens:
            decoded_tokens = [
                tokenizer.convert_ids_to_tokens(input_ids)
                for input_ids in tokenIdsDict["input_ids"]
            ]

            arr = encSelectedLayers[0]
            n_tok_segments = len(decoded_tokens)
            n_emb_segments = arr.shape[0]

            if debug:
                lens_tokens = [len(toks) for toks in decoded_tokens]
                print(f"decoded_tokens: n_segments = {len(decoded_tokens)}, "
                      f"lengths (first 5): {lens_tokens[:5]}")
                print(f"encSelectedLayers[0] segments = {n_emb_segments}")

            # align segments
            if n_tok_segments != n_emb_segments:
                n_common = min(n_tok_segments, n_emb_segments)
                if debug:
                    print(
                        f"*** Segment mismatch in batch {batch_idx}: "
                        f"n_tok_segments={n_tok_segments}, "
                        f"n_emb_segments={n_emb_segments}. "
                        f"Keeping first {n_common} segments."
                    )
                decoded_tokens = decoded_tokens[:n_common]
                arr = arr[:n_common, :, :]
            else:
                n_common = n_tok_segments

            token_embeddings = []
            for i_seg in range(n_common):
                n_tok_seg = len(decoded_tokens[i_seg])
                seg_emb = arr[i_seg, :n_tok_seg, :].tolist()
                token_embeddings.append(seg_emb)

            token_embeddings_all.extend(token_embeddings)
        # import pdb; pdb.set_trace()
    
    msg_embeddings = [msg_embeddings[i].tolist() for i in range(len(msg_embeddings))]
    cf_embeddings = [cf_embeddings[i].tolist() if isinstance(cf_embeddings[i], np.ndarray) else cf_embeddings[i] for i in range(len(cf_embeddings))] 

    if return_tokens: 
        # msg_ids_all = list(map(lambda x: x[0], msgId_seq))
        # print (msg_ids_all, msgId_seq)
        # assert len(msg_ids_all) == len(token_embeddings_all) == len(tokens_all), "Length of msg_ids_all, token_embeddings_all, tokens_all must be equal"
        # Return cf_embeddings, msg ids, token embeddings and tokens.
        # But msg_ids should be a unique list (currently msg_ids_all has repeated msg_ids whenever a message is split into multiple submessages).
        # Then group token_embeddings for indices with same message_id, same for tokens

        # Group token_embeddings for indices with same message_id, same for tokens
        token_embeddings_all_grouped = {}
        tokens_all_grouped = {}
        for i in range(len(msg_ids_all)):
            if msg_ids_all[i] not in token_embeddings_all_grouped:
                token_embeddings_all_grouped[msg_ids_all[i]] = []
                tokens_all_grouped[msg_ids_all[i]] = []
            
            token_embeddings_all_grouped[msg_ids_all[i]].extend(token_embeddings_all[i])
            tokens_all_grouped[msg_ids_all[i]].extend(tokens_all[i])
    
        msg_ids_all_grouped = sorted(token_embeddings_all_grouped.keys())
        token_embeddings_all_grouped = [token_embeddings_all_grouped[msg_id] for msg_id in msg_ids_all_grouped]
        tokens_all_grouped = [tokens_all_grouped[msg_id] for msg_id in msg_ids_all_grouped]
        return cf_embeddings, msg_ids_all_grouped, token_embeddings_all_grouped, tokens_all_grouped
    return cf_ids_all, cf_embeddings

def hgTokenizerGetTokens(text_strings,
                              model = 'bert-large-uncased',
                              max_token_to_sentence = 4,
                              device = 'cpu',
                              tokenizer_parallelism = False,
                              model_max_length = None,
                              hg_gated = False,
                              hg_token = "",
                              trust_remote_code = False,
                              logging_level = 'warning'):
    """
    Simple Python method for embedding text with pretained Hugging Face models

    Parameters
    ----------
    text_strings : list
        list of strings, each is embedded separately
    model : str
        shortcut name for Hugging Face pretained model
        Full list https://huggingface.co/transformers/pretrained_models.html
    max_token_to_sentence : int
        maximum number of tokens in a string to handle before switching to embedding text
        sentence by sentence
    device : str
        name of device: 'cpu', 'gpu', or 'gpu:k' where k is a specific device number
    tokenizer_parallelism :  bool
        something
    model_max_length : int
        maximum length of the tokenized text
    hg_gated : bool
        Set to True if the accessed model is gated
    hg_token: str
        The token to access the gated model gen in hg website
    trust_remote_code : bool
        use a model with custom code on the Huggingface Hub
    logging_level : str
        set logging level, options: critical, error, warning, info, debug

    Returns
    -------
    all_embs : list
        embeddings for each item in text_strings
    all_toks : list, optional
        tokenized version of text_strings
    """
    try:
        set_logging_level(logging_level)
    except NameError:
        pass
    set_tokenizer_parallelism(tokenizer_parallelism)
    device, device_num = get_device(device)

    tokenizer = get_model(model, tokenizer_only=True, hg_gated=hg_gated, hg_token=hg_token, trust_remote_code=trust_remote_code)

    if device != 'cpu':
        tokenizer.to(device)

    max_tokens = tokenizer.max_len_sentences_pair

    # check and adjust input types
    if not isinstance(text_strings, list):
        text_strings = [text_strings]

    all_toks = []

    for text_string in text_strings:
        # if length of text_string is > max_token_to_sentence*4
        # embedd each sentence separately
        if len(text_string) > max_token_to_sentence*4:
            sentence_batch = [s for s in sent_tokenize(text_string)]
            if model_max_length is None:
                batch = tokenizer(sentence_batch, padding=True, truncation=True, add_special_tokens=True)
            else:
                batch = tokenizer(sentence_batch, padding=True, truncation=True, add_special_tokens=True, max_length=model_max_length)
            input_ids = torch.tensor(batch["input_ids"])

            tokens = []
            for ids in input_ids:
                tokens.extend([token for token in tokenizer.convert_ids_to_tokens(ids) if token != '[PAD]' and token != '<pad>'])
            all_toks.append(tokens)

        else:
            input_ids = tokenizer.encode(text_string, add_special_tokens=True)
            tokens = tokenizer.convert_ids_to_tokens(input_ids)

            all_toks.append(tokens)

    return all_toks


    
### EXAMPLE TEST CODE:
#if __name__   == '__main__':
#   embeddings, tokens = hgTransformerGetEmbedding("Here is one sentence.", layers=[0,10], device="gpu", logging_level="warn")
#   print(np.array(embeddings).shape)
#   print(tokens)
#
#   embeddings, tokens = hgTransformerGetEmbedding("Here is more sentences. But why is not . and , and ? indicated with SEP?", layers=[0,10], device="gpu")
#   print(np.array(embeddings).shape)
#   print(tokens)
