#!/usr/bin/env python
"""
    Main training workflow
"""
from __future__ import division

import argparse
import glob
import os
import random
import signal
import time
import warnings
os.environ["CUDA_VISIBLE_DEVICES"] = '3'
warnings.filterwarnings('ignore')
import torch
import distributed
from transformers import AlbertConfig
from transformers import BertConfig
from models import data_loader, model_builder, Aldata_loader,clf_loader,Com_loader,Pointer_loader
# from models.Aldata_loader import load_dataset
from models.clf_loader import load_dataset
from models.model_builder import Summarizer,GCNSummarizer,AlbertSummarizer,ClfSummarizer,ComSummarizer,PointerSummarizer
from models.trainer import build_trainer
from others.logging import logger, init_logger
# python train.py -mode train -encoder classifier -dropout 0.1 -bert_data_path ../clf_bert/cnndm -model_path ../models/bert_re_seg/ -lr 2e-3 -visible_gpus 0,1,2,3 -gpu_ranks 0,1,2,3 -world_size 4 -report_every 50 -save_checkpoint_steps 1000 -batch_size 5000 -decay_method noam -train_steps 50000 -accum_count 6 -log_file ../logs/bert_re_seg -use_interval True -warmup_steps 10000 -process False -model clf_albert -reproduce True

model_flags = ['hidden_size', 'ff_size', 'heads', 'inter_layers','encoder','ff_actv', 'use_interval','rnn_size']


