import pprint

import mxnet as mx
from mxnet import gluon
from mxnet import init

from lib.core.get_optimizer import *
from lib.core.metric import MPJPEMetric
from lib.core.loss import MeanSquareLoss
from lib.core.loader import JointsDataIter
from lib.network import get_net
from lib.net_module import *
from lib.utils import *
from lib.dataset.hm36 import hm36

from config import config, gen_config, update_config_from_args, s_args
config = update_config_from_args(config, s_args)

def main():
    # Parse config and mkdir output
    logger, final_Model_path = create_logger(config)
    config.final_Model_path = final_Model_path
    gen_config(os.path.join(final_Model_path, 'hyperParams.yaml'))
    logger.info('Training config:{}\n'.format(pprint.pformat(config)))

    # define context
    if config.useGPU:
        ctx = [mx.gpu(int(i)) for i in config.gpu.split(',')]
    else:
        ctx = mx.cpu()
    logger.info("Using context:", ctx)

    # dataset, generate trainset/ validation set
    train_imdbs = []
    valid_imdbs = []
    for i in range(len(config.DATASET.train_image_set)):
        logger.info("Construct Dataset:", config.DATASET.dbname[i], ", Dataset Path:", config.DATASET.dataset_path[i])
        train_imdbs.append(eval(config.DATASET.dbname[i])(config.DATASET.train_image_set[i],
                                                          config.DATASET.root_path[i],
                                                          config.DATASET.dataset_path[i]))
        valid_imdbs.append(eval(config.DATASET.dbname[i])(config.DATASET.valid_image_set[i],
                                                          config.DATASET.root_path[i],
                                                          config.DATASET.dataset_path[i],
                                                          config.final_Model_path))
    data_names  = ['hm36data']
    label_names = ['hm36label']
    train_data_iter = JointsDataIter(train_imdbs[0], runmode=0,
                                    data_names = data_names, label_names=label_names,
                                    shuffle=config.TRAIN.SHUFFLE, batch_size=len(ctx)*config.TRAIN.batchsize, logger=logger)
    valid_data_iter = JointsDataIter(valid_imdbs[0], runmode=1,
                                    data_names = data_names, label_names=label_names,
                                    shuffle=False, batch_size=len(ctx)*config.TEST.batchsize, logger=logger)

    assert train_data_iter.get_meanstd()['mean3d'].all() == valid_data_iter.get_meanstd()['mean3d'].all()

    # network
    net = get_net(config)
    if config.resume:
        ckp_path = os.path.join(config.resumeckp)
        net.collect_params().load(ckp_path, ctx=ctx)
    else:
        net.initialize(init=init.MSRAPrelu(), ctx=ctx)

    if config.NETWORK.hybrid:
        net.hybridize()

    logger.info(net)

    # define loss and metric
    mean3d = train_data_iter.get_meanstd()['mean3d']
    std3d  = train_data_iter.get_meanstd()['std3d']
    train_metric = MPJPEMetric('train_metric', mean3d, std3d)
    eval_metric  = MPJPEMetric('valid_metric', mean3d, std3d)
    loss         = MeanSquareLoss()

    # optimizer
    optimizer, optimizer_params = get_optimizer(config, ctx)

    # train and valid
    TrainDBsize = train_data_iter.get_size()
    ValidDBsize = valid_data_iter.get_size()
    logger.info("Train DB size:", TrainDBsize, "Valid DB size:",ValidDBsize)

    if not isinstance(train_data_iter, mx.io.PrefetchingIter):
        train_data_iter = mx.io.PrefetchingIter(train_data_iter)

    trainer = gluon.Trainer(net.collect_params(), optimizer, optimizer_params)
    for epoch in range(config.TRAIN.begin_epoch, config.TRAIN.end_epoch):
        trainNet(net, trainer, train_data_iter, loss, train_metric, epoch, config, logger=logger, ctx=ctx)
        validNet(net, valid_data_iter, loss, eval_metric, epoch, config, logger=logger, ctx=ctx)

    logger.kill()

if __name__ == '__main__':
    main()
[(17, 'config.update_config_from_args', 'update_config_from_args', False, 'from config import config, gen_config, update_config_from_args, s_args\n'), (57, 'lib.network.get_net', 'get_net', False, 'from lib.network import get_net\n'), (72, 'lib.core.metric.MPJPEMetric', 'MPJPEMetric', False, 'from lib.core.metric import MPJPEMetric\n'), (73, 'lib.core.metric.MPJPEMetric', 'MPJPEMetric', False, 'from lib.core.metric import MPJPEMetric\n'), (74, 'lib.core.loss.MeanSquareLoss', 'MeanSquareLoss', False, 'from lib.core.loss import MeanSquareLoss\n'), (30, 'mxnet.cpu', 'mx.cpu', True, 'import mxnet as mx\n'), (85, 'mxnet.io.PrefetchingIter', 'mx.io.PrefetchingIter', True, 'import mxnet as mx\n'), (24, 'pprint.pformat', 'pprint.pformat', False, 'import pprint\n'), (62, 'mxnet.init.MSRAPrelu', 'init.MSRAPrelu', False, 'from mxnet import init\n')]