#!/usr/bin/env python

import os
import sys
import argparse
import numpy as np
import torch
from torch.utils.data import DataLoader, SequentialSampler, TensorDataset
from tqdm import tqdm
from transformers import BertModel, BertTokenizer
from transformers import XLMRobertaTokenizer, XLMRobertaModel


def convert_sent_to_input(sents, tokenizer, max_seq_length):
    input_ids = []
    mask = []
    for sent in sents:
        ids = tokenizer.convert_tokens_to_ids(sent)
        mask.append([1] * (len(ids) + 2) + [0] * (max_seq_length - len(ids)))
        input_ids.append([tokenizer.cls_token_id] + ids + [tokenizer.eos_token_id] + [tokenizer.pad_token_id] * (max_seq_length - len(ids)))
        #input_ids.append(tokenizer.convert_tokens_to_ids(['[CLS]']) + ids + tokenizer.convert_tokens_to_ids(['[SEP]']) + tokenizer.convert_tokens_to_ids(['[PAD]']) * (
        #            max_seq_length - len(ids)))
    return torch.tensor(input_ids, dtype=torch.long), torch.tensor(mask, dtype=torch.long)


def convert_words_to_bpe(para, tokenizer):
    
    bpe_para, bpe_table = [], []

    for s in para:
        
        src_sent, tgt_sent = s[0], s[1]
        src_bpe_table, tgt_bpe_table = [], []
        src_sent_bpe, tgt_sent_bpe = [], []
        
        for word in src_sent:
            token = tokenizer.tokenize(word)
            word2bpe_map = []
            for i in range(len(token)):
                word2bpe_map.append(len(src_sent_bpe)+i)
            src_sent_bpe.extend(token)
            src_bpe_table.append(word2bpe_map)

        for word in tgt_sent:
            token = tokenizer.tokenize(word)
            word2bpe_map = []
            for i in range(len(token)):
                word2bpe_map.append(len(tgt_sent_bpe)+i)
            tgt_sent_bpe.extend(token)
            tgt_bpe_table.append(word2bpe_map)
            
        bpe_para.append([src_sent_bpe, tgt_sent_bpe])
        bpe_table.append([src_bpe_table, tgt_bpe_table])
        
    return bpe_para, bpe_table


def read_parallel(file_name, reverse):
    para = list()
    with open(file_name, "rb") as f:
        for line in f:
            src, tgt = line.decode().split("|||")
            para.append((tgt.split(), src.split()) if reverse else (src.split(), tgt.split()))
    return para


def read_align(file_name, reverse):
    aligned = list()
    with open(file_name, "rb") as f:
        for line in f:
            pairs = list()
            for pair in line.split():
                pair = tuple(map(int, reversed(pair.split(b"-")) if reverse else pair.split(b"-")))
                pairs.append(pair)
            aligned.append(pairs)
    return sum(len(pairs) for pairs in aligned), aligned


