from __future__ import absolute_import, division, print_function

import argparse
import glob
import logging
import os
import random
import json
import numpy as np
from collections import defaultdict
import torch
os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
torch.cuda.empty_cache()
try:
    from torch.utils.tensorboard import SummaryWriter
except:
    from tensorboardX import SummaryWriter

from tqdm import tqdm, trange
from transformers import (WEIGHTS_NAME, BertConfig, BertTokenizer)
import torch

os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
torch.cuda.empty_cache()
from transformers import AdamW, get_linear_schedule_with_warmup
from models import TokenClsModel, SentClsModel
import threading
from queue import Queue
import task_utils
import pickle, time
from utils import DataProcessor_Data, CustomDataset, set_seed, read_categories
logger = logging.getLogger(__name__)

ALL_MODELS = sum((tuple(conf.pretrained_config_archive_map.keys()) for conf in (BertConfig, )), ())

MODEL_CLASSES = {
    'token-cls': (BertConfig, TokenClsModel, BertTokenizer),
    'sent-cls': (BertConfig, SentClsModel, BertTokenizer),
}

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

def create_dataset(args, model_name, part):
    input_file = part + args.data + '_' + model_name + '.pkl'
    input_file = os.path.join(args.data_dir, input_file)
    dataset = pickle.load(open(input_file, 'rb'))
    return dataset

class DataProcessor(threading.Thread):
    """docstring for DataProcessor"""
    def __init__(self, args, dataset, tokenizer, batch_processor, is_eval=False, max_queue_size=5):
        super(DataProcessor, self).__init__()
        self.args = args
        self.stream = torch.cuda.Stream()
        self.queue = Queue(maxsize=max_queue_size)
        self.device = args.device
        self.to_gpu_keys = args.to_gpu_keys
        self.dataset = dataset
        self.tokenizer = tokenizer
        self.batch_processor = batch_processor
        self.is_eval = is_eval

    def make_features(self, examples):
        max_length = -1
        for example in examples:
            max_length = max(max_length, len(example['input_ids']))
        max_length = min(max_length, self.args.max_seq_length)

        batch = {}
        for key in ['input_ids', 'attention_mask', 'dep_att_mask', 'labels']:
            batch[key] = []

        if len(examples[0]['labels']) > 1:
            batch['loss_mask'] = []

        for example in examples:
            input_ids = example['input_ids'][:max_length]
            dep_att_mask = example['dep_att_mask'][:max_length]
            tokens = example['tokens'][:max_length]
            seq_length = len(input_ids)

            assert seq_length == len(dep_att_mask)

            padding_length = max_length - seq_length
            batch['input_ids'].append(input_ids + ([0] * padding_length))

            batch['attention_mask'].append(([1] * seq_length) + ([0] * padding_length))

            for i in range(seq_length):
                dep_att_mask[i] = dep_att_mask[i][:max_length]
                if padding_length > 0:
                    dep_att_mask[i] = dep_att_mask[i] + ([0] * padding_length)

            for i in range(padding_length):
                dep_att_mask.append([0] * max_length)

            batch['dep_att_mask'].append(dep_att_mask)

            if len(example['labels']) > 1:
                labels = example['labels'][:max_length] + ([0] * padding_length)
                loss_mask = example['loss_mask'][:max_length] + ([0] * padding_length)
                batch['labels'].append(labels)
                batch['loss_mask'].append(loss_mask)
            else:
                batch['labels'].append(example['labels'][0])

        for key in batch:
            if key != 'dep_att_mask':
                batch[key] = torch.tensor(batch[key], dtype=torch.long)
            else:
                batch[key] = torch.tensor(batch[key], dtype=torch.float)

        batch = self.batch_processor(batch, self.tokenizer)

        for key in self.to_gpu_keys:
            if key in batch:
                batch[key] = batch[key].to(device=self.device)
        return batch

    def run(self):
        data_size = len(self.dataset)
        if not self.is_eval:
            random.shuffle(self.dataset)
            bs = self.args.train_batch_size
        else:
            bs = self.args.eval_batch_size

        with torch.cuda.stream(self.stream):
            for i in range(0, data_size, bs):
                self.queue.put(self.make_features(self.dataset[i:min(i+bs,data_size)]))

def train(args, model, tokenizer, batch_processor):
    """ Train the model """
    print('Training\n----------------------------------------')
    args.train_batch_size = args.per_gpu_train_batch_size
    if args.local_rank == -1:
        train_dataset = create_dataset(args, args.model_name_or_path, 'train')
    else:
        total_dataset = create_dataset(args, args.model_name_or_path, 'train')
        start_position = 0
        width = torch.distributed.get_world_size()
        train_dataset = []
        while start_position + width <= len(total_dataset):
            train_dataset.append(total_dataset[start_position+args.local_rank])
            start_position += width

    if args.max_steps > 0:
        t_total = args.max_steps
        args.num_train_epochs = args.max_steps // (len(train_dataset) // args.train_batch_size // args.gradient_accumulation_steps) + 1
    else:
        t_total = len(train_dataset) // args.train_batch_size // args.gradient_accumulation_steps * args.num_train_epochs

    no_decay = ['bias', 'LayerNorm.weight']
    need_optimized_parameters = [(n, p) for n, p in model.named_parameters() if p.requires_grad]
    optimizer_grouped_parameters = [
        {'params': [p for n, p in need_optimized_parameters if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
        {'params': [p for n, p in need_optimized_parameters if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]

    optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total)
    if args.fp16:
        try:
            from apex import amp
            amp.register_half_function(torch, 'einsum')
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)

    if args.local_rank != -1:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
                                                          output_device=args.local_rank,
                                                          find_unused_parameters=True)

    global_step = 0
    tr_loss, logging_loss = 0.0, 0.0
    model.zero_grad()
    train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0])
    set_seed(args)

    for _ in train_iterator:
        data_helper = DataProcessor(args, train_dataset, tokenizer, batch_processor, False)
        data_helper.start()
        time.sleep(2)
        step = 0
        while not data_helper.queue.empty():
            batch = data_helper.queue.get()
            model.train()
            loss = model(batch)[0]
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            if args.fp16:
                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                loss.backward()

            tr_loss += loss.item()
            if (step + 1) % args.gradient_accumulation_steps == 0:
                if args.fp16:
                    torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                else:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

                optimizer.step()
                scheduler.step()  # Update learning rate schedule
                model.zero_grad()
                global_step += 1

                if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
                    logs = {}
                    loss_scalar = (tr_loss - logging_loss) / args.logging_steps
                    learning_rate_scalar = scheduler.get_lr()[0]
                    logs['learning_rate'] = learning_rate_scalar
                    logs['loss'] = loss_scalar
                    logging_loss = tr_loss

                if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
                    output_dir = os.path.join(args.output_dir, 'saved_ckpt-{}'.format(global_step))
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                    model_to_save = model.module if hasattr(model,
                                                            'module') else model  # Take care of distributed/parallel training
                    model_to_save.save_pretrained(output_dir)
                    torch.save(args, os.path.join(output_dir, 'training_args.bin'))
            if args.max_steps > 0 and global_step > args.max_steps:
                break
            step += 1
        if args.max_steps > 0 and global_step > args.max_steps:
            train_iterator.close()
            break
    return global_step, tr_loss / global_step

def evaluate(args, model, tokenizer, batch_processor, prefix=""):
    eval_dataset = create_dataset(args, args.model_name_or_path, 'dev')
    args.eval_batch_size = args.per_gpu_eval_batch_size

    eval_loss = 0.0
    nb_eval_steps = 0

    eval_data = {
        'labels': [],
        'preds': [],
        'loss_mask': []
    }
    data_helper = DataProcessor(args, eval_dataset, tokenizer, batch_processor, True)
    data_helper.start()
    time.sleep(2)

    while not data_helper.queue.empty():
        batch = data_helper.queue.get()
        model.eval()
        with torch.no_grad():
            outputs = model(batch['input_ids'], batch['attention_mask'], batch['dep_att_mask'], labels= batch['labels'])
            tmp_eval_loss, logits = outputs[:2]
            eval_loss += tmp_eval_loss.mean().item()
        nb_eval_steps += 1
        preds = logits.detach().cpu().numpy()
        preds = np.argmax(preds, axis=-1).tolist()
        labels = batch['labels'].detach().cpu().numpy().tolist()
        eval_data['labels'] += labels
        eval_data['preds'] += preds
        if 'loss_mask' in batch:
            eval_data['loss_mask'] += batch['loss_mask'].detach().cpu().numpy().tolist()

    results = task_utils.evaluate(args.task, eval_data)
    for key in sorted(results.keys()):
        logger.info("  %s = %s", key, str(results[key]))
    return results

