# Copyright (c) 2022-2023, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import configparser
import dataclasses
import json
import os
import pathlib
import time

import numpy as np
import torch
import torch.distributed as dist
import typing
from tqdm import tqdm

from omegaconf.omegaconf import OmegaConf, open_dict
from nemo.collections.nlp.data.glue_benchmark.glue_benchmark_dataset import (
    TextToTextGLUEDataset,
    TextToTextXNLIDataset,
)
from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer

from examples.pytorch.t5.utils.ft_encoder import FTT5EncoderWeight, FTT5Encoder
from examples.pytorch.t5.utils.ft_decoding import FTT5DecodingWeight, FTT5Decoding, FTT5


def _build_dataset(data_cfg, tokenizer):
    if data_cfg.task_name == 'xnli':
        dataset = TextToTextXNLIDataset(
            data_cfg.file_path,
            task_name=data_cfg.task_name,
            tokenizer=tokenizer,
            max_seq_length=data_cfg.max_seq_length,
            lang_list=data_cfg.eval_languages,
        )
    else:
        dataset = TextToTextGLUEDataset(
            data_cfg.file_path,
            task_name=data_cfg.task_name,
            tokenizer=tokenizer,
            max_seq_length=data_cfg.max_seq_length,
        )
    return dataset


def preds_and_labels_to_text(tokenizer, preds, labels):
    preds = preds.cpu().numpy().tolist()
    labels = labels.cpu().numpy().tolist()
    preds = [pred[0] for pred in preds]

    preds_text, labels_text = [], []
    for _, (pred, label) in enumerate(zip(preds, labels)):
        if tokenizer.eos_id in pred:
            idx = pred.index(tokenizer.eos_id)
            pred = pred[:idx]

        # Legacy sentencepiece detokenization still preserves special tokens which messes up exact string match.
        if hasattr(tokenizer, 'special_token_to_id'):
            pred = [id for id in pred if id not in tokenizer.special_token_to_id.values()]
            label = [id for id in label if id not in tokenizer.special_token_to_id.values()]
        pred = tokenizer.ids_to_text(pred)
        label = tokenizer.ids_to_text(label)
        preds_text.append(pred)
        labels_text.append(label)

    return preds_text, labels_text


def accuracy_score(pred, ref):
    assert len(pred) == len(ref)
    total = len(pred)
    correct = 0
    for p, r in zip(pred, ref):
        if p == r:
            correct += 1
        # else:
        #     print(f"[pred]: {p} [label]: {r}")
    accuracy = correct / total
    print(f"[accuracy]: {accuracy}")
    return accuracy


@dataclasses.dataclass
class Metric:
    acc: float


@dataclasses.dataclass
class RequestAndResult:
    model_answer: str
    target: str
    metrics: Metric


class InputToken:
    def __init__(self, input_ids, attention_mask):
        self.input_ids = input_ids
        self.attention_mask = attention_mask


class EncoderDecoderConfig:
    def __init__(self, d_model, vocab_size, num_heads, d_kv, d_ff, num_layers,
                 relative_attention_num_buckets_or_max_pos_seq_len,
                 adapter_inter_size, adapter_norm_position, decoder_start_token_id=0, decoder_end_token_id=1):
        self.d_model = d_model
        self.vocab_size = vocab_size
        self.num_heads = num_heads
        self.d_kv = d_kv
        self.d_ff = d_ff
        self.num_layers = num_layers
        self.relative_attention_num_buckets = relative_attention_num_buckets_or_max_pos_seq_len
        self.adapter_inter_size = adapter_inter_size
        self.adapter_norm_position = adapter_norm_position
        self.decoder_start_token_id = decoder_start_token_id
        self.decoder_end_token_id = decoder_end_token_id


data_type_mapping = {"fp32": 0, "fp16": 1, "bf16": 2}

