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, 9, 17, 48), 'config.update_config_from_args', 'update_config_from_args', ({(17, 33, 17, 39): 'config', (17, 41, 17, 47): 's_args'}, {}), '(config, s_args)', False, 'from config import config, gen_config, update_config_from_args, s_args\n'), ((57, 10, 57, 25), 'lib.network.get_net', 'get_net', ({(57, 18, 57, 24): 'config'}, {}), '(config)', False, 'from lib.network import get_net\n'), ((72, 19, 72, 61), 'lib.core.metric.MPJPEMetric', 'MPJPEMetric', ({(72, 31, 72, 45): '"""train_metric"""', (72, 47, 72, 53): 'mean3d', (72, 55, 72, 60): 'std3d'}, {}), "('train_metric', mean3d, std3d)", False, 'from lib.core.metric import MPJPEMetric\n'), ((73, 19, 73, 61), 'lib.core.metric.MPJPEMetric', 'MPJPEMetric', ({(73, 31, 73, 45): '"""valid_metric"""', (73, 47, 73, 53): 'mean3d', (73, 55, 73, 60): 'std3d'}, {}), "('valid_metric', mean3d, std3d)", False, 'from lib.core.metric import MPJPEMetric\n'), ((74, 19, 74, 35), 'lib.core.loss.MeanSquareLoss', 'MeanSquareLoss', ({}, {}), '()', False, 'from lib.core.loss import MeanSquareLoss\n'), ((30, 14, 30, 22), 'mxnet.cpu', 'mx.cpu', ({}, {}), '()', True, 'import mxnet as mx\n'), ((85, 26, 85, 64), 'mxnet.io.PrefetchingIter', 'mx.io.PrefetchingIter', ({(85, 48, 85, 63): 'train_data_iter'}, {}), '(train_data_iter)', True, 'import mxnet as mx\n'), ((24, 46, 24, 68), 'pprint.pformat', 'pprint.pformat', ({(24, 61, 24, 67): 'config'}, {}), '(config)', False, 'import pprint\n'), ((62, 28, 62, 44), 'mxnet.init.MSRAPrelu', 'init.MSRAPrelu', ({}, {}), '()', False, 'from mxnet import init\n')]