"""
This script provides an example to wrap TencentPretrain for speech-to-text inference.
"""
import sys
import os
import tqdm
import argparse
import math
import torch
import torch.nn.functional as F

tencentpretrain_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.append(tencentpretrain_dir)

from tencentpretrain.utils.constants import *
from tencentpretrain.utils import *
from tencentpretrain.utils.config import load_hyperparam
from tencentpretrain.model_loader import load_model
from tencentpretrain.opts import infer_opts, tokenizer_opts
from finetune.run_speech2text import Speech2text, read_dataset
from inference.run_classifier_infer import batch_loader


def main():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    infer_opts(parser)

    tokenizer_opts(parser)

    parser.add_argument("--beam_width", type=int, default=10,
                        help="Beam width.")
    parser.add_argument("--tgt_seq_length", type=int, default=100,
                        help="inference step.")
    args = parser.parse_args()

    # Load the hyperparameters from the config file.
    args = load_hyperparam(args)

    # Build tokenizer.
    args.tokenizer = str2tokenizer[args.tokenizer](args)

    # Build s2t model.
    model = Speech2text(args)
    model = load_model(model, args.load_model_path)

    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(args.device)

    if torch.cuda.device_count() > 1:
        print("{} GPUs are available. Let's use them.".format(torch.cuda.device_count()))
        model = torch.nn.DataParallel(model)

    dataset = read_dataset(args, args.test_path)

    src_audio = torch.stack([sample[0] for sample in dataset], dim=0)
    seg_audio = torch.LongTensor([sample[3] for sample in dataset])

    batch_size = args.batch_size
    beam_width=args.beam_width
    instances_num = src_audio.size()[0]
    tgt_seq_length = args.tgt_seq_length

    PAD_ID = args.tokenizer.convert_tokens_to_ids([PAD_TOKEN])
    SEP_ID = args.tokenizer.convert_tokens_to_ids([SEP_TOKEN])
    CLS_ID = args.tokenizer.convert_tokens_to_ids([CLS_TOKEN])
    print("The number of prediction instances: ", instances_num)

    model.eval()
    with open(args.prediction_path, mode="w", encoding="utf-8") as f:
        for i, (src_batch, seg_batch) in tqdm.tqdm(enumerate(batch_loader(batch_size, src_audio, seg_audio))):
            src_batch = src_batch.to(args.device)
            seg_batch = seg_batch.to(args.device)

            seq_length = seg_batch.sum(dim=-1).max()
            src_batch = src_batch[:,:seq_length * 4,:]
            seg_batch = seg_batch[:,:seq_length]

            tgt_in_batch = torch.zeros(src_batch.size()[0], 1, dtype = torch.long, device = args.device)
            current_batch_size=tgt_in_batch.size()[0]
            for j in range(current_batch_size):
                tgt_in_batch[j][0] = torch.LongTensor(SEP_ID) #torch.LongTensor(CLS_ID) for tencentpretrain model

            with torch.no_grad():
                memory_bank, emb = model(src_batch, None, seg_batch, None, only_use_encoder=True)

            step = 0
            scores = torch.zeros(current_batch_size, beam_width, tgt_seq_length)
            tokens = torch.zeros(current_batch_size, beam_width, tgt_seq_length+1, dtype = torch.long)
            tokens[:,:,0] = torch.LongTensor(args.tokenizer.convert_tokens_to_ids([SEP_TOKEN])) #2
            emb = emb.repeat(1, beam_width, 1).reshape(current_batch_size * beam_width, -1, int(args.conv_channels[-1] / 2)) #same batch nearby
            memory_bank = memory_bank.repeat(1, beam_width, 1).reshape(current_batch_size * beam_width, -1, args.emb_size) 
            tgt_in_batch = tgt_in_batch.repeat(beam_width, 1)
            while step < tgt_seq_length and step < seq_length:
                with torch.no_grad():
                    outputs = model(emb, (tgt_in_batch, None, None), None, None, memory_bank=memory_bank)

                vocab_size = outputs.shape[-1]
                log_prob = F.log_softmax(outputs[:, -1, :], dim=-1) #(B*beam_size) * 1 * vocab_size
                log_prob = log_prob.squeeze() #(B*beam_size) * vocab_size

                log_prob[:,PAD_ID] = -math.inf # do not select pad
                if step == 0:
                    log_prob[:,SEP_ID] = -math.inf # </s>

                log_prob_beam = log_prob.reshape(current_batch_size, beam_width, -1) # B * beam * vocab_size

                if step == 0:
                    log_prob_beam = log_prob_beam[:, ::beam_width, :].contiguous().to(scores.device)
                else:
                    log_prob_beam = log_prob_beam.to(scores.device) + scores[:,:, step-1].unsqueeze(-1)
                
                top_prediction_prob, top_prediction_indices = torch.topk(log_prob_beam.view(current_batch_size, -1), k=beam_width)
                beams_buf = torch.div(top_prediction_indices, vocab_size).trunc().long()
                beams_buf = beams_buf + torch.arange(current_batch_size).repeat(beam_width).reshape(beam_width,-1).transpose(0,1) * beam_width
                top_prediction_indices = top_prediction_indices.fmod(vocab_size)
                
                scores[:, :, step] = top_prediction_prob
                tokens[:, :, step+1] = top_prediction_indices

                if step > 0 and current_batch_size == 1:
                    tokens[:, :, :step+1] = torch.index_select(tokens, dim=1, index=beams_buf.squeeze())[:, :, :step+1]
                elif step > 0:
                    tokens[:, :, step+1] = torch.index_select(tokens.reshape(-1,tokens.shape[2]), dim=0, index=beams_buf.reshape(-1)).reshape(current_batch_size, -1, tokens.shape[2])[:, :, step+1]
                tgt_in_batch = tokens[:, :, :step+2].view(current_batch_size * beam_width, -1)
                tgt_in_batch = tgt_in_batch.long().to(emb.device)

                step = step + 1
            for i in range(current_batch_size):
                for j in range(1):
                    res = "".join([args.tokenizer.inv_vocab[token_id.item()] for token_id in tokens[i,j,:]])
                    res = res.split(SEP_TOKEN)[1].split(CLS_TOKEN)[0] # res.split(CLS_TOKEN)[1].split(SEP_TOKEN)[0] for tencentpretrain model
                    res = res.replace('▁',' ')

                    f.write(res)
                    f.write("\n")


if __name__ == "__main__":
    main()