"""
This module contains the helper functions to get the word alignment mapping between two sentences.
"""

import torch
import itertools
import transformers
from transformers import logging

# Set the verbosity to error, so that the warning messages are not printed
logging.set_verbosity_warning()
logging.set_verbosity_error()


def select_model(model_name):
    """
    Select Model
    """
    if model_name == "Google-mBERT (Base-Multilingual)":
        model_name="bert-base-multilingual-cased"
    elif model_name == "Neulab-AwesomeAlign (Bn-En-0.5M)":
        model_name="musfiqdehan/bn-en-word-aligner"
    elif model_name == "BUET-BanglaBERT (Large)":
        model_name="csebuetnlp/banglabert_large"
    elif model_name == "SagorSarker-BanglaBERT (Base)":
        model_name="sagorsarker/bangla-bert-base"
    elif model_name == "SentenceTransformers-LaBSE (Multilingual)":
        model_name="sentence-transformers/LaBSE"

    return model_name


def get_alignment_mapping(source="", target="", model_name=""):
    """
    Get Aligned Words
    """
    model_name = select_model(model_name)

    model = transformers.BertModel.from_pretrained(model_name)
    tokenizer = transformers.BertTokenizer.from_pretrained(model_name)

    # pre-processing
    sent_src, sent_tgt = source.strip().split(), target.strip().split()

    token_src, token_tgt = [tokenizer.tokenize(word) for word in sent_src], [
        tokenizer.tokenize(word) for word in sent_tgt]
    
    wid_src, wid_tgt = [tokenizer.convert_tokens_to_ids(x) for x in token_src], [
        tokenizer.convert_tokens_to_ids(x) for x in token_tgt]
    
    ids_src, ids_tgt = tokenizer.prepare_for_model(list(itertools.chain(*wid_src)), return_tensors='pt', model_max_length=tokenizer.model_max_length, truncation=True)['input_ids'], tokenizer.prepare_for_model(list(itertools.chain(*wid_tgt)), return_tensors='pt', truncation=True, model_max_length=tokenizer.model_max_length)['input_ids']
    sub2word_map_src = []

    for i, word_list in enumerate(token_src):
        sub2word_map_src += [i for x in word_list]

    sub2word_map_tgt = []

    for i, word_list in enumerate(token_tgt):
        sub2word_map_tgt += [i for x in word_list]

    # alignment
    align_layer = 8

    threshold = 1e-3

    model.eval()

    with torch.no_grad():
        out_src = model(ids_src.unsqueeze(0), output_hidden_states=True)[
            2][align_layer][0, 1:-1]
        out_tgt = model(ids_tgt.unsqueeze(0), output_hidden_states=True)[
            2][align_layer][0, 1:-1]

        dot_prod = torch.matmul(out_src, out_tgt.transpose(-1, -2))

        softmax_srctgt = torch.nn.Softmax(dim=-1)(dot_prod)
        softmax_tgtsrc = torch.nn.Softmax(dim=-2)(dot_prod)

        softmax_inter = (softmax_srctgt > threshold) * \
            (softmax_tgtsrc > threshold)

    align_subwords = torch.nonzero(softmax_inter, as_tuple=False)

    align_words = set()

    for i, j in align_subwords:
        align_words.add((sub2word_map_src[i], sub2word_map_tgt[j]))

    return sent_src, sent_tgt, align_words



def get_word_mapping(source="", target="", model_name=""):
    """
    Get Word Aligned Mapping Words
    """
    sent_src, sent_tgt, align_words = get_alignment_mapping(
        source=source, target=target, model_name=model_name)

    result = []

    for i, j in sorted(align_words):
        result.append(f'bn:({sent_src[i]}) -> en:({sent_tgt[j]})')

    return result



def get_word_index_mapping(source="", target="", model_name=""):
    """
    Get Word Aligned Mapping Index
    """
    sent_src, sent_tgt, align_words = get_alignment_mapping(
        source=source, target=target, model_name=model_name)

    result = []

    for i, j in sorted(align_words):
        result.append(f'bn:({i}) -> en:({j})')

    return result


def get_alignments_table(
        source="", 
        target="", 
        model_name=""):
    """Get Spacy PoS Tags and return a Markdown table"""

    sent_src, sent_tgt, align_words = get_alignment_mapping(
        source=source, target=target, model_name=model_name
    )

    mapped_sent_src = []

    html_table = '''
                    <table>
                        <thead>
                            <th>Source</th>
                            <th>Target</th>
                        </thead>
                '''

    for i, j in sorted(align_words):
        punc = r"""!()-[]{}।;:'"\,<>./?@#$%^&*_~"""
        if sent_src[i] in punc or sent_tgt[j] in punc:
            mapped_sent_src.append(sent_src[i])

            html_table += f'''
                            <tbody>
                                <tr>
                                    <td> {sent_src[i]} </td>
                                    <td> {sent_tgt[j]} </td>
                                </tr>
                            '''
        else:
            mapped_sent_src.append(sent_src[i])

            html_table += f'''
                            <tr>
                                <td> {sent_src[i]} </td>
                                <td> {sent_tgt[j]} </td>
                            </tr>
                            '''

    unks = list(set(sent_src).difference(set(mapped_sent_src)))
    for word in unks:

        html_table += f'''
                        <tr>
                            <td> {word} </td>
                            <td> N/A </td>
                        </tr>                         
                    '''
        
    html_table += '''
                        </tbody>
                    </table>
                '''
    
    pos_accuracy = ((len(sent_src) - len(unks)) / len(sent_src))
    pos_accuracy = f"{pos_accuracy:0.2%}"

    return html_table, pos_accuracy