def read_categories(file_name):
    dict_cat = defaultdict(list)
    file_cat = open(file_name, 'r', encoding='utf8')
    for fil in file_cat.readlines():
        fil = fil.strip('\n').split('\t')
        dict_cat[fil[1]].append(fil[0])
    return dict_cat

def evaluate_category(args, model, tokenizer, batch_processor, prefix=''):
    dict_cat = read_categories('file_categories.txt')
    results_list = defaultdict(defaultdict)
    for key in dict_cat.keys():
        eval_dataset = create_dataset(args, args.model_name_or_path, 'test', cat=key)
        args.eval_batch_size = args.per_gpu_eval_batch_size

        logger.info("***** Running evaluation {} *****".format(prefix))
        logger.info("  Num examples = %d", len(eval_dataset))
        logger.info("  Batch size = %d", args.eval_batch_size)

        eval_loss = 0.0
        nb_eval_steps = 0

        eval_data = {
        'labels': [],
        'preds': [],
        'loss_mask': []
    }
        data_helper = DataProcessor(args, eval_dataset, tokenizer, batch_processor, True)
        data_helper.start()
        time.sleep(2)

        while not data_helper.queue.empty():
            batch = data_helper.queue.get()
            model.eval()
            with torch.no_grad():
                input_ids = batch['input_ids']
                attention_mask = batch['attention_mask']
                labels = batch['labels']
                dep_att_mask = batch['dep_att_mask']
                loss_mask = batch['loss_mask']

                outputs = model(input_ids, attention_mask, dep_att_mask, labels=labels)
                tmp_eval_loss, logits = outputs[:2]
                eval_loss += tmp_eval_loss.mean().item()
            nb_eval_steps += 1
            preds = logits.detach().cpu().numpy()
            preds = np.argmax(preds, axis=-1).tolist()
            labels = batch['labels'].detach().cpu().numpy().tolist()
            eval_data['labels'] += labels
            eval_data['preds']  += preds
            if 'loss_mask' in batch:
                eval_data['loss_mask'] += batch['loss_mask'].detach().cpu().numpy().tolist()

        results_list[key] = task_utils.evaluate(args.task, eval_data)
        for key in sorted(results.keys()):
            logger.info("  %s = %s", key, str(results[key]))
    return results_list

def evaluate_test(args, model, tokenizer, batch_processor, prefix=""):
    eval_dataset = create_dataset(args, args.model_name_or_path, 'test')

    args.eval_batch_size = 1

    eval_loss = 0.0
    nb_eval_steps = 0

    eval_data = {
        'labels': [],
        'preds': [],
        'loss_mask': [],
	 'tokens': []
    }
    data_helper = DataProcessor(args, eval_dataset, tokenizer, batch_processor, True)
    data_helper.start()
    time.sleep(2)
    sentences = []

    out_file = open('outputs/predictions.txt', 'w', encoding='utf8')

    while not data_helper.queue.empty():
        batch = data_helper.queue.get()
        model.eval()
        with torch.no_grad():
            input_ids = batch['input_ids']
            attention_mask = batch['attention_mask']
            labels = batch['labels']
            dep_att_mask = batch['dep_att_mask']
            loss_mask = batch['loss_mask']
            tokens = tokenizer.decode(input_ids)

            outputs = model(input_ids, attention_mask, dep_att_mask, labels=labels)
            tmp_eval_loss, logits = outputs[:2]
            eval_loss += tmp_eval_loss.mean().item()
        nb_eval_steps += 1
        preds = logits.detach().cpu().numpy()
        preds = np.argmax(preds, axis=-1).tolist()
        labels = batch['labels'].detach().cpu().numpy().tolist()
        eval_data['labels'] += labels
        eval_data['preds'] += preds
        eval_data['tokens'] += tokens

        if 'loss_mask' in batch:
            eval_data['loss_mask'] += batch['loss_mask'].detach().cpu().numpy().tolist()

        results = task_utils.evaluate(args.task, eval_data)

        preds = eval_data['preds']
        labels = eval_data['labels']
        loss_mask = eval_data['loss_mask']
        tokens = eval_data['tokens']

        pred_tags, true_tags, final_tokens = [], [], []
        label_list = task_utils.FceProcessor().get_labels()
        for i in range(len(preds)):
            for pid, tid, mask,tok in zip(preds[i], labels[i],  loss_mask[i], tokens[i].split()[1:-1] ):
                if mask == 1:
                    pred_tags.append(label_list[pid])
                    true_tags.append(label_list[tid])
                    final_tokens.append(tok)
                    out_file.write(tok + '\t' + true_tags[-1] + '\t' + pred_tags[-1] + '\n')

        #print(len(pred_tags), len(tokens))
    return results