def str2bool(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 multi_main(args):
    """ Spawns 1 process per GPU """
    init_logger()

    nb_gpu = args.world_size
    mp = torch.multiprocessing.get_context('spawn')

    # Create a thread to listen for errors in the child processes.
    error_queue = mp.SimpleQueue()
    error_handler = ErrorHandler(error_queue)

    # Train with multiprocessing.
    procs = []
    for i in range(nb_gpu):
        device_id = i
        procs.append(mp.Process(target=run, args=(args,
            device_id, error_queue,), daemon=True))
        procs[i].start()
        logger.info(" Starting process pid: %d  " % procs[i].pid)
        error_handler.add_child(procs[i].pid)
    for p in procs:
        p.join()



def run(args, device_id, error_queue):

    """ run process """
    setattr(args, 'gpu_ranks', [int(i) for i in args.gpu_ranks])

    try:
        gpu_rank = distributed.multi_init(device_id, args.world_size, args.gpu_ranks)
        print('gpu_rank %d' %gpu_rank)
        if gpu_rank != args.gpu_ranks[device_id]:
            raise AssertionError("An error occurred in \
                  Distributed initialization")

        train(args,device_id)
    except KeyboardInterrupt:
        pass  # killed by parent, do nothing
    except Exception:
        # propagate exception to parent process, keeping original traceback
        import traceback
        error_queue.put((args.gpu_ranks[device_id], traceback.format_exc()))


class ErrorHandler(object):
    """A class that listens for exceptions in children processes and propagates
    the tracebacks to the parent process."""

    def __init__(self, error_queue):
        """ init error handler """
        import signal
        import threading
        self.error_queue = error_queue
        self.children_pids = []
        self.error_thread = threading.Thread(
            target=self.error_listener, daemon=True)
        self.error_thread.start()
        signal.signal(signal.SIGUSR1, self.signal_handler)

    def add_child(self, pid):
        """ error handler """
        self.children_pids.append(pid)

    def error_listener(self):
        """ error listener """
        (rank, original_trace) = self.error_queue.get()
        self.error_queue.put((rank, original_trace))
        os.kill(os.getpid(), signal.SIGUSR1)

    def signal_handler(self, signalnum, stackframe):
        """ signal handler """
        for pid in self.children_pids:
            os.kill(pid, signal.SIGINT)  # kill children processes
        (rank, original_trace) = self.error_queue.get()
        msg = """\n\n-- Tracebacks above this line can probably
                 be ignored --\n\n"""
        msg += original_trace
        raise Exception(msg)



def wait_and_validate(args, device_id):

    timestep = 0
    if (args.test_all):
        cp_files = sorted(glob.glob(os.path.join(args.model_path, 'model_step_*.pt')))
        cp_files.sort(key=os.path.getmtime)
        xent_lst = []
        for i, cp in enumerate(cp_files):
            step = int(cp.split('.')[-2].split('_')[-1])
            if (args.test_start_from != -1 and step < args.test_start_from):
                xent_lst.append((1e6, cp))
                continue
            xent = validate(args,  device_id, cp, step)
            xent_lst.append((xent, cp))
            max_step = xent_lst.index(min(xent_lst))
            if (i - max_step > 10):
                break

        xent_lst = sorted(xent_lst, key=lambda x: x[0])[:3]
        logger.info('PPL %s' % str(xent_lst))
        for xent, cp in xent_lst:
            step = int(cp.split('.')[-2].split('_')[-1])
            args.train=False
            test(args,  device_id, cp, step)
    else:
        while (True):
            cp_files = sorted(glob.glob(os.path.join(args.model_path, 'model_step_*.pt')))
            cp_files.sort(key=os.path.getmtime)
            if (cp_files):
                cp = cp_files[-1]
                time_of_cp = os.path.getmtime(cp)
                if (not os.path.getsize(cp) > 0):
                    time.sleep(60)
                    continue
                if (time_of_cp > timestep):
                    timestep = time_of_cp
                    step = int(cp.split('.')[-2].split('_')[-1])
                    validate(args,  device_id, cp, step)
                    test(args,  device_id, cp, step)

            cp_files = sorted(glob.glob(os.path.join(args.model_path, 'model_step_*.pt')))
            cp_files.sort(key=os.path.getmtime)
            if (cp_files):
                cp = cp_files[-1]
                time_of_cp = os.path.getmtime(cp)
                if (time_of_cp > timestep):
                    continue
            else:
                time.sleep(300)


def validate(args,  device_id, pt, step):
    device = "cpu" if args.visible_gpus == '-1' else "cuda"
    if (pt != ''):
        test_from = pt
    else:
        test_from = args.test_from
    logger.info('Loading checkpoint from %s' % test_from)
    checkpoint = torch.load(test_from, map_location=lambda storage, loc: storage)
    opt = vars(checkpoint['opt'])
    for k in opt.keys():
        if (k in model_flags):
            setattr(args, k, opt[k])
    print(args)
    if args.model=='bert':
        config = BertConfig.from_json_file(args.bert_config_path)#from_pretrained
        model = Summarizer(args, device, load_pretrained_bert=False, bert_config = config)#Albert
        valid_iter = data_loader.Dataloader(args, load_dataset(args, 'valid', shuffle=False),
                                            args.batch_size, device,
                                            shuffle=False, is_test=False)
    elif args.model=='albert':
        config = AlbertConfig.from_pretrained('albert-base-v2')  # from_pretrained
        model = AlbertSummarizer(args, device, load_pretrained_bert=False, Albert_config=config)  #
        valid_iter = Aldata_loader.Dataloader(args, load_dataset(args, 'valid', shuffle=False),
                                            args.batch_size, device,
                                            shuffle=False, is_test=False)
    elif args.model=='pointer':
        config = AlbertConfig.from_pretrained('albert-base-v2')  # from_pretrained
        model = PointerSummarizer(args, device, load_pretrained_bert=False, Albert_config=config)  #
        valid_iter = Pointer_loader.Dataloader(args, load_dataset(args, 'valid', shuffle=False),
                                            args.batch_size, device,
                                            shuffle=False, is_test=False)
    elif args.model=='albert_com':
        config = AlbertConfig.from_pretrained('albert-base-v2')  # from_pretrained
        model = ComSummarizer(args, device, load_pretrained_bert=False, Albert_config=config)  #
        valid_iter = Com_loader.Dataloader(args, load_dataset(args, 'valid', shuffle=False),
                                            args.batch_size, device,
                                            shuffle=False, is_test=False)
    elif args.model=='gcn':
        config = AlbertConfig.from_pretrained('albert-base-v2')  # from_pretrained
        model = GCNSummarizer(args, device, load_pretrained_bert=False, Albert_config=config)  # Albert
        valid_iter = Aldata_loader.Dataloader(args, load_dataset(args, 'valid', shuffle=False),
                                            args.batch_size, device,
                                            shuffle=False, is_test=False)
    elif args.model=='clf_albert':
        config = AlbertConfig.from_pretrained('albert-base-v2')  # from_pretrained
        # model = ClfSummarizer(args, device, load_pretrained_bert=False, Albert_config=config)
        if args.reproduce:
            model = ReSummarizer(args, device, load_pretrained_bert=True)
        else:
            model = ClfSummarizer(args, device, load_pretrained_bert=True)
        valid_iter = clf_loader.Dataloader(args, load_dataset(args, 'valid', shuffle=False),
                                            args.batch_size, device,
                                            shuffle=False, is_test=False)
    model.load_cp(checkpoint)
    model.eval()
    trainer = build_trainer(args, device_id, model, None)
    stats = trainer.validate(valid_iter, step)
    return stats.xent()

def test(args, device_id, pt, step):

    device = "cpu" if args.visible_gpus == '-1' else "cuda"
    if (pt != ''):
        test_from = pt
    else:
        test_from = args.test_from
    logger.info('Loading checkpoint from %s' % test_from)
    checkpoint = torch.load(test_from, map_location=lambda storage, loc: storage)
    opt = vars(checkpoint['opt'])
    for k in opt.keys():
        if (k in model_flags):
            setattr(args, k, opt[k])
    print(args)
    if args.model=='bert':
        config = BertConfig.from_json_file(args.bert_config_path)#from_pretrained
        test_iter = data_loader.Dataloader(args, load_dataset(args, 'test', shuffle=False),
                                           args.batch_size, device,
                                           shuffle=False, is_test=True)
        model = Summarizer(args, device, load_pretrained_bert=False, bert_config = config)#Albert
    elif args.model=='albert':
        config = AlbertConfig.from_pretrained('albert-base-v2')  # from_pretrained
        model = AlbertSummarizer(args, device, load_pretrained_bert=False, Albert_config=config)  #
        test_iter = Aldata_loader.Dataloader(args, load_dataset(args, 'test', shuffle=False),
                                           args.batch_size, device,
                                           shuffle=False, is_test=True)
    elif args.model=='albert_com':
        config = AlbertConfig.from_pretrained('albert-base-v2')  # from_pretrained
        model = ComSummarizer(args, device, load_pretrained_bert=False, Albert_config=config)  #
        test_iter = Com_loader.Dataloader(args, load_dataset(args, 'test', shuffle=False),
                                           args.batch_size, device,
                                           shuffle=False, is_test=True)
    elif args.model=='gcn':
        config = AlbertConfig.from_pretrained('albert-base-v2')   # from_pretrained
        model = GCNSummarizer(args, device, load_pretrained_bert=False, Albert_config=config)  # Albert
        test_iter = Aldata_loader.Dataloader(args, load_dataset(args, 'test', shuffle=False),
                                           args.batch_size, device,
                                           shuffle=False, is_test=True)
    elif args.model=='pointer':
        config = AlbertConfig.from_pretrained('albert-base-v2')   # from_pretrained
        model = PointerSummarizer(args, device, load_pretrained_bert=False, Albert_config=config)  # Albert
        test_iter = Pointer_loader.Dataloader(args, load_dataset(args, 'test', shuffle=False),
                                           args.batch_size, device,
                                           shuffle=False, is_test=True)
    elif args.model == 'clf_albert':
        # model = ClfSummarizer(args, device, load_pretrained_bert=True)
        model = ClfSummarizer(args, device, load_pretrained_bert=True)
        test_iter = clf_loader.Dataloader(args, load_dataset(args, 'test', shuffle=False),
                                           args.batch_size, device,
                                           shuffle=False, is_test=True)
    print('-'*30)
    print(model)
    model.load_cp(checkpoint)
    model.eval()

    #Al
    trainer = build_trainer(args, device_id, model, None)
    trainer.test(test_iter,step)

def process(args, device_id, pt, step):

    device = "cpu" if args.visible_gpus == '-1' else "cuda"
    if (pt != ''):
        test_from = pt
    else:
        test_from = args.test_from
    logger.info('Loading checkpoint from %s' % test_from)
    checkpoint = torch.load(test_from, map_location=lambda storage, loc: storage)
    opt = vars(checkpoint['opt'])
    for k in opt.keys():
        if (k in model_flags):
            setattr(args, k, opt[k])
    print(args)
    if args.model=='bert':
        config = BertConfig.from_json_file(args.bert_config_path)#from_pretrained
        model = Summarizer(args, device, load_pretrained_bert=False, bert_config = config)#Albert
    elif args.model=='albert':
        config = AlbertConfig.from_pretrained('albert-base-v2')  # from_pretrained
        model = AlbertSummarizer(args, device, load_pretrained_bert=False, Albert_config=config)  #
    elif args.model == 'albert_com':
        config = AlbertConfig.from_pretrained('albert-base-v2')  # from_pretrained
        model = ComSummarizer(args, device, load_pretrained_bert=False, Albert_config=config)
    elif args.model=='gcn':
        config = AlbertConfig.from_pretrained('albert-base-v2')   # from_pretrained
        model = GCNSummarizer(args, device, load_pretrained_bert=False, Albert_config=config)  # Albert
    model.load_cp(checkpoint)
    model.eval()

    test_iter =Aldata_loader.Dataloader(args, load_dataset(args, 'train', shuffle=False),
                                  args.batch_size, device,
                                  shuffle=False, is_test=True)#Al
    trainer = build_trainer(args, device_id, model, None)
    trainer.process(test_iter,step)


def baseline(args, cal_lead=False, cal_oracle=False):

    test_iter =data_loader.Dataloader(args, load_dataset(args, 'test', shuffle=False),
                                  args.batch_size, device,
                                  shuffle=False, is_test=True)#Al

    trainer = build_trainer(args, device_id, None, None)
    #
    if (cal_lead):
        trainer.test(test_iter, 0, cal_lead=True)
    elif (cal_oracle):
        trainer.test(test_iter, 0, cal_oracle=True)


def train(args, device_id):
    init_logger(args.log_file)

    device = "cpu" if args.visible_gpus == '-1' else "cuda"
    logger.info('Device ID %d' % device_id)
    logger.info('Device %s' % device)
    torch.manual_seed(args.seed)
    random.seed(args.seed)
    torch.backends.cudnn.deterministic = True

    if device_id >= 0:
        torch.cuda.set_device(device_id)
        torch.cuda.manual_seed(args.seed)


    torch.manual_seed(args.seed)
    random.seed(args.seed)
    torch.backends.cudnn.deterministic = True

    def train_iter_fct():
        if args.model=='clf_albert':
            return clf_loader.Dataloader(args, load_dataset(args, 'train', shuffle=True), args.batch_size, device,
                                                 shuffle=True, is_test=False)
        elif args.model=='albert_com':
            return Com_loader.Dataloader(args, load_dataset(args, 'train', shuffle=True), args.batch_size, device,
                                                 shuffle=True, is_test=False)
        elif args.model=='albert':
            return Aldata_loader.Dataloader(args, load_dataset(args, 'train', shuffle=True), args.batch_size, device,
                                                 shuffle=True, is_test=False)
        elif args.model=='pointer':
            return Pointer_loader.Dataloader(args, load_dataset(args, 'train', shuffle=True), args.batch_size, device,
                                                 shuffle=True, is_test=False)
        else:
            return data_loader.Dataloader(args, load_dataset(args, 'train', shuffle=True), args.batch_size, device,
                                            shuffle=True, is_test=False)
    if args.model=='bert':
        model = Summarizer(args, device, load_pretrained_bert=True)#Albert
    elif args.model=='albert':
        model = AlbertSummarizer(args, device, load_pretrained_bert=True)
    elif args.model=='pointer':
        model = PointerSummarizer(args, device, load_pretrained_bert=True)
    elif args.model=='albert_com':
        model = ComSummarizer(args, device, load_pretrained_bert=True)
    elif args.model=='gcn':
        model = GCNSummarizer(args, device, load_pretrained_bert=True)  # Albert
    elif args.model == 'clf_albert':
        if args.reproduce:
            model = ReSummarizer(args,device,load_pretrained_bert=True)
        else:
            model = ClfSummarizer(args, device, load_pretrained_bert=True)
    if args.train_from != '':
        logger.info('Loading checkpoint from %s' % args.train_from)
        checkpoint = torch.load(args.train_from,
                                map_location=lambda storage, loc: storage)
        opt = vars(checkpoint['opt'])
        for k in opt.keys():
            if (k in model_flags):
                setattr(args, k, opt[k])
        model.load_cp(checkpoint)
        optim = model_builder.build_optim(args, model, checkpoint)
    else:
        optim = model_builder.build_optim(args, model, None)

    logger.info(model)
    trainer = build_trainer(args, device_id, model, optim)
    trainer.train(train_iter_fct, args.train_steps)



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



    parser.add_argument("-encoder", default='classifier', type=str, choices=['classifier','transformer','rnn','baseline'])
    parser.add_argument("-model", default='bert', type=str, choices=['bert','gcn','albert','clf_albert','albert_com','pointer'])
    parser.add_argument("-mode", default='train', type=str, choices=['train','validate','test','lead','oracle'])
    parser.add_argument("-bert_data_path", default='/home/gwy/ALL/gcn_data/cnndm')
    parser.add_argument("-model_path", default='../GCNmodels/')
    parser.add_argument("-result_path", default='../results/GCNcnndm')
    parser.add_argument("-temp_dir", default='../temp')
    parser.add_argument("-bert_config_path", default='../bert_config_uncased_base.json')

    parser.add_argument("-batch_size", default=1000, type=int)

    parser.add_argument("-use_interval", type=str2bool, nargs='?',const=True,default=True)
    parser.add_argument("-truncate", type=str2bool, nargs='?',const=True,default=True)
    parser.add_argument("-hidden_size", default=768, type=int)
    parser.add_argument("-ff_size", default=512, type=int)
    parser.add_argument("-heads", default=4, type=int)
    parser.add_argument("-inter_layers", default=2, type=int)
    parser.add_argument("-rnn_size", default=768, type=int)

    parser.add_argument("-param_init", default=0, type=float)
    parser.add_argument("-param_init_glorot", type=str2bool, nargs='?',const=True,default=True)
    parser.add_argument("-dropout", default=0, type=float)
    parser.add_argument("-optim", default='adam', type=str)
    parser.add_argument("-lr", default=1, type=float)
    parser.add_argument("-beta1", default= 0.9, type=float)
    parser.add_argument("-beta2", default=0.999, type=float)
    parser.add_argument("-rate", default=0.2, type=float)
    parser.add_argument("-decay_method", default='', type=str)
    parser.add_argument("-warmup_steps", default=8000, type=int)
    parser.add_argument("-max_grad_norm", default=0, type=float)

    parser.add_argument("-save_checkpoint_steps", default=5, type=int)
    parser.add_argument("-accum_count", default=1, type=int)
    parser.add_argument("-world_size", default=1, type=int)
    parser.add_argument("-report_every", default=1, type=int)
    parser.add_argument("-train_steps", default=1000, type=int)
    parser.add_argument("-recall_eval", type=str2bool, nargs='?',const=True,default=False)
    parser.add_argument("-reproduce", type=str2bool, nargs='?',const=True,default=True)
    parser.add_argument("-test_start_from", default=-1, type=int)


    parser.add_argument('-visible_gpus', default='-1', type=str)
    parser.add_argument('-gpu_ranks', default='0', type=str)
    parser.add_argument('-log_file', default='../logs/cnndm.log')
    parser.add_argument('-dataset', default='')
    parser.add_argument('-seed', default=666, type=int)

    parser.add_argument("-test_all", type=str2bool, nargs='?',const=True,default=False)
    parser.add_argument("-test_from", default='')
    parser.add_argument("-train_from", default='')
    parser.add_argument("-report_rouge", type=str2bool, nargs='?',const=True,default=True)
    parser.add_argument("-train", type=str2bool, nargs='?',const=True,default=False)
    parser.add_argument("-block_trigram", type=str2bool, nargs='?', const=True, default=True)
    parser.add_argument("-process", type=str2bool, nargs='?',const=True,default=False)
    parser.add_argument("-hidden_dim", default=128, type=int)
    parser.add_argument("-out_channel", default=64, type=int)
    args = parser.parse_args()
    args.gpu_ranks = [int(i) for i in args.gpu_ranks.split(',')]
    os.environ["CUDA_VISIBLE_DEVICES"] = args.visible_gpus

    init_logger(args.log_file)
    device = "cpu" if args.visible_gpus == '-1' else "cuda"
    device_id = 0 if device == "cuda" else -1

    if(args.world_size>1):
        multi_main(args)
    elif (args.mode == 'train'):
        train(args, device_id)
    elif (args.mode == 'validate'):
        wait_and_validate(args, device_id)
    elif (args.mode == 'lead'):
        baseline(args, cal_lead=True)
    elif (args.mode == 'oracle'):
        baseline(args, cal_oracle=True)
    elif (args.mode == 'test'):
        cp = args.test_from
        try:
            step = int(cp.split('.')[-2].split('_')[-1])
        except:
            step = 0
        if args.process:
            pts = sorted(glob.glob(args.bert_data_path + '.' + 'train' + '.[0-9]*.pt'))
            # t=int(len(pts))
            for i in pts:
                args.bert_data_path=i
                print(args.bert_data_path)
                name = args.bert_data_path.split('/')[-1]
                if not os.path.exists(os.path.join('/home/gwy/ALL/BertSum/clf_data', name)):
                    name = args.bert_data_path.split('/')[-1]
                    torch.save({'a':1}, os.path.join('/home/gwy/ALL/BertSum/clf_data', name))
                    process(args, device_id, cp, step)
                # args.bert_data_path = i
                # print(args.bert_data_path)
                # name = args.bert_data_path.split('/')[-1]
                # if not os.path.exists(os.path.join('/home/gwy/ALL/BertSum/clf_data', name)):
                #     process(args, device_id, cp, step)
        else:
            test(args, device_id, cp, step)