def main(arguments):

    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    
    parser.add_argument('--para_file', type=str, 
                        help="Input parallel file")
    parser.add_argument('--align_file', type=str, 
                        help="Input alignment file")
    parser.add_argument('--reverse', type=int, default=0,
                        help="Wheter para and align is reversed")

    parser.add_argument('--align_mode', choices=['word_align', 'sent_avg'],
                        default='word_align', help=" \
                        word_align uses the alignment from align_file \
                        sent_avg uses the averaged embeddign for each sentence")

    parser.add_argument("--bert_model", default=None, type=str, 
                    help="Bert pre-trained model selected in the list: bert-base-uncased, "
                    "bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased, "
                    "bert-base-multilingual-cased, bert-base-chinese.")
    parser.add_argument("--do_lower_case", action='store_true',
                        help="Set this flag if you are using an uncased model.")
    parser.add_argument("--cache_dir", default="./", type=str,
                        help="Where do you want to store the pre-trained models downloaded from s3")
    parser.add_argument('--layer', type=int, default=11,
                        help="layer of bert outputs to be used")

    parser.add_argument('--batch_size', type=int, default=16)
    parser.add_argument('--max_seq_length', type=int, default=175)

    parser.add_argument('--output_file', help="Output file path", \
                        default='output')

    args = parser.parse_args()

    para = read_parallel(args.para_file, reverse=args.reverse)

    align_cnt, align_pairs = read_align(args.align_file, reverse=args.reverse)
    print("Read {} aligns".format(align_cnt))

    tokenizer = XLMRobertaTokenizer.from_pretrained(args.bert_model, do_lower_case=args.do_lower_case)
    #tokenizer = BertTokenizer.from_pretrained(args.bert_model, do_lower_case=args.do_lower_case)
    
    bpe_para, bpe_table = convert_words_to_bpe(para, tokenizer)

    # filter long/empty sentences
    fltr_para = []
    fltr_align_pairs = []
    fltr_bpe_table = []
    align_cnt = 0

    for cnt, [src, trg] in enumerate(bpe_para):
        if len(src) > args.max_seq_length or len(trg) > args.max_seq_length \
            or len(src) == 0 or len(trg) == 0:
            continue
        fltr_para.append([src, trg])
        fltr_align_pairs.append(align_pairs[cnt])
        fltr_bpe_table.append(bpe_table[cnt])
        align_cnt += len(align_pairs[cnt])

    para = fltr_para
    align_pairs = fltr_align_pairs
    print("After filtering {} parallel sentences remains".format(len(para)))

    # print the n-th sentence and alignment
    n = 0
    print("{} ||| {}".format(" ".join(para[n][0]), " ".join(para[n][1])))
    for a in align_pairs[n]:
        print(' '.join([para[n][0][i] for i in fltr_bpe_table[n][0][a[0]]]) + ' : ' + ' '.join([para[n][1][i] for i in fltr_bpe_table[n][1][a[1]]]))

    # bert
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    en_model = XLMRobertaModel.from_pretrained(args.bert_model, cache_dir=args.cache_dir, output_hidden_states=True)
    #en_model = BertModel.from_pretrained(args.bert_model, cache_dir=args.cache_dir, output_hidden_states=True)
    en_model.to(device)

    src_sents = [s[0] for s in para]
    tgt_sents = [s[1] for s in para]
    src_input, src_mask = convert_sent_to_input(src_sents, tokenizer, args.max_seq_length)
    tgt_input, tgt_mask = convert_sent_to_input(tgt_sents, tokenizer, args.max_seq_length)

    src_data = TensorDataset(src_input, src_mask)
    src_sampler = SequentialSampler(src_data)
    src_dataloader = DataLoader(src_data, sampler=src_sampler, batch_size=args.batch_size)

    tgt_data = TensorDataset(tgt_input, tgt_mask)
    tgt_sampler = SequentialSampler(tgt_data)
    tgt_dataloader = DataLoader(tgt_data, sampler=tgt_sampler,batch_size=args.batch_size)
    print(tokenizer.convert_ids_to_tokens(src_input[0]))
    en_model.eval()

    src_embed = []
    tgt_embed = []

    for step, batch in enumerate(tqdm(src_dataloader, desc="Iteration")):
        input_ids, input_mask = batch
        input_ids = input_ids.to(device)
        input_mask = input_mask.to(device)
        
        with torch.no_grad():
            all_encoder_layers = en_model(input_ids=input_ids, attention_mask=input_mask)
            all_encoder_layers = all_encoder_layers[2]

        src_embed.append(all_encoder_layers[args.layer][:,1:].detach().to('cpu').numpy())#remove CLS

    for step, batch in enumerate(tqdm(tgt_dataloader, desc="Iteration")):
        input_ids, input_mask = batch
        input_ids = input_ids.to(device)
        input_mask = input_mask.to(device)
        
        with torch.no_grad():
            all_encoder_layers = en_model(input_ids, attention_mask=input_mask)[2]
            
        tgt_embed.append(all_encoder_layers[args.layer][:,1:].detach().to('cpu').numpy())

    src_embed = np.concatenate(src_embed)
    tgt_embed = np.concatenate(tgt_embed)
    print(src_embed.shape)
    print(tgt_embed.shape)

    feature_size = src_embed.shape[2]
    
    if args.align_mode == 'word_align':
        final_res = [np.zeros((align_cnt, feature_size)),
                     np.zeros((align_cnt, feature_size))]
        cnt = 0
    elif args.align_mode == 'sent_avg':
        sent_cnt = len(src_embed)
        final_res = [np.zeros((sent_cnt, feature_size)),
                     np.zeros((sent_cnt, feature_size))]

    for i, pairs in enumerate(align_pairs):
        try:
            if args.align_mode == 'word_align':
                for a in pairs:
                    if len(fltr_bpe_table[i][0][a[0]]) > 0 and len(fltr_bpe_table[i][1][a[1]]) > 0: # token alignment (0,0)
                        src_word_avg_embed = np.zeros((1, feature_size))
                        
                        for j in fltr_bpe_table[i][0][a[0]]:
                            src_word_avg_embed += src_embed[i][j,:]
                        final_res[0][cnt,:] = src_word_avg_embed / len(fltr_bpe_table[i][0][a[0]])

                        tgt_word_avg_embed = np.zeros((1, feature_size))
                        for j in fltr_bpe_table[i][1][a[1]]:
                            tgt_word_avg_embed += tgt_embed[i][j,:]

                        final_res[1][cnt,:] = tgt_word_avg_embed / len(fltr_bpe_table[i][1][a[1]])
                        cnt += 1
            elif args.align_mode == 'sent_avg':
                final_res[0][i,:] = np.mean(src_embed[i], axis=0)
                final_res[1][i,:] = np.mean(tgt_embed[i], axis=0)
        except IndexError:
            continue

    suffix = ['.src', '.trg']
    for j in [0, 1]:
        np.savetxt(args.output_file+suffix[j], final_res[j], delimiter=' ', fmt='%0.6f')

if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))