if __name__=='__main__':
    parser = argparse.ArgumentParser()

    ## Required parameters
    parser.add_argument("--task", default='quantity', type=str)
    parser.add_argument("--data", default='MeasEval', type=str)
    parser.add_argument("--model_type", default='token-cls', type=str,
                        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()))
    parser.add_argument("--output_dir", default='outputs/', type=str,
                        help="The output directory where the model predictions and checkpoints will be written.")

    ## Other parameters
    parser.add_argument("--model_name_or_path", default='bert-base-uncased', type=str,
                        help="Path to pre-trained model or shortcut name selected in the list: " + ", ".join(ALL_MODELS))
    parser.add_argument("--data_dir", default="data", type=str,
                        help="The input data dir. Should contain the .tsv files (or other data files) for the task.")
    parser.add_argument("--config_name", default="", type=str,
                        help="Pretrained config name or path if not the same as model_name")
    parser.add_argument("--tokenizer_name", default="", type=str,
                        help="Pretrained tokenizer name or path if not the same as model_name")
    parser.add_argument("--cache_dir", default="cache", type=str,
                        help="Where do you want to store the pre-trained models downloaded from s3")
    parser.add_argument("--max_seq_length", default=128, type=int,
                        help="The maximum total input sequence length after tokenization. Sequences longer "
                             "than this will be truncated, sequences shorter will be padded.")

    parser.add_argument("--train_file", default="train_STEM-ECR_bert-base.pkl", type=str)
    parser.add_argument("--dev_file", default="dev_STEM-ECR_bert-base.pkl", type=str)

    parser.add_argument("--do_train", action='store_true', default=True,
                        help="Whether to run training.")
    parser.add_argument("--do_eval", action='store_true', default=False,
                        help="Whether to run eval on the dev set.")
    parser.add_argument("--do_predict", action='store_true', default=False)
    parser.add_argument("--do_lower_case", action='store_true',
                        help="Set this flag if you are using an uncased model.")

    parser.add_argument("--order_metric", default='f1', type=str)
    parser.add_argument("--metric_reverse", action='store_true')
    parser.add_argument("--num_save_ckpts", default=4, type=int)
    parser.add_argument("--remove_unused_ckpts", action='store_true')

    parser.add_argument("--per_gpu_train_batch_size", default=8, type=int,
                        help="Batch size per GPU/CPU for training.")
    parser.add_argument("--per_gpu_eval_batch_size", default=8, type=int,
                        help="Batch size per GPU/CPU for evaluation.")
    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
                        help="Number of updates steps to accumulate before performing a backward/update pass.")
    parser.add_argument("--learning_rate", default=5e-5, type=float,
                        help="The initial learning rate for Adam.")
    parser.add_argument("--weight_decay", default=0, type=float,
                        help="Weight decay if we apply some.")
    parser.add_argument("--adam_epsilon", default=1e-8, type=float,
                        help="Epsilon for Adam optimizer.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float,
                        help="Max gradient norm.")
    parser.add_argument("--num_train_epochs", default=5, type=float,
                        help="Total number of training epochs to perform.")
    parser.add_argument("--max_steps", default=-1, type=int,
                        help="If > 0: set total number of training steps to perform. Override num_train_epochs.")
    parser.add_argument("--warmup_steps", default=0, type=int,
                        help="Linear warmup over warmup_steps.")

    parser.add_argument('--logging_steps', type=int, default=100,
                        help="Log every X updates steps.")
    parser.add_argument('--save_steps', type=int, default=400,
                        help="Save checkpoint every X updates steps.")
    parser.add_argument("--eval_all_checkpoints", action='store_true',
                        help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number")
    parser.add_argument("--no_cuda", action='store_true',
                        help="Avoid using CUDA when available")
    parser.add_argument('--overwrite_output_dir', action='store_true',
                        help="Overwrite the content of the output directory")
    parser.add_argument('--overwrite_cache', action='store_true',
                        help="Overwrite the cached training and evaluation sets")
    parser.add_argument('--seed', type=int, default=42,
                        help="random seed for initialization")

    parser.add_argument('--fp16', action='store_true',
                        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit")
    parser.add_argument('--fp16_opt_level', type=str, default='O1',
                        help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
                             "See details at https://nvidia.github.io/apex/amp.html")
    parser.add_argument("--local_rank", type=int, default=-1,
                        help="For distributed training: local_rank")
    args = parser.parse_args()

    if os.path.exists(args.output_dir) and os.listdir(args.output_dir) and args.do_train and not args.overwrite_output_dir:
        raise ValueError("Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(args.output_dir))

    if args.local_rank == -1 or args.no_cuda:
        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
        args.n_gpu = torch.cuda.device_count()
    else:
        torch.cuda.set_device(args.local_rank)
        device = torch.device("cuda", args.local_rank)
        torch.distributed.init_process_group(backend='nccl')
        args.n_gpu = 1
    args.device = device
    args.to_gpu_keys = ['input_ids', 'attention_mask', 'dep_att_mask', 'labels', 'loss_mask']

    # Setup logging
    logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                        datefmt = '%m/%d/%Y %H:%M:%S',
                        level = logging.INFO if args.local_rank in [-1, 0] else logging.WARN)

    set_seed(args)

    args.model_type = args.model_type.lower()
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path,
                                          selection_mode='soft',
                                          num_labels=task_utils.num_labels[args.task],
                                          cache_dir=args.cache_dir if args.cache_dir else None,
                                          output_selections=False)
    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
                                                do_lower_case=args.do_lower_case,
                                                cache_dir=args.cache_dir if args.cache_dir else None)
    model = model_class.from_pretrained(args.model_name_or_path, config=config,
                                        cache_dir=args.cache_dir if args.cache_dir else None)
    model.to(args.device)
    batch_processor = model_class.build_batch

    # Training
    if args.do_train:
        global_step, tr_loss = train(args, model, tokenizer, batch_processor)

    if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
        for i in range(5):
            set_seed(seed * (i + 1))
            if not os.path.exists(args.output_dir +'_'+i) and args.local_rank in [-1, 0]:
                os.makedirs(args.output_dir +'_'+i)

            model_to_save = model.module if hasattr(model, 'module') else model
            model_to_save.save_pretrained(args.output_dir)
            tokenizer.save_pretrained(args.output_dir)

            torch.save(args, os.path.join(args.output_dir +'_'+i, 'training_args.bin'))

            model = model_class.from_pretrained(args.output_dir +'_'+i)
            tokenizer = tokenizer_class.from_pretrained(args.output_dir +'_'+i)
            model.to(args.device)

    if args.do_predict:
        for i in range(5):
            checkpoints = [args.output_dir +'_'+i]
            checkpoint = [args.output_dir +'_'+i][-1]
            global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
            prefix = checkpoint.split('/')[-1] if checkpoint.find('saved_ckpt-') != -1 else ""
            data_dir = '../data_files'

            train_docs, train_annotations = os.path.join(data_dir, 'train/text'), os.path.join(data_dir, 'train/tsv')
            test_docs, test_annotations = os.path.join(data_dir, 'eval/text'), os.path.join(data_dir, 'eval/tsv')

            data_processor = DataProcessor_Data(train_docs)

            test_sentences, test_labels = data_processor.get_quantities_labels(test_docs, test_annotations, cat_eval=False, category=[], part='test')
            tokenizer = BertTokenizer.from_pretrained(args.output_dir +'_'+i)
            model = BertForTokenClassification.from_pretrained(checkpoint)
            model.to(args.device)
            explainer = TokenClassificationExplainer(model=model, tokenizer=tokenizer)
            predict(test_sentences, test_labels, explainer)
    # Evaluation
    if args.do_eval and args.local_rank in [-1, 0]:
        for i in range(5):
            checkpoints = [args.output_dir +'_'+i]
            if args.eval_all_checkpoints:
                checkpoints = list(os.path.dirname(c) for c in sorted(glob.glob(args.output_dir +'_'+i + '/**/' + WEIGHTS_NAME, recursive=True)))
                logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)

            best_ckpt_dirs = []
            results = {}

            for checkpoint in checkpoints:
                global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
                prefix = checkpoint.split('/')[-1] if checkpoint.find('saved_ckpt-') != -1 else ""

                model = model_class.from_pretrained(checkpoint, config=config)
                model.to(args.device)
                result = evaluate(args, model, tokenizer, batch_processor, prefix=prefix)
                best_ckpt_dirs.append((result[args.order_metric], checkpoint))

                result = dict((k + '_{}'.format(global_step), v) for k, v in result.items())
                results.update(result)

            best_ckpt_dirs = sorted(best_ckpt_dirs, reverse=args.metric_reverse)[:args.num_save_ckpts]
            best_ckpt_dirs = [d for v, d in best_ckpt_dirs]

            json.dump(results, open(os.path.join(args.output_dir, 'dev_results.json'), 'w', encoding='utf-8'), indent=4)

            if args.remove_unused_ckpts:
                for checkpoint in checkpoints:
                    prefix = checkpoint.split('/')[-1] if checkpoint.find('saved_ckpt-') != -1 else ""
                    if len(prefix) != 0 and (checkpoint not in best_ckpt_dirs):
                        os.system("rm -rf %s"%checkpoint)

            checkpoints = list(os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + '/**/' + WEIGHTS_NAME, recursive=True)))

            results_dev = {}
            results_test = {}
            results_cat = []
            for checkpoint in checkpoints:
                global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
                prefix = checkpoint.split('/')[-1] if checkpoint.find('saved_ckpt-') != -1 else ""

                model = model_class.from_pretrained(checkpoint, config=config)
                model.to(args.device)
                result = evaluate(args, model, tokenizer, batch_processor, prefix=prefix)
                result = dict((k + '_{}'.format(global_step), v) for k, v in result.items())
                results_dev.update(result)
                result = evaluate_test(args, model, tokenizer, batch_processor, prefix=prefix)
                results_test.update(result)

            json.dump(results_dev, open(os.path.join(args.output_dir, 'predict_results_dev.json'), 'w', encoding='utf-8'), indent=4)
            json.dump(results_test, open(os.path.join(args.output_dir, 'predict_results_test.json'), 'w', encoding='utf-8'), indent=4)


from simpy import Store, Environment

from simulator.factory_models import Storage

class StorageResource(Store):
    """
    貯蔵庫、倉庫クラス
    仕掛品、半製品、製品をＦＩＦＯで管理する
    """
    def __init__(
            self,
            storage: Storage,
            env: Environment,
            capacity=float('inf')
    ):
        """
        コンストラクタ
        :param storage: 貯蔵庫情報
        :param env: Simpy.Environment
        :param capacity: 貯蔵庫容量（省略時は無限)
        """
        super(StorageResource, self).__init__(env=env, capacity=capacity)
        self.storage = storage

    def get_compound_id(self) -> str:
        return self.storage.get_compound_id()


"""A pytest plugin which helps with integration testing for Kubernetes deployments.

This plugin helps with the managing of the Kubernetes cluster and provides
useful test fixtures and functionality in order to interact with and test
the state of the cluster.
"""

# import os
import logging

import kubernetes
import pytest

from kubetest import markers
from kubetest.manager import KubetestManager

GOOGLE_APPLICATION_CREDENTIALS = 'GOOGLE_APPLICATION_CREDENTIALS'

log = logging.getLogger('kubetest')

# A global instance of the KubetestManager. This will be used by various
# pytest hooks and fixtures in order to create and manage the TestClient
# instances. The manager is in charge of creating new TestClients for each
# test that requests one via the `k8s` fixture, and to manage the lifecycle
# of the test namespace.
manager = KubetestManager()

# ********** pytest hooks **********

def pytest_addoption(parser):
    """Add options to pytest to configure kubetest.

    See Also:
        https://docs.pytest.org/en/latest/reference.html#_pytest.hookspec.pytest_addoption
    """

    group = parser.getgroup('kubetest', 'kubernetes integration test support')
    group.addoption(
        '--kube-config',
        action='store',
        metavar='path',
        default=None,
        help='the kubernetes config for kubetest'
    )
    group.addoption(
        '--kube-disable',
        action='store_true',
        default=False,
        help='disable automatic configuration with the kubeconfig file'
    )

    # FIXME (etd) - this was an attempt to fix occassional permissions errors
    # (https://github.com/vapor-ware/kubetest/issues/11) but in doing so, it looks
    # like I hosed my permissions, so I'm just commenting all of this out for now...
    # group.addoption(
    #     '--google-application-credentials',
    #     action='store',
    #     metavar='path',
    #     default=None,
    #     help='google application credentials for GKE clusters'
    # )

    group.addoption(
        '--kube-log-level',
        action='store',
        default='warning',
        help='log level for the kubetest logger'
    )

    group.addoption(
        '--kube-error-log-lines',
        action='store',
        default=50,
        help='set the number of lines to tail from container logs on error. '
             'to show all lines, set this to -1.'
    )

