# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import argparse
import logging
import time
import mxnet as mx
import numpy as np
from get_data import get_movielens_iter, get_movielens_data
from model import matrix_fact_model_parallel_net


logging.basicConfig(level=logging.DEBUG)

parser = argparse.ArgumentParser(description="Run model parallel version of matrix factorization",
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--num-epoch', type=int, default=3,
                    help='number of epochs to train')
parser.add_argument('--batch-size', type=int, default=256,
                    help='number of examples per batch')
parser.add_argument('--print-every', type=int, default=100,
                    help='logging interval')
parser.add_argument('--factor-size', type=int, default=128,
                    help="the factor size of the embedding operation")
parser.add_argument('--num-gpus', type=int, default=2,
                    help="number of gpus to use")

MOVIELENS = {
    'dataset': 'ml-10m',
    'train': './ml-10M100K/r1.train',
    'val': './ml-10M100K/r1.test',
    'max_user': 71569,
    'max_movie': 65135,
}

if __name__ == '__main__':
    head = '%(asctime)-15s %(message)s'
    logging.basicConfig(level=logging.INFO, format=head)

    # arg parser
    args = parser.parse_args()
    logging.info(args)
    num_epoch = args.num_epoch
    batch_size = args.batch_size
    optimizer = 'sgd'
    factor_size = args.factor_size
    print_every = args.print_every
    num_gpus = args.num_gpus    
 
    momentum = 0.9
    learning_rate = 0.1

    # prepare dataset and iterators
    max_user = MOVIELENS['max_user']
    max_movies = MOVIELENS['max_movie']
    get_movielens_data(MOVIELENS['dataset'])
    train_iter = get_movielens_iter(MOVIELENS['train'], batch_size)
    val_iter = get_movielens_iter(MOVIELENS['val'], batch_size)

    # construct the model
    net = matrix_fact_model_parallel_net(factor_size, factor_size, max_user, max_movies)

    # construct the module
    # map the ctx_group attribute to the context assignment
    group2ctxs={'dev1':[mx.cpu()]*num_gpus, 'dev2':[mx.gpu(i) for i in range(num_gpus)]}

    # Creating a module by passing group2ctxs attribute which maps
    # the ctx_group attribute to the context assignment
    mod = mx.module.Module(symbol=net, context=[mx.cpu()]*num_gpus, data_names=['user', 'item'],
        label_names=['score'], group2ctxs=group2ctxs)
    
    # the initializer used to initialize the parameters
    initializer = mx.init.Xavier(factor_type="in", magnitude=2.34)
    
    # the parameters for the optimizer constructor
    optimizer_params = {
        'learning_rate': learning_rate,
        'wd': 1e-4,
        'momentum': momentum,
        'rescale_grad': 1.0/batch_size}

    # use MSE as the metric
    metric = mx.gluon.metric.create(['MSE'])
    
    speedometer = mx.callback.Speedometer(batch_size, print_every)
    
    # start training
    mod.fit(train_iter,
            val_iter,
            eval_metric        = metric,
            num_epoch          = num_epoch,
            optimizer          = optimizer,
            optimizer_params   = optimizer_params,
            initializer        = initializer,
            batch_end_callback = speedometer) 

[((27, 0, 27, 40), 'logging.basicConfig', 'logging.basicConfig', (), '', False, 'import logging\n'), ((29, 9, 30, 88), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((52, 4, 52, 56), 'logging.basicConfig', 'logging.basicConfig', (), '', False, 'import logging\n'), ((56, 4, 56, 22), 'logging.info', 'logging.info', ({(56, 17, 56, 21): 'args'}, {}), '(args)', False, 'import logging\n'), ((70, 4, 70, 44), 'get_data.get_movielens_data', 'get_movielens_data', ({(70, 23, 70, 43): "MOVIELENS['dataset']"}, {}), "(MOVIELENS['dataset'])", False, 'from get_data import get_movielens_iter, get_movielens_data\n'), ((71, 17, 71, 67), 'get_data.get_movielens_iter', 'get_movielens_iter', ({(71, 36, 71, 54): "MOVIELENS['train']", (71, 56, 71, 66): 'batch_size'}, {}), "(MOVIELENS['train'], batch_size)", False, 'from get_data import get_movielens_iter, get_movielens_data\n'), ((72, 15, 72, 63), 'get_data.get_movielens_iter', 'get_movielens_iter', ({(72, 34, 72, 50): "MOVIELENS['val']", (72, 52, 72, 62): 'batch_size'}, {}), "(MOVIELENS['val'], batch_size)", False, 'from get_data import get_movielens_iter, get_movielens_data\n'), ((75, 10, 75, 88), 'model.matrix_fact_model_parallel_net', 'matrix_fact_model_parallel_net', ({(75, 41, 75, 52): 'factor_size', (75, 54, 75, 65): 'factor_size', (75, 67, 75, 75): 'max_user', (75, 77, 75, 87): 'max_movies'}, {}), '(factor_size, factor_size, max_user, max_movies)', False, 'from model import matrix_fact_model_parallel_net\n'), ((87, 18, 87, 66), 'mxnet.init.Xavier', 'mx.init.Xavier', (), '', True, 'import mxnet as mx\n'), ((97, 13, 97, 44), 'mxnet.gluon.metric.create', 'mx.gluon.metric.create', ({(97, 36, 97, 43): "['MSE']"}, {}), "(['MSE'])", True, 'import mxnet as mx\n'), ((99, 18, 99, 66), 'mxnet.callback.Speedometer', 'mx.callback.Speedometer', ({(99, 42, 99, 52): 'batch_size', (99, 54, 99, 65): 'print_every'}, {}), '(batch_size, print_every)', True, 'import mxnet as mx\n'), ((79, 52, 79, 61), 'mxnet.gpu', 'mx.gpu', ({(79, 59, 79, 60): 'i'}, {}), '(i)', True, 'import mxnet as mx\n'), ((79, 24, 79, 32), 'mxnet.cpu', 'mx.cpu', ({}, {}), '()', True, 'import mxnet as mx\n'), ((83, 48, 83, 56), 'mxnet.cpu', 'mx.cpu', ({}, {}), '()', True, 'import mxnet as mx\n')]