from transformers import AutoTokenizer ,AutoModelForCausalLM
import re
# Speller and punctuation:
import os
import yaml
import torch
from torch import package
# not very necessary
import textwrap
from textwrap3 import wrap

# util function to get expected len after tokenizing
def get_length_param(text: str, tokenizer) -> str:
    tokens_count = len(tokenizer.encode(text))
    if tokens_count <= 15:
        len_param = '1'
    elif tokens_count <= 50:
        len_param = '2'
    elif tokens_count <= 256:
        len_param = '3'
    else:
        len_param = '-'
    return len_param

def remove_duplicates(S):
    S = re.sub(r'[a-zA-Z]+', '', S) #Remove english
    S = S.split()
    result = ""
    for subst in S:
        if subst not in result:
            result += subst+" "
    return result.rstrip()

def removeSigns(S):
    last_index = max(S.rfind("."), S.rfind("!"))
    if last_index >= 0:
        S = S[:last_index+1]
    return S

def prepare_punct():
    torch.hub.download_url_to_file('https://raw.githubusercontent.com/snakers4/silero-models/master/models.yml',
                                   'latest_silero_models.yml',
                                   progress=False)

    with open('latest_silero_models.yml', 'r') as yaml_file:
        models = yaml.load(yaml_file, Loader=yaml.SafeLoader)
    model_conf = models.get('te_models').get('latest')

    # Prepare punctuation fix
    model_url = model_conf.get('package')

    model_dir = "downloaded_model"
    os.makedirs(model_dir, exist_ok=True)
    model_path = os.path.join(model_dir, os.path.basename(model_url))

    if not os.path.isfile(model_path):
        torch.hub.download_url_to_file(model_url,
                                       model_path,
                                       progress=True)

    imp = package.PackageImporter(model_path)
    model_punct = imp.load_pickle("te_model", "model")

    return model_punct

def initialize():
    """ Loading the model """
    torch.backends.quantized.engine = 'qnnpack' # Just for the specific machine architecture
    fit_checkpoint = "WarBot"
    tokenizer = AutoTokenizer.from_pretrained(fit_checkpoint)
    model = AutoModelForCausalLM.from_pretrained(fit_checkpoint)
    model_punсt = prepare_punct()
    return (model,tokenizer,model_punсt)

def split_string(string,n=256):
    return [string[i:i+n] for i in range(0, len(string), n)]

def get_response(quote:str,model,tokenizer,model_punct):
    # encode the input, add the eos_token and return a tensor in Pytorch
    user_inpit_ids = tokenizer.encode(f"|0|{get_length_param(quote, tokenizer)}|" \
                                                  + quote + tokenizer.eos_token, return_tensors="pt")

    chat_history_ids = user_inpit_ids # To be changed

    tokens_count = len(tokenizer.encode(quote))
    if tokens_count < 15:
        no_repeat_ngram_size = 2
    else:
        no_repeat_ngram_size = 1

    output_id = model.generate(
                chat_history_ids,
                num_return_sequences=1, # use for more variants, but have to print [i]
                max_length=200, #512
                no_repeat_ngram_size=no_repeat_ngram_size, #3
                do_sample=True, #True
                top_k=50,#50
                top_p=0.9, #0.9
                temperature = 0.4, # was 0.6, 0 for greedy
                #mask_token_id=tokenizer.mask_token_id,
                eos_token_id=tokenizer.eos_token_id,
                #unk_token_id=tokenizer.unk_token_id,
                pad_token_id=tokenizer.pad_token_id,
                #pad_token_id=tokenizer.eos_token_id,
                #device='cpu'
            )

    response = tokenizer.decode(output_id[0], skip_special_tokens=True)
    response = removeSigns(response)
    response = response.split(quote)[-1]  # Remove the Quote
    response = re.sub(r'[^0-9А-Яа-яЁёa-zA-z;., !()/\-+:?]', '',
                      response)  # Clear the response, remains only alpha-numerical values
    response = remove_duplicates(re.sub(r"\d{4,}", "", response))  # Remove the consequent numbers with 4 or more digits
    response = re.sub(r'\.\.+', '', response) # Remove the "....." thing

    maxLen = 170

    try:
        if len(response)>maxLen: # We shall play with it
            resps = wrap(response,maxLen)
            for i in range(len(resps)):
                resps[i] = model_punct.enhance_text(resps[i], lan='ru')
                response = ''.join(resps)
        else:
            response = model_punct.enhance_text(response, lan='ru')
    except:
	    pass # sometimes the string is getting too long

    response = re.sub(r'[UNK]', '', response)  # Remove the [UNK] thing
    return response

#if __name__ == '__main__':
    #model,tokenizer,model_punct = initialize()
    #quote = "Это хорошо, но глядя на ролик, когда ефиопские толпы в Израиле громят машины и нападают на улице на израильтян - задумаешься, куда все движется"
    #print('please wait...')
    #response = wrap(get_response(quote,model,tokenizer,model_punct),60)
    #for phrase in response:
    #    print(phrase)