def pytest_report_header(config):
    """Augment the pytest report header with kubetest info.

    See Also:
        https://docs.pytest.org/en/latest/reference.html#_pytest.hookspec.pytest_report_header
    """
    disabled = config.getoption('kube_disable')
    if not disabled:
        config_file = config.getoption('kube_config')
        if config_file is None:
            config_file = 'default'

        return [
            'kubetest config file: {}'.format(config_file),
        ]

def pytest_configure(config):
    """Configure pytest with kubetest additions.

    This registers the kubetest markers with pytest.

    See Also:
        https://docs.pytest.org/en/latest/reference.html#_pytest.hookspec.pytest_configure
    """
    # Register kubetest markers
    markers.register(config)

def pytest_sessionstart(session):
    """Configure kubetest for the test session.

    Kubetest setup happens at session start (pytest_sessionstart) rather
    than on configuration (pytest_configure) so that we only have the
    expectation of a cluster config and available cluster when there are
    actually tests available. For example, if we are simply calling any of

        pytest --help
        pytest --markers
        pytest --fixtures

    we do not want to configure kubetest and force the expectation of
    cluster availability.

    See Also:
        https://docs.pytest.org/en/latest/reference.html#_pytest.hookspec.pytest_sessionstart
    """
    # Setup the kubetest logger
    log_level = session.config.getoption('kube_log_level')
    level = logging._nameToLevel.get(log_level.upper(), logging.WARNING)
    logger = logging.getLogger('kubetest')
    logger.setLevel(level)

    # Configure kubetest with the kubernetes config, if not disabled.
    disabled = session.config.getoption('kube_disable')
    if not disabled:
        # # Set the GOOGLE_APPLICATION_CREDENTIALS environment variable. For more, see:
        # # https://cloud.google.com/docs/authentication/production
        # gke_creds = config.getoption('google_application_credentials')
        # if gke_creds is not None:
        #     # If application credentials already exist, store them so they can be
        #     # reset after testing.
        #     old_creds = os.environ.get(GOOGLE_APPLICATION_CREDENTIALS)
        #     if old_creds:
        #         os.environ['OLD_' + GOOGLE_APPLICATION_CREDENTIALS] = old_creds
        #     os.environ[GOOGLE_APPLICATION_CREDENTIALS] = gke_creds
        pass

# def pytest_sessionfinish(session):
#     """Unconfigure the test session for kubetest.
#
#     See Also:
#         https://docs.pytest.org/en/latest/reference.html#_pytest.hookspec.pytest_sessionfinish
#     """
#     # Unset/reset the GOOGLE_APPLICATION_CREDENTIALS environment variable. For more,
#     # see: https://cloud.google.com/docs/authentication/production
#     gke_creds = config.getoption('google_application_credentials')
#     if gke_creds is not None:
#         # If there are old credentials stored, reset them. Otherwise, just unset
#         # the environment variable.
#         old = os.environ.get('OLD_' + GOOGLE_APPLICATION_CREDENTIALS)
#         if old:
#             os.environ[GOOGLE_APPLICATION_CREDENTIALS] = old
#             del os.environ['OLD_' + GOOGLE_APPLICATION_CREDENTIALS]
#         else:
#             del os.environ[GOOGLE_APPLICATION_CREDENTIALS]

def pytest_runtest_setup(item):
    """Run setup actions to prepare the test case.

    See Also:
        https://docs.pytest.org/en/latest/reference.html#_pytest.hookspec.pytest_runtest_setup
    """
    # Register a new test case with the manager and setup the test case state.
    test_case = manager.new_test(
        node_id=item.nodeid,
        test_name=item.name,
    )

    # FIXME (etd) - not sure this is really what we want to do. does it make sense
    # to entirely disable the plugin just be specifying the disable flag? probably..
    # but there must be a better way than adding this check (perhaps unregistering the
    # plugin in the pytest_configure hook?)
    disabled = item.config.getoption('kube_disable')
    if not disabled:

        # Register test case state based on markers on the test case
        test_case.register_rolebindings(
            *markers.rolebindings_from_marker(item, test_case.ns)
        )
        test_case.register_clusterrolebindings(
            *markers.clusterrolebindings_from_marker(item, test_case.ns)
        )

        # Apply manifests for the test case, if any are specified.
        markers.apply_manifest_from_marker(item, test_case)

def pytest_runtest_teardown(item):
    """Run teardown actions to clean up the test client.

    See Also:
        https://docs.pytest.org/en/latest/reference.html#_pytest.hookspec.pytest_runtest_teardown
    """
    disabled = item.config.getoption('kube_disable')
    if not disabled:
        manager.teardown(item.nodeid)

def pytest_runtest_makereport(item, call):
    """Create a test report for the test case. If the test case was found
    to fail, this will log out the container logs to provide more debugging
    context.

    See Also:
        https://docs.pytest.org/en/latest/reference.html#_pytest.hookspec.pytest_runtest_makereport
    """
    disabled = item.config.getoption('kube_disable')
    if not disabled:
        if call.when == 'call':
            if call.excinfo is not None:
                tail_lines = item.config.getoption('kube_error_log_lines')
                if tail_lines != 0:
                    test_case = manager.get_test(item.nodeid)
                    logs = test_case.yield_container_logs(
                        tail_lines=tail_lines
                    )
                    for container_log in logs:
                        # Add a report section to the test output
                        item.add_report_section(
                            when=call.when,
                            key='kubernetes container logs',
                            content=container_log
                        )

def pytest_keyboard_interrupt():
    """Clean up the cluster from kubetest artifacts if the tests
    are manually terminated via keyboard interrupt.

    See Also:
        https://docs.pytest.org/en/latest/reference.html#_pytest.hookspec.pytest_keyboard_interrupt
    """
    try:
        namespaces = kubernetes.client.CoreV1Api().list_namespace()
        for ns in namespaces.items:
            # if the namespace has a 'kubetest-' prefix, remove it.
            name = ns.metadata.name
            status = ns.status
            if (
                    name.startswith('kubetest-') and
                    status is not None and
                    status.phase.lower() == 'active'
            ):
                print('keyboard interrupt: cleaning up namespace "{}"'.format(name))
                kubernetes.client.CoreV1Api().delete_namespace(
                    body=kubernetes.client.V1DeleteOptions(),
                    name=name,
                )

        crbs = kubernetes.client.RbacAuthorizationV1Api().list_cluster_role_binding()
        for crb in crbs.items:
            # if the cluster role binding has a 'kubetest:' prefix, remove it.
            name = crb.metadata.name
            if name.startswith('kubetest:'):
                print(
                    'keyboard interrupt: cleaning up clusterrolebinding "{}"'.format(crb)
                )
                kubernetes.client.RbacAuthorizationV1Api().delete_cluster_role_binding(
                    body=kubernetes.client.V1DeleteOptions(),
                    name=name,
                )
    except Exception as e:
        log.error(
            'Failed to clean up kubetest artifacts from cluster on keyboard interrupt. '
            'You may need to manually remove items from your cluster. Check for '
            'namespaces with the "kubetest-" prefix and cluster role bindings with '
            'the "kubetest:" prefix. ({})'.format(e)
        )

# ********** pytest fixtures **********

@pytest.fixture
def kubeconfig(request):
    config_file = request.session.config.getoption('kube_config')
    return config_file

@pytest.fixture()
def kube(kubeconfig, request):
    """Return a client for managing a Kubernetes cluster for testing."""
    test_case = manager.get_test(request.node.nodeid)
    if test_case is None:
        logging.getLogger('kubetest').warning(
            'No kubetest test client found for test using the "kube" fixture. '
            '(are you running with the --kube-disable flag?)'
        )
        return None

    # Setup the test case. This will create the namespace and any other
    # objects (e.g. role bindings) that the test case will need.
    kubernetes.config.load_kube_config(config_file=kubeconfig)
    test_case.setup()

    return test_case.client


from checkov.common.models.enums import CheckResult, CheckCategories
from checkov.terraform.checks.resource.base_resource_check import BaseResourceCheck
from checkov.common.util.secrets import string_has_secrets, AWS, GENERAL
from checkov.common.util.type_forcers import force_list

class LambdaEnvironmentCredentials(BaseResourceCheck):

    def __init__(self):
        name = "Ensure no hard-coded secrets exist in lambda environment"
        id = "CKV_AWS_45"
        supported_resources = ['aws_lambda_function']
        categories = [CheckCategories.SECRETS]
        super().__init__(name=name, id=id, categories=categories, supported_resources=supported_resources)

    def scan_resource_conf(self, conf):
        self.evaluated_keys = 'environment/[0]/variables'
        if len(conf.get('environment', [])) > 0:
            if isinstance(conf['environment'][0], dict):
                if 'variables' in conf['environment'][0]:
                    if isinstance(force_list(conf['environment'][0]['variables'])[0], dict):
                        # variables can be a string, which in this case it points to a variable
                        for values in list(force_list(conf['environment'][0]['variables'])[0].values()):
                            for value in list(filter(lambda value: isinstance(value, str), force_list(values))):
                                if string_has_secrets(value,AWS,GENERAL):
                                    return CheckResult.FAILED
        return CheckResult.PASSED

check = LambdaEnvironmentCredentials()


from django.contrib import admin
from django.urls import path,include
from myapp import views

urlpatterns = [
   path('',views.index),
   path('about/',views.about),
   path('contact/',views.contact),
   path('blog/',views.blog),
   path('products/',views.products),
]


from api.views import LicenseInfoView, get_organizations, get_software_recommendations, leaderboard, \
    get_org_software_names, \
    get_org_software_users, get_licenses_associated_with_user, get_reallocatabe_by_software_name, \
    get_org_software_users_by_name, software_counts, GetUserInfo, check_if_unused, get_potential_savings
from django.urls import path

urlpatterns = [
    path('organizations/', get_organizations, name='organizations'),
    path('recommendations/', get_software_recommendations, name='recommendations'),
    path('software/', get_org_software_names, name='software'),
    path('applications/', get_org_software_users, name='get_applications_by_user'),
    path('userlicenses/', get_licenses_associated_with_user,
         name='get_licenses_associated_with_user'),
    path('softwarebyname/<str:software>', get_reallocatabe_by_software_name,
         name='get_reallocatabe_by_software_name'),
    path('softwarebyuser/', get_org_software_users_by_name, name='get_org_software_users_by_name'),
    path('count', software_counts, name='software_counts'),
    path('user/', GetUserInfo.as_view(), name='user'),
    path('licenseinfo/', LicenseInfoView.as_view(), name='licenseinfo'),
    path('leaderboard/', leaderboard, name='leaderboard'),
    path('check/', check_if_unused, name='check_if_unused'),
    path('moneysaved/', get_potential_savings, name='get_potential_savings'),

]


import argparse
import json
import logging
import math
import os
import random
import time
import warnings
warnings.filterwarnings('ignore')

os.environ["HOME"] = "/home/nsml"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import numpy as np
from sklearn.model_selection import KFold

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from torch.optim.lr_scheduler import LambdaLR
import transformers
from transformers import BertTokenizer, get_linear_schedule_with_warmup

from model import TransformerModel, EncoderDecoder, BertGpt2, TransformerS2Model
from tokenizer import CharTokenizer, CustomTokenizer
from dataset import TextDataset, collate_fn, collate_fn_pretrain, collate_fn_transformer, collate_fn_pretrain_transformer, collate_fn_mlm_pretrain_transformer
from data_loader import read_strings
from meter import Meter
from evaluation import em, gleu

import larva
from larva import LarvaTokenizer, LarvaModel, list_larva_models

import nsml
from nsml import DATASET_PATH

logging.basicConfig(format='%(asctime)s -  %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO)
logger = logging.getLogger(__name__)

def get_args():
    parser = argparse.ArgumentParser()
    # data
    parser.add_argument("--data_dir", type=str, default=os.path.join(DATASET_PATH, 'train'))
    parser.add_argument("--val_ratio", type=float, default=0.2)
    # parser.add_argument("--num_val_data", type=int, default=1000)

    # tokenizer
    parser.add_argument("--limit_alphabet", type=int, default=6000)
    parser.add_argument("--tokenizer_vocab_size", type=int, default=6000)

    # model
    parser.add_argument("--model", type=str, default='transformer')
    parser.add_argument("--vocab_size", type=int, default=1600)
    parser.add_argument("--hidden_size", type=int, default=512)
    parser.add_argument("--num_attention_heads", type=int, default=8)
    parser.add_argument("--num_encoder_layers", type=int, default=6)
    parser.add_argument("--num_decoder_layers", type=int, default=6)
    parser.add_argument("--intermediate_size", type=int, default=1024)
    parser.add_argument("--ffn_size", type=int, default=2048)
    parser.add_argument("--share_decoder_input_output_embed", type=str2bool, default=False)

    # bert
    parser.add_argument("--bert", type=str, default='original')
    parser.add_argument("--bert_hidden_size", type=int, default=768)
    parser.add_argument("--bert_num_attention_heads", type=int, default=12)
    parser.add_argument("--num_hidden_layers", type=int, default=12)
    parser.add_argument("--bert_intermediate_size", type=int, default=3072)

    # training
    parser.add_argument("--kfold", type=str2bool, default=False)
    parser.add_argument("--n_splits", type=int, default=5)

    parser.add_argument("--mlm_pre", type=str2bool, default=False)
    parser.add_argument("--num_pretrain_epochs", type=int, default=5)

    parser.add_argument("--train_all", type=str2bool, default=False)

    parser.add_argument("--max_seq_length", type=int, default=128)
    parser.add_argument("--train_batch_size", type=int, default=128)
    parser.add_argument("--eval_batch_size", type=int, default=64)

    parser.add_argument("--lr", type=float, default=5e-4)
    parser.add_argument("--tokenizer", type=str, default='char')
    parser.add_argument("--optimizer", type=str, default='Adam')
    parser.add_argument("--scheduler", type=str, default='lambda')
    parser.add_argument("--adam_betas", type=str, default="(0.9, 0.98)")
    parser.add_argument("--eps", type=float, default=1e-9)
    parser.add_argument("--weight_decay", type=float, default=0.0001)
    parser.add_argument("--max_grad_norm", type=float, default=1.0)
    parser.add_argument("--dropout", type=float, default=0.3)

    parser.add_argument("--max_steps", type=int, default=10000)
    parser.add_argument("--num_warmup_steps", type=int, default=4000)
    parser.add_argument("--warmup_ratio", type=float, default=0.1)
    parser.add_argument("--log_interval", type=int, default=100)
    parser.add_argument("--eval_interval", type=int, default=1000)

    parser.add_argument('--seed', type=int, default=2021)
    parser.add_argument('--num_workers', type=int, default=20)

    # nsml
    parser.add_argument('--mode', type=str, default="train")
    parser.add_argument('--pause', type=int, default=0)
    parser.add_argument('--iteration', type=str, default="0")

    args = parser.parse_args()
    return args

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)

def str2bool(v : str) -> bool:
    """convert string argument to boolean
    Args:
        v (str)
    Raises:
        argparse.ArgumentTypeError: [Boolean value expected]
    Returns:
        bool
    """
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

def postprocess_state(sentence : str) -> str:
    """TRADE, SUMBT postprocessing
    Args:
        state (List[str]): state prediction
    Returns:
        List[str]: postprocessing state
    """
    sentence = sentence.replace(" : ", ":").replace(" , ", ", ").replace('( ', '(').replace(' )', ')').replace(' & ', '&').replace(' = ', '=')
    sentence = sentence.replace(" % ", "%").replace(' ~ ', '~').replace(' ^ ', '^')
    if sentence.endswith(' ~'):
        sentence = sentence.replace(' ~', '~')
    if sentence.endswith(' ^^'):
        sentence = sentence.replace(' ^^', '^^')
    if sentence.endswith(' ^'):
        sentence = sentence.replace(' ^', '^')
    if sentence.endswith('......'):
        sentence = sentence.replace('......', ' ......')
    sentence = sentence.replace(') 에', ')에').replace('곳 (', '곳(').replace('부터~트', '부터~ 트').replace('# 정왕동', '#정왕동')
    sentence = sentence.replace('쨘 -', '쨘-').replace('해드리겠습니다!', '해드리겠습니다 !').replace('6 / 6', '6/6').replace('6 / 4', '6/4')
    sentence = sentence.replace('> ㅋ', '>ㅋ').replace('이상~헤', '이상~ 헤').replace('6 / 6', '6/6').replace('6 / 4', '6/4')

    return sentence

def get_lr(optimizer : transformers.optimization) -> float:
    """Get learning_rate.

    Args:
        optimizer (transformers.optimization).

    Return:
        Learning_Rate (float).
    """
    for param_group in optimizer.param_groups:
        return param_group['lr']