def mnli_task(args_dict):
    torch.set_printoptions(precision=6)
    batch_size = args_dict['batch_size']
    beam_size = args_dict['beam_width']
    max_output_len = args_dict['max_output_len']
    beam_search_diversity_rate = args_dict['beam_search_diversity_rate']
    topk = args_dict['sampling_topk']
    topp = args_dict['sampling_topp']
    tensor_para_size = args_dict['tensor_para_size']
    pipeline_para_size = args_dict['pipeline_para_size']

    if args_dict['ckpt_path'] is None:
        raise Exception("Megatron T5 model needs to specify checkpoint path !")

    if dist.is_mpi_available():
        try:
            dist.init_process_group(backend='mpi')
            rank = dist.get_rank()
        except:
            rank = dist.get_rank()
    else:
        rank = 0

    assert dist.get_world_size() == tensor_para_size * pipeline_para_size

    ckpt_path = pathlib.Path(args_dict['ckpt_path'])
    ## read checkpoint config if exists
    ckpt_config = configparser.ConfigParser()

    vocab_path = ckpt_path / "vocab.txt"
    ckpt_config_path = ckpt_path / "config.ini"
    if ckpt_config_path.is_file():
        ckpt_config.read(ckpt_config_path)
        ## update structure config
        t5_with_bias = ckpt_config.getboolean('structure', 't5_with_bias')
        ## megatron with bias and use absolute position embedding
        ## relative position embedding -> 0, absolute position embedding -> 1
        position_embedding_type = 0 if ckpt_config.get('structure', 'position_embedding_type') == 'relative' else 1
        use_gated_activation = ckpt_config.getboolean('structure', 'use_gated_activation')
        weight_data_type = {"fp16": np.float16, "fp32": np.float32}[ckpt_config.get("encoder", "weight_data_type")]
        activation_type = ckpt_config.get('encoder', 'feed_forward_proj')
        assert ckpt_config.getint("encoder", "tensor_para_size") == tensor_para_size
    else:
        raise Exception("config file does exist with the ckpt !")

    if rank == 0:
        print("\n=============== Argument ===============")
        for key in args_dict:
            print("{}: {}".format(key, args_dict[key]))
        print("========================================")

    lib_path = args_dict['lib_path']

    ## build tokenizer, dataset, dataloader
    tokenizer_t5 = get_nmt_tokenizer(
                    library='megatron',
                    model_name='BertWordPieceCase',
                    tokenizer_model=None,
                    vocab_file=vocab_path.as_posix(),
                    merges_file=None,
                    legacy=False,
                    )

    assert tokenizer_t5.bos_id == ckpt_config.getint("decoder", "decoder_start_token_id")
    assert tokenizer_t5.eos_id == ckpt_config.getint("decoder", "eos_token_id")

    token_params = {
        tokenizer_t5.bos_token: tokenizer_t5.bos_id,
        tokenizer_t5.eos_token: tokenizer_t5.eos_id,
        tokenizer_t5.pad_token: tokenizer_t5.pad_id,
    }
    print(f"tokenizer special tokens: {token_params}")
    mnli_cfg = OmegaConf.create({
        "file_path": args_dict['data_path'],
        "task_name": "mnli",
        "max_seq_length": 512
    })

    mnli_dataset = _build_dataset(mnli_cfg, tokenizer_t5)

    data_loader = torch.utils.data.DataLoader(
                mnli_dataset,
                collate_fn=mnli_dataset.collate_fn,
                batch_size=batch_size,
                num_workers=8,
                pin_memory=True,
                drop_last=True)
    
    q_scaling = 1.0

    encoder_config = EncoderDecoderConfig(ckpt_config.getint('encoder', 'd_model'),
                                          ckpt_config.getint('encoder', 'vocab_size'),
                                          ckpt_config.getint('encoder', 'num_heads'),
                                          ckpt_config.getint('encoder', 'd_kv'),
                                          ckpt_config.getint('encoder', 'd_ff'),
                                          ckpt_config.getint('encoder', 'num_layers'),
                                          ckpt_config.getint('encoder',
                                                             'relative_attention_num_buckets_or_max_pos_seq_len'),
                                          ckpt_config.getint('encoder', 'adapter_inter_size', fallback=0),
                                          ckpt_config.get('encoder', 'adapter_norm_position', fallback="pre")
                                          )

    decoder_config = EncoderDecoderConfig(ckpt_config.getint('decoder', 'd_model'),
                                          ckpt_config.getint('decoder', 'vocab_size'),
                                          ckpt_config.getint('decoder', 'num_heads'),
                                          ckpt_config.getint('decoder', 'd_kv'),
                                          ckpt_config.getint('decoder', 'd_ff'),
                                          ckpt_config.getint('decoder', 'num_layers'),
                                          ckpt_config.getint('decoder',
                                                             'relative_attention_num_buckets_or_max_pos_seq_len'),
                                          ckpt_config.getint('decoder', 'adapter_inter_size', fallback=0),
                                          ckpt_config.get('decoder', 'adapter_norm_position', fallback="pre"),
                                          tokenizer_t5.bos_id,
                                          tokenizer_t5.eos_id
                                          )

    ## run gemm test
    if os.path.isfile("gemm_config.in") and rank == 0:
        cmd = f"rm gemm_config.in"
        print(f"Run {cmd}")
        os.system(cmd)
    if rank == 0:
        data_type = data_type_mapping[args_dict['data_type']]
        cmd = f"./bin/t5_gemm {batch_size // pipeline_para_size} {beam_size} {128} " \
            f"{encoder_config.d_model} {encoder_config.num_heads} {encoder_config.d_kv} {encoder_config.d_ff} " \
            f"{decoder_config.d_model} {decoder_config.num_heads} {decoder_config.d_kv} {decoder_config.d_ff} " \
            f"{decoder_config.vocab_size} {data_type} {tensor_para_size} 1 > .tmp_gemm.log"
        print(f"Run gemm test: {cmd}")
        os.system(cmd)

    dist.barrier()

    ft_encoder_weight = FTT5EncoderWeight(
        encoder_config,
        tensor_para_size,
        pipeline_para_size,
        t5_with_bias=t5_with_bias,
        use_gated_activation=use_gated_activation,
        position_embedding_type=position_embedding_type,
        weight_data_type=weight_data_type,
    )
    ft_decoding_weight = FTT5DecodingWeight(
        decoder_config,
        tensor_para_size,
        pipeline_para_size,
        t5_with_bias=t5_with_bias,
        use_gated_activation=use_gated_activation,
        position_embedding_type=position_embedding_type,
        weight_data_type=weight_data_type,
    )

    ft_encoder_weight.load_from_bin(ckpt_path.as_posix(), "Megatron")
    ft_decoding_weight.load_from_bin(ckpt_path.as_posix(), "Megatron")

    if args_dict['data_type'] == 'fp16':
        ft_encoder_weight.to_half()
        ft_decoding_weight.to_half()
    elif args_dict['data_type'] == 'fp32':
        ft_encoder_weight.to_single()
        ft_decoding_weight.to_single()
    elif args_dict['data_type'] == 'bf16':
        ft_encoder_weight.to_bfloat16()
        ft_decoding_weight.to_bfloat16()

    remove_padding = True if batch_size > 32 else False
    ft_encoder = FTT5Encoder(ft_encoder_weight.w, lib_path, encoder_config.num_heads,
                             encoder_config.d_kv, encoder_config.d_ff,
                             encoder_config.d_model, remove_padding, encoder_config.num_layers,
                             encoder_config.relative_attention_num_buckets,
                             0,  # num_experts
                             [],  # moe_layer_index
                             128, False,
                             q_scaling, tensor_para_size, pipeline_para_size, t5_with_bias,
                             position_embedding_type, 0, activation_type,
                             adapter_inter_size=encoder_config.adapter_inter_size,
                             adapter_norm_position=encoder_config.adapter_norm_position)
    ft_decoding = FTT5Decoding(ft_decoding_weight.w, lib_path,
                               decoder_config.num_heads, decoder_config.d_kv,
                               decoder_config.d_ff, encoder_config.d_model,
                               decoder_config.d_model, decoder_config.num_layers,
                               decoder_config.decoder_start_token_id, decoder_config.decoder_end_token_id,
                               decoder_config.vocab_size,
                               q_scaling,
                               decoder_config.relative_attention_num_buckets,
                               0,  # num_experts
                               [],  # moe_layer_index
                               max_distance=128,
                               tensor_para_size=tensor_para_size, pipeline_para_size=pipeline_para_size,
                               moe_k=0, activation_type=activation_type,
                               t5_with_bias=t5_with_bias, position_embedding_type=position_embedding_type,
                               adapter_inter_size=decoder_config.adapter_inter_size,
                               adapter_norm_position=decoder_config.adapter_norm_position)

    ft_t5 = FTT5(ft_encoder, ft_decoding)

    preds_list = []
    labels_list = []
    results_list = []
    start = time.time()
    for idx, batch in tqdm(enumerate(data_loader)):
        input_token = InputToken(batch['text_enc'], batch['enc_mask'])
        ft_decoding_outputs, ft_decoding_seq_lens = ft_t5(input_token,
                                                          None,
                                                          beam_size,
                                                          max_output_len,
                                                          topk,
                                                          topp,
                                                          beam_search_diversity_rate=beam_search_diversity_rate,
                                                          is_return_output_log_probs=args_dict["return_output_log_probs"],
                                                          is_return_cum_log_probs=args_dict["return_cum_log_probs"])
        
        preds, labels = preds_and_labels_to_text(tokenizer_t5, torch.IntTensor(ft_decoding_outputs), batch['labels'])

        labels_list += labels
        preds_list += preds
        results_list.extend([
            RequestAndResult(
                model_answer=pred,
                target=label,
                metrics=Metric(acc=pred == label)
            )
            for pred, label in zip(preds, labels)
        ])

    end = time.time()
    if rank == 0:
        print(f"\n[Elapsed Time]: {end - start} seconds")

        accuracy = accuracy_score(preds_list, labels_list)
        output_path = args_dict.get("output_path")
        if output_path is not None and rank == 0:
            output_path = pathlib.Path(output_path)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            with output_path.open("w") as output_file:
                results = {
                    "results": {
                        "mnli": {
                            "acc": accuracy
                        }
                    },
                    "output": {
                        "mnli": [
                            dataclasses.asdict(r) for r in results_list
                        ]
                    }
                }
                json.dump(results, output_file)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-batch', '--batch_size', type=int, default=1, metavar='NUMBER',
                        help='batch size (default: 1)')
    parser.add_argument('-beam', '--beam_width', type=int, default=4, metavar='NUMBER',
                        help='beam width (default: 4)')
    parser.add_argument('-s', '--max_output_len', type=int, default=10, metavar='NUMBER',
                        help='max output length (default: 10)')
    parser.add_argument('-diversity_rate', '--beam_search_diversity_rate', type=float, default=0.0, metavar='NUMBER',
                        help='deviersity rate of beam search. default is 0. When diversity rate = 0, it is equivalent to the naive beam search.')
    parser.add_argument('-topk', '--sampling_topk', type=int, default=1, metavar='NUMBER',
                        help='Candidate (k) value of top k sampling in decoding. Default is 1.')
    parser.add_argument('-topp', '--sampling_topp', type=float, default=0.0, metavar='NUMBER',
                        help='Probability (p) value of top p sampling in decoding. Default is 0.0. ')
    parser.add_argument('-d', '--data_type', type=str, default="fp32", metavar='STRING',
                        help='data type (default: fp32)', choices=['fp32', 'fp16', 'bf16'])
    parser.add_argument('-lib_path', '--lib_path', type=str, default="/workspace/FasterTransformer/build/lib/libth_transformer.so", metavar='STRING',
                        help='the path of FasterTransformer pytorch t5 op library.')
    parser.add_argument('-data_path', '--data_path', type=str, required=True, help="the MNLI task data path")
    parser.add_argument('-tensor_para_size', '--tensor_para_size', type=int, default=1, metavar='NUMBER',
                    help='size of tensor parallelism (default: 1)')
    parser.add_argument('-pipeline_para_size', '--pipeline_para_size', type=int, default=1, metavar='NUMBER',
                        help='size of pipeline parallelism (default: 1)')
    # assume checkpoint config is also in the same path
    parser.add_argument('--ckpt_path', type=str, help='path to the checkpoint file.')
    parser.add_argument('--output_path', help='path to results file with calculated metrics.')
    parser.add_argument('--return_output_log_probs', action='store_true',
                        help='Return the log probability of generated tokens.')
    parser.add_argument('--return_cum_log_probs', action='store_true',
                        help='Return the cumulative log probability of generated tokens.')
    args = parser.parse_args()

    mnli_task(vars(args))