def calc_loss(model, batch, args):
    src, tgt, src_mask, tgt_mask, tgt_label = batch
    if args.model == 'transformer':
        output = model(src=src, tgt=tgt, src_key_padding_mask=~src_mask, tgt_key_padding_mask=~tgt_mask)
    elif args.model in ['encoderdecoder', 'BertGpt2']:
        output = model(src=src, tgt=tgt, src_key_padding_mask=src_mask, tgt_key_padding_mask=tgt_mask)
    elif args.model == 'bertfuse':
        output = model(src_tokens=src, prev_output_tokens=tgt, bert_input=src)

    if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
        mask_tgt_label = tgt_label.clone()
        mask_tgt_label[tgt_label.eq(1)] = -100
    else:
        mask_tgt_label = tgt_label.clone()
        mask_tgt_label[tgt_label.eq(0)] = -100

    raw_loss = F.cross_entropy(output.view(-1, output.size(-1)), mask_tgt_label.view(-1), reduction='none')
    # raw_loss = F.cross_entropy(output.view(-1, output.size(-1)), tgt_label.view(-1), reduction='none')

    if args.model == 'transformer':
        bsz = tgt.size(1)
        raw_loss = raw_loss.view(-1, bsz)
    else:
        bsz = tgt.size(0)
        raw_loss = raw_loss.view(bsz, -1)

    if args.model == 'transformer':
        loss = (raw_loss * tgt_mask.float()).sum(0).mean()
    else:
        loss = (raw_loss * tgt_mask.float()).sum(-1).mean()

    items = [loss.data.item(), bsz, tgt_mask.sum().item()]
    return loss, items

def calc_loss1(model, batch, args):
    src, tgt, src_mask, tgt_mask, tgt_label = batch

    print(f"source padded is {src}, source padded's shape is {src.shape}")
    print(f"target padded before -1 is {tgt}, target padded before -1 's shape is {tgt.shape}")
    print(f"source padding mask is {src_mask}, source padding mask's shape is {src_mask.shape}")
    print(f"target padding mask is {tgt_mask}, target padding mask's shape is {tgt_mask.shape}")
    print(f"target padded after 1 is {tgt_label}, target padded after 1 's shape is {tgt_label.shape}")

    if args.model == 'transformer':
        output = model.forward1(src=src, tgt=tgt, src_key_padding_mask=~src_mask, tgt_key_padding_mask=~tgt_mask)
    elif args.model in ['encoderdecoder', 'BertGpt2']:
        output = model.forward1(src=src, tgt=tgt, src_key_padding_mask=src_mask, tgt_key_padding_mask=tgt_mask)
    elif args.model == 'bertfuse':
        output = model(src_tokens=src, prev_output_tokens=tgt, bert_input=src)

    print(f"output shape is {output.shape}")

    if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
        mask_tgt_label = tgt_label.clone()
        mask_tgt_label[tgt_label.eq(1)] = -100
    else:
        mask_tgt_label = tgt_label.clone()
        mask_tgt_label[tgt_label.eq(0)] = -100

    print(f"cross entropy first input is {output.view(-1, output.size(-1)).shape}, second input is {tgt_label.view(-1).shape}")

    raw_loss = F.cross_entropy(output.view(-1, output.size(-1)), mask_tgt_label.view(-1), reduction='none')
    # raw_loss = F.cross_entropy(output.view(-1, output.size(-1)), tgt_label.view(-1), reduction='none')

    if args.model == 'transformer':
        bsz = tgt.size(1)
        raw_loss = raw_loss.view(-1, bsz)
    else:
        bsz = tgt.size(0)
        raw_loss = raw_loss.view(bsz, -1)

    print(f"raw_loss is {raw_loss}, raw_loss 's shape is {raw_loss.shape}")
    if args.model == 'transformer':
        loss = (raw_loss * tgt_mask.float()).sum(0).mean()
    else:
        print(f"sum of raw_loss is {(raw_loss * tgt_mask.float()).sum(-1)}, sum of raw_loss's shape is {(raw_loss * tgt_mask.float()).sum(-1).shape}")
        loss = (raw_loss * tgt_mask.float()).sum(-1).mean()

    print(f"loss is {loss}, loss 's shape is {loss.shape}")

    items = [loss.data.item(), bsz, tgt_mask.sum().item()]
    return loss, items

def evaluate(model, data_loader, args):
    model.eval()
    meter = Meter()
    with torch.no_grad():
        for batch in data_loader:
            batch = tuple(t.to(args.device) for t in batch)
            _, items = calc_loss(model, batch, args)
            meter.add(*items)
    return meter.average(), meter.print_str(False)

def correct(model, tokenizer, test_data, args):
    model.eval()
    prediction = []
    for i in range(0, len(test_data), args.eval_batch_size):
        batch = test_data[i:i + args.eval_batch_size]

        if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
            # src_token_ids = [tokenizer(x) for x in batch]
            src_token_ids = [[2] + tokenizer(x) + [3] for x in batch]
        elif args.tokenizer in ['dsksd', 'plus', 'wordpiece']:
            src_token_ids = tokenizer(batch).input_ids

        if args.model == 'transformer':
            src_seq_length = [len(x) for x in src_token_ids]
            src_max_seq_length = max(src_seq_length)
            src_padded = []
            src_padding_mask = []
            for x in src_token_ids:
                x = x[:src_max_seq_length]
                src_pad_length = src_max_seq_length - len(x)
                src_padded.append(x + [tokenizer.pad_token_id] * src_pad_length)
                src_padding_mask.append([1] * len(x) + [0] * src_pad_length)

            src_padded = torch.tensor(src_padded).t().contiguous().to(args.device)
            src_padding_mask = torch.tensor(src_padding_mask).bool().t().to(args.device)

            memory = model(src=src_padded, src_key_padding_mask=~src_padding_mask)

            tgt_token_ids = [[2] for _ in batch]
            end = [False for _ in batch]
            for l in range(src_max_seq_length + 20):
                tgt = torch.tensor(tgt_token_ids).t().contiguous().to(args.device)
                output = model(tgt=tgt, memory=memory, memory_key_padding_mask=~src_padding_mask)
                top1 = output[-1].argmax(-1).tolist()
                for i, tok in enumerate(top1):
                    if tok == 3 or l >= src_seq_length[i] + 20:
                        end[i] = True
                    tgt_token_ids[i].append(tok if not end[i] else 3)
                if all(end):
                    break
        elif args.model in ['encoderdecoder', 'BertGpt2']:
            src_seq_length = [len(x) for x in src_token_ids]
            src_max_seq_length = max(src_seq_length)

            src_padded = []
            for x in src_token_ids:
                x = x[:src_max_seq_length]
                src_pad_length = src_max_seq_length - len(x)
                src_padded.append(x + [tokenizer.pad_token_id] * src_pad_length)

            src_padded = torch.LongTensor(src_padded).contiguous().to(args.device)
            src_padding_mask = src_padded.ne(tokenizer.pad_token_id).to(args.device)

            memory = model.encoder(
                input_ids=src_padded,
                attention_mask=src_padding_mask,
            )

            tgt_token_ids = [[2] for _ in batch]
            end = [False for _ in batch]
            for l in range(src_max_seq_length + 20):
                tgt_padded = torch.LongTensor(tgt_token_ids).contiguous().to(args.device)
                tgt_padding_mask = tgt_padded.ne(tokenizer.pad_token_id).to(args.device)

                # output = model(src=src_padded, tgt=tgt_padded, src_key_padding_mask=src_padding_mask, tgt_key_padding_mask=tgt_padding_mask)
                output = model(tgt=tgt_padded, memory=memory, src_key_padding_mask=src_padding_mask, tgt_key_padding_mask=tgt_padding_mask)
                top1 = output[:,-1,:].argmax(-1).tolist()

                for i, tok in enumerate(top1):
                    if tok == 3 or l >= src_seq_length[i] + 20:
                        end[i] = True
                    tgt_token_ids[i].append(tok if not end[i] else 3)
                if all(end):
                    break
        elif args.model == 'bertfuse':
            src_seq_length = [len(x) for x in src_token_ids]
            src_max_seq_length = max(src_seq_length)

            src_padded = []
            for x in src_token_ids:
                x = x[:src_max_seq_length]
                src_pad_length = src_max_seq_length - len(x)
                src_padded.append(x + [tokenizer.pad_token_id] * src_pad_length)

            src_padded = torch.LongTensor(src_padded).contiguous().to(args.device)
            bert_encoder_padding_mask = src_padded.eq(tokenizer.pad_token_id)

            bert_encoder_out =  model.bert_encoder(src_padded, attention_mask=~bert_encoder_padding_mask)
            bert_encoder_out = bert_encoder_out.last_hidden_state

            bert_encoder_out = bert_encoder_out.permute(1,0,2).contiguous()
            bert_encoder_out = {
                'bert_encoder_out': bert_encoder_out,
                'bert_encoder_padding_mask': bert_encoder_padding_mask,
            }

            encoder_out = model.encoder(src_padded, bert_encoder_out=bert_encoder_out)

            tgt_token_ids = [[2] for _ in batch]
            end = [False for _ in batch]
            for l in range(src_max_seq_length + 20):
                tgt_padded = torch.LongTensor(tgt_token_ids).contiguous().to(args.device)

                output = model.decoder(tgt_padded, encoder_out=encoder_out, bert_encoder_out=bert_encoder_out)[0]

                top1 = output[:,-1,:].argmax(-1).tolist()

                for i, tok in enumerate(top1):
                    if tok == 3 or l >= src_seq_length[i] + 20:
                        end[i] = True
                    tgt_token_ids[i].append(tok if not end[i] else 3)
                if all(end):
                    break

        if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
            # prediction.extend([''.join([tokenizer.i2c[tok] for tok in x if tok >= 5]) for x in tgt_token_ids])
            prediction.extend([postprocess_state(''.join([tokenizer.i2c[tok] for tok in x if tok >= 5])) for x in tgt_token_ids])
        elif args.tokenizer in ['dsksd', 'plus', 'wordpiece']:
            prediction.extend([postprocess_state(tokenizer.decode([tok for tok in x if tok >= 5])) for x in tgt_token_ids])
    return prediction

def mlm_pretrain(args : argparse.Namespace, model : nn.Module, loader : DataLoader, optimizer : transformers.optimization, scheduler,
                tokenizer : BertTokenizer, loss_fnc_pretrain : nn.CrossEntropyLoss, epoch : int) -> None:
    """Masked Language Model Pretrain
    Args:
        args (argparse.Namespace)
        model (nn.Module)
        loader (DataLoader)
        optimizer (transformers.optimization)
        tokenizer (BertTokenizer)
        loss_fnc_pretrain (nn.CrossEntropyLoss): mlm pretrain loss
        epoch (int)
    """
    model.train()
    for step, batch in enumerate(loader):
        src, tgt, src_mask, tgt_mask, tgt_label = tuple(t.to(args.device) for t in batch)

        logits, labels = model.forward_pretrain(src, src_mask, tokenizer)
        loss = loss_fnc_pretrain(logits.view(-1, args.vocab_size), labels.view(-1))

        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
        optimizer.step()
        optimizer.zero_grad()
        scheduler.step()

        if step % 100 == 0:
            current_lr = get_lr(optimizer)
            epochs = args.num_pretrain_epochs
            print(
                f"[{epoch + 1}/{epochs}] [{step}/{len(loader)}] loss : {loss.item()}  lr : {current_lr}"
            )

def train(model, tokenizer, train_data, valid_data, corpuses, args, index = None):
    global best_val_gleu

    train_dataset = TextDataset(train_data)

    print(f"\nfirst train_dataset is {train_dataset[0]}")

    if args.model == 'transformer':
        if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
            collate_func = collate_fn
        else:
            collate_func = collate_fn_pretrain
    elif args.model in ['encoderdecoder', 'BertGpt2', 'bertfuse']:
        if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
            collate_func = collate_fn_transformer
        elif args.tokenizer in ['plus', 'wordpiece']:
            collate_func = collate_fn_pretrain_transformer

    train_dataloader = DataLoader(
        train_dataset,
        sampler=RandomSampler(train_dataset),
        batch_size=args.train_batch_size,
        num_workers=args.num_workers,
        collate_fn=lambda x: collate_func(x, tokenizer, args.max_seq_length),
    )

    if valid_data is not None:
        valid_dataset = TextDataset(valid_data)
        valid_dataloader = DataLoader(
            valid_dataset,
            sampler=SequentialSampler(valid_dataset),
            batch_size=args.eval_batch_size,
            num_workers=args.num_workers,
            collate_fn=lambda x: collate_func(x, tokenizer, args.max_seq_length),
        )

        valid_noisy = [x['noisy'] for x in valid_data]
        valid_clean = [x['clean'] for x in valid_data]

    epochs = (args.max_steps - 1) // len(train_dataloader) + 1
    optimizer = getattr(torch.optim, args.optimizer)(
        model.parameters(),
        lr=args.lr,
        betas=eval(args.adam_betas),
        eps=args.eps,
        weight_decay=args.weight_decay
    )

    if args.scheduler == 'lambda':
        lr_lambda = lambda x: x / args.num_warmup_steps if x <= args.num_warmup_steps else (x / args.num_warmup_steps) ** -0.5
        scheduler = LambdaLR(optimizer, lr_lambda)
    elif args.scheduler == 'linear':
        warmup_steps = args.max_steps * args.warmup_ratio
        scheduler = get_linear_schedule_with_warmup(
            optimizer, num_warmup_steps=warmup_steps, num_training_steps=args.max_steps
        )

    if args.mlm_pre and (index is None or index == 0):
        print('mlm_pretrain!!!')
        pre_optimizer = getattr(torch.optim, args.optimizer)(
            model.parameters(),
            lr=args.lr,
            betas=eval(args.adam_betas),
            eps=args.eps,
            weight_decay=args.weight_decay
        )
        loss_fnc_pretrain = nn.CrossEntropyLoss()

        if args.tokenizer == 'char':
            t_total = len(train_dataloader) * args.num_pretrain_epochs
            warmup_steps = t_total * args.warmup_ratio
            pre_scheduler = get_linear_schedule_with_warmup(
                pre_optimizer, num_warmup_steps=warmup_steps, num_training_steps=t_total
            )
            for epoch in range(args.num_pretrain_epochs):
                mlm_pretrain(args, model, train_dataloader, pre_optimizer, pre_scheduler, tokenizer, loss_fnc_pretrain, epoch)
        elif args.tokenizer == 'pre_char':
            corpus_dataset = TextDataset(corpuses)
            corpus_dataloader = DataLoader(
                corpus_dataset,
                sampler=RandomSampler(corpus_dataset),
                batch_size=args.train_batch_size,
                num_workers=args.num_workers,
                collate_fn=lambda x: collate_fn_pretrain(x, tokenizer, args.max_seq_length),
            )
            t_total = len(corpus_dataloader) * args.num_pretrain_epochs
            warmup_steps = t_total * args.warmup_ratio
            pre_scheduler = get_linear_schedule_with_warmup(
                pre_optimizer, num_warmup_steps=warmup_steps, num_training_steps=t_total
            )
            for epoch in range(args.num_pretrain_epochs):
                mlm_pretrain(args, model, corpus_dataloader, pre_optimizer, pre_scheduler, tokenizer, loss_fnc_pretrain, epoch)
        elif args.tokenizer == 'wordpiece':
            corpus_dataset = TextDataset(corpuses)
            corpus_dataloader = DataLoader(
                corpus_dataset,
                sampler=RandomSampler(corpus_dataset),
                batch_size=args.train_batch_size,
                num_workers=args.num_workers,
                collate_fn=lambda x: collate_fn_mlm_pretrain_transformer(x, tokenizer, args.max_seq_length),
            )
            t_total = len(corpus_dataloader) * args.num_pretrain_epochs
            warmup_steps = t_total * args.warmup_ratio
            pre_scheduler = get_linear_schedule_with_warmup(
                pre_optimizer, num_warmup_steps=warmup_steps, num_training_steps=t_total
            )
            for epoch in range(args.num_pretrain_epochs):
                mlm_pretrain(args, model, corpus_dataloader, pre_optimizer, pre_scheduler, tokenizer, loss_fnc_pretrain, epoch)

    step = 0
    if index is None or index == 0:
        best_val_gleu = -float("inf")
    meter = Meter()

    for epoch in range(1, epochs + 1):
        for batch in train_dataloader:
            model.train()
            step += 1
            batch = tuple(t.to(args.device) for t in batch)

            if step < 4:
                loss, items = calc_loss1(model, batch, args)
            else:
                loss, items = calc_loss(model, batch, args)

            if step < 4:
                print(f"first item loss is {items[0]}")
                print(f"second item bsz is {items[1]}")
                print(f"third item target mask sum is {items[2]}")
                print()

            meter.add(*items)

            loss.backward()
            if args.max_grad_norm > 0:
                nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
            optimizer.step()
            model.zero_grad()
            scheduler.step()

            if step % args.log_interval == 0:
                lr = scheduler.get_lr()[0]
                loss_sent, loss_token = meter.average()

                logger.info(f' [{step:5d}] lr {lr:.6f} | {meter.print_str(True)}')
                nsml.report(
                    step=step,
                    scope=locals(),
                    summary=True,
                    train__lr=lr,
                    train__loss_sent=loss_sent,
                    train__token_ppl=math.exp(loss_token)
                )
                meter.init()

            if step % args.eval_interval == 0 and valid_data is not None:
                start_eval = time.time()
                (val_loss, val_loss_token), valid_str = evaluate(model, valid_dataloader, args)
                prediction = correct(model, tokenizer, valid_noisy, args)

                for (i, (pred, clean)) in enumerate(zip(prediction, valid_clean)):
                    if i == 5:
                        break
                    print(f"{i + 1}'s pred sentence is {pred}, clean sentence is {clean}")

                val_em = em(prediction, valid_clean)
                val_gleu = gleu(prediction, valid_clean)

                logger.info('-' * 89)
                logger.info(f' [{step:6d}] valid | {valid_str} | em {val_em:5.2f} | gleu {val_gleu:5.2f}')
                logger.info('-' * 89)
                nsml.report(
                    step=step,
                    scope=locals(),
                    summary=True,
                    valid__loss_sent=val_loss,
                    valid__token_ppl=math.exp(val_loss_token),
                    valid__em=val_em,
                    valid__gleu=val_gleu
                )

                if val_gleu > best_val_gleu:
                    if args.kfold:
                        if index == args.n_splits - 1:
                            print('------------save model------------')
                            nsml.save(f"model_{epoch}")
                    else:
                        print('------------save model------------')
                        best_val_gleu = val_gleu
                        nsml.save("best")
                meter.start += time.time() - start_eval

            elif step % args.eval_interval == 0 and valid_data is None:
                if epoch >= 5:
                    print('------------save model------------')
                    nsml.save(f"model_{epoch}")

            if step >= args.max_steps:
                break
        if step >= args.max_steps:
            break

def bind_nsml(model, tokenizer=None, args=None):
    def save(path, **kwargs):
        torch.save(model.state_dict(), open(os.path.join(path, 'model.pt'), 'wb'))
        if tokenizer is not None and args.tokenizer in ['char', 'pre_char', 'kfold_char']:
            tokenizer.save(os.path.join(path, 'vocab.txt'))
        elif args.tokenizer == 'wordpiece':
            tokenizer.save(os.path.join(path, 'wordpiece_tokenizer'))

    def load(path, **kwargs):
        model.load_state_dict(torch.load(open(os.path.join(path, 'model.pt'), 'rb'),
                                         map_location=lambda storage, loc: storage))
        if tokenizer is not None and args.tokenizer in ['char', 'pre_char', 'kfold_char']:
            tokenizer.load(os.path.join(path, 'vocab.txt'))
        elif args.tokenizer == 'wordpiece':
            tokenizer.load(os.path.join(path, 'wordpiece_tokenizer'))

    def infer(test_data, **kwargs):
        '''
        :param test_data: list of noisy sentences
        :return: list of corrected sentences
        '''
        return correct(model, tokenizer, test_data, args)

    import nsml
    nsml.bind(save, load, infer)

def main():
    args = get_args()
    logger.info(f"args: {json.dumps(args.__dict__, indent=2, sort_keys=True)}")

    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    args.n_gpu = torch.cuda.device_count()
    set_seed(args)

    if args.tokenizer in ['char', 'pre_char', 'kfold_char']:
        tokenizer = CharTokenizer([])
    elif args.tokenizer == 'dsksd':
        tokenizer = BertTokenizer.from_pretrained('dsksd/bert-ko-small-minimal')
        print(f"vocab size is {len(tokenizer)}")
        args.vocab_size = len(tokenizer)
    elif args.tokenizer == 'plus':
        tokenizer = LarvaTokenizer.from_pretrained("larva-kor-plus-base-cased")
        print(f"vocab size is {len(tokenizer)}")
        args.vocab_size = len(tokenizer)
    elif args.tokenizer == 'wordpiece':
        tokenizer = CustomTokenizer(args)

    if args.model == 'transformer':
        model = TransformerModel(
            vocab_size=args.vocab_size,
            hidden_size=args.hidden_size,
            num_attention_heads=args.num_attention_heads,
            num_encoder_layers=args.num_encoder_layers,
            num_decoder_layers=args.num_decoder_layers,
            intermediate_size=args.intermediate_size,
            dropout=args.dropout,
        ).to(args.device)
    elif args.model == 'encoderdecoder':
        model = EncoderDecoder(
            vocab_size=args.vocab_size,
            hidden_size=args.hidden_size,
            num_attention_heads=args.num_attention_heads,
            num_encoder_layers=args.num_encoder_layers,
            num_decoder_layers=args.num_decoder_layers,
            intermediate_size=args.intermediate_size,
            dropout=args.dropout,
        ).to(args.device)
    elif args.model == 'BertGpt2':
        model = BertGpt2(
            vocab_size=args.vocab_size,
            hidden_size=args.hidden_size,
            num_attention_heads=args.num_attention_heads,
            num_encoder_layers=args.num_encoder_layers,
            num_decoder_layers=args.num_decoder_layers,
            intermediate_size=args.intermediate_size,
            dropout=args.dropout,
        ).to(args.device)
    elif args.model == 'bertfuse':
        model = TransformerS2Model(tokenizer, args=args).to(args.device)

    logger.info(f"# of model parameters: {sum(p.numel() for p in model.parameters()) * 1e-6:.2f}M")

    bind_nsml(model, tokenizer, args)
    if args.pause:
        nsml.paused(scope=locals())

    if args.mode == "train":
        noisy_sents = read_strings(os.path.join(args.data_dir, "train_data", "train_data"))
        annotations = read_strings(os.path.join(args.data_dir, "train_data", "train_annotation"))
        corpuses = read_strings(os.path.join(args.data_dir, "train_data", "train_corpus"))
        clean_sents = read_strings(os.path.join(args.data_dir, "train_label"))

        print(f"noisy_sents length is {len(noisy_sents)}")
        print(f"annotations length is {len(annotations)}")
        print(f"clean_sents length is {len(clean_sents)}")
        print(f"corpuses length is {len(corpuses)}")

        for (i, noisy_sent) in enumerate(noisy_sents[:5]):
            print(f"{i + 1}'s noisy sentence is {noisy_sent}")

        for (i, clean_sent) in enumerate(clean_sents[:5]):
            print(f"{i + 1}'s clean sentence is {clean_sent}")

        for (i, annotation) in enumerate(annotations[:5]):
            print(f"{i + 1}'s annotation is {annotation}")

        pairs = [{"noisy": noisy, "clean": clean} for noisy, clean in zip(noisy_sents, clean_sents)]

        if args.train_all:
            train_data = pairs
            valid_data = None
            logger.info(f"# of train data: {len(train_data)}")
        else:
            train_data, valid_data = pairs[:-int(len(pairs) * args.val_ratio)], pairs[-int(len(pairs) * args.val_ratio):]
            logger.info(f"# of train data: {len(train_data)}")
            logger.info(f"# of valid data: {len(valid_data)}")

        print(f"train_data is {train_data[0]}")

        train_sents = [x['noisy'] for x in train_data] + [x['clean'] for x in train_data]

        print(f"first train_sentences is {train_sents[0]}")
        print(f"last train_sentences is {train_sents[-1]}")

        if args.tokenizer == 'char':
            tokenizer = CharTokenizer.from_strings(train_sents, args.vocab_size)
        elif args.tokenizer == 'pre_char':
            tokenizer = CharTokenizer.from_strings(corpuses + clean_sents + noisy_sents, args.vocab_size)
        elif args.tokenizer == 'kfold_char':
            tokenizer = CharTokenizer.from_strings(clean_sents + noisy_sents, args.vocab_size)
        elif args.tokenizer == 'wordpiece':
            tokenizer.get_tokenizer()
            print(f"vocab size is {len(tokenizer)}")
            args.vocab_size = len(tokenizer)

        bind_nsml(model, tokenizer, args)

    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model, dim=1)

    if args.mode == "train":
        if not args.kfold:
            train(model, tokenizer, train_data, valid_data, corpuses, args)
        else:
            kfold = KFold(n_splits=args.n_splits)
            for index, (train_idx, valid_idx) in enumerate(kfold.split(range(len(pairs)))):
                train_data, valid_data = [pairs[idx] for idx in train_idx], [pairs[idx] for idx in valid_idx]
                train(model, tokenizer, train_data, valid_data, corpuses, args, index)

if __name__ == "__main__":
    main()