def test_fullyconnected_with_type():
    sym = mx.sym.FullyConnected(num_hidden=3, name='inner')
    ctx_list = [{'ctx': mx.gpu(0), 'inner_data': (2, 10), 'type_dict': {'inner_data': np.float64}},
                {'ctx': mx.gpu(0), 'inner_data': (2, 10), 'type_dict': {'inner_data': np.float32}},
                {'ctx': mx.cpu(0), 'inner_data': (2, 10), 'type_dict': {'inner_data': np.float64}},
                {'ctx': mx.cpu(0), 'inner_data': (2, 10), 'type_dict': {'inner_data': np.float32}}]
    check_consistency(sym, ctx_list)

def test_convolution_with_type():
    sym = mx.sym.Convolution(num_filter=3, kernel=(3,3), name='conv')
    ctx_list = [{'ctx': mx.gpu(0), 'conv_data': (2, 2, 10, 10), 'type_dict': {'conv_data': np.float64}},
                {'ctx': mx.gpu(0), 'conv_data': (2, 2, 10, 10), 'type_dict': {'conv_data': np.float32}},
                {'ctx': mx.cpu(0), 'conv_data': (2, 2, 10, 10), 'type_dict': {'conv_data': np.float64}},
                {'ctx': mx.cpu(0), 'conv_data': (2, 2, 10, 10), 'type_dict': {'conv_data': np.float32}}]
    check_consistency(sym, ctx_list)

def test_activation_with_type():
    sym = mx.sym.Activation(name='act', act_type='sigmoid')
    ctx_list = [{'ctx': mx.gpu(0), 'act_data': (2, 2, 10, 10), 'type_dict': {'act_data': np.float64}},
                {'ctx': mx.gpu(0), 'act_data': (2, 2, 10, 10), 'type_dict': {'act_data': np.float32}},
                {'ctx': mx.cpu(0), 'act_data': (2, 2, 10, 10), 'type_dict': {'act_data': np.float64}},
                {'ctx': mx.cpu(0), 'act_data': (2, 2, 10, 10), 'type_dict': {'act_data': np.float32}}]
    check_consistency(sym, ctx_list)

def make_image(img, m, color_ref=None):
    generator = symbol.generator_symbol(m, 'style')
    args = mx.nd.load('args%s_style.nd'%img)
    for i in range(m):
        args['znoise_%d'%i] = mx.nd.zeros([1,1,16*2**i,16*2**i], mx.gpu())
        args['znoise_%d'%i][:] = np.random.uniform(-250,250,[1,1,16*2**i,16*2**i])
    auxs = mx.nd.load('auxs%s_style.nd'%img)
    with open('models/model%s.pkl'%img, 'w') as f:
        pickle.dump([args, auxs, generator], f)

    for i in range(m):
        args['znoise_%d'%i] = mx.nd.zeros([1,1,16*2**i,16*2**i], mx.gpu())
        args['zim_%d'%i] = mx.nd.zeros([1,3,16*2**i, 16*2**i], mx.gpu())
    gene_executor = generator.bind(ctx=mx.gpu(), args=args, aux_states=mx.nd.load('auxs%s_style.nd'%img))
    for test_im in os.listdir('test_pics'): 
        print test_im
        if color_ref:
            color_ref = cv2.cvtColor(crop_img('test_pics/%s'%test_im, 8*2**m), cv2.COLOR_RGB2HSV)
        for i in range(m):
            gene_executor.arg_dict['zim_%d'%i][:] = preprocess_img('test_pics/%s'%test_im, 16*2**i)
        for ii in range(4):
            t = time.clock()
            for i in range(m):
                gene_executor.arg_dict['znoise_%d'%i][:] = np.random.uniform(-150*ii,150*ii,[1,1,16*2**i,16*2**i])
            gene_executor.forward(is_train=True)
            out = gene_executor.outputs[0].asnumpy()
            im = postprocess_img(out, color_ref)
            cv2.imwrite('models/%s_%s_%d.jpg'%(test_im.split('.')[0], img, ii), im)

def init(args):
    layers = [784] + [args.hidden_size] * args.num_hidden + [10]
    biases = [mx.nd.zeros((1, x), ctx=mx.gpu(0)) for x in layers[1:]]
    weights = [
        mx.nd.zeros((x, y), ctx=mx.gpu(0)) for x, y in zip(layers[:-1], layers[1:])
    ]
    return weights, biases

def init_yolo():
#     initializer = mx.init.Xavier(factor_type="in", magnitude=2.34)
    initializer = mx.init.Uniform(5e-4)
    pretrained_model = mx.model.FeedForward.load(PRETRAINED[0], PRETRAINED[1], ctx=mx.gpu())
#     args = mx.nd.load('args2.nd')
#     auxs = mx.nd.load('auxs2.nd')
    arg_params = pretrained_model.arg_params
    aux_params = pretrained_model.aux_params
    symbol = get_yolo_symbol()

    arg_shapes, output_shapes, aux_shapes = symbol.infer_shape(data=(BATCHSIZE, 3, 448, 448))
    arg_names = symbol.list_arguments()
    aux_names = symbol.list_auxiliary_states()
    arg_dict = dict(zip(arg_names, [mx.nd.zeros(shape, ctx=mx.gpu()) for shape in arg_shapes]))
    aux_dict = dict(zip(aux_names, [mx.nd.zeros(shape, ctx=mx.gpu()) for shape in aux_shapes]))
    grad_dict = dict(zip(arg_names, [mx.nd.zeros(shape, ctx=mx.gpu()) for shape in arg_shapes]))
    for name in arg_dict:
        if name.endswith('label'):
            continue
        if name.startswith('data'):
            continue
        if name in arg_params and not name.startswith('fullyconnected'):
            arg_params[name].copyto(arg_dict[name])
        else:
            print name
            initializer(name, arg_dict[name])
    for name in aux_dict:
        if 0 and name in aux_params and not name.startswith('fullyconnected'):
            aux_params[name].copyto(aux_dict[name])
        else:
            initializer(name, aux_dict[name])
    executor = symbol.bind(ctx=mx.gpu(), args=arg_dict, args_grad=grad_dict, aux_states=aux_dict, grad_req='write')
#     executor = symbol.bind(ctx=mx.gpu(), args=args, args_grad=grad_dict, aux_states=auxs, grad_req='write')
    return executor

def handleFolder(GUPid,tasks):
    #synset = [l.strip() for l in open(args.synset).readlines()]
    prefix = "full-resnet-152"
    num_round = 0	
    model = mx.model.FeedForward.load( prefix, num_round, ctx=mx.gpu(GUPid),numpy_batch_size=1)
    internals = model.symbol.get_internals()
    fea_symbol = internals["pool1_output"]	  
    feature_extractor = mx.model.FeedForward( ctx=mx.gpu(GUPid), symbol=fea_symbol, numpy_batch_size=1, \
            arg_params=model.arg_params, aux_params=model.aux_params, allow_extra_params=True)
			
    #subfolders = [ fold for fold in os.listdir(folder)]			
    k = 0	
    for subfolder in tasks:
      workspace_folder = os.path.join(folder,subfolder)
      print "extract label####",subfolder,"---GPU: ",GUPid," process: ",k,"/",len(tasks)
      i = 0
      k +=1	  
      feature_array = []
      for filename in os.listdir(workspace_folder):
        if '.jpg'	in filename or '.JPEG' in filename:
          i +=1		
          m = cv2.imread(os.path.join(workspace_folder,filename),1)	  
          img = cv2.cvtColor(m, cv2.COLOR_BGR2RGB)
          img = cv2.resize(img, (224, 224))  # resize to 224*224 to fit model
          img = np.swapaxes(img, 0, 2)
          img = np.swapaxes(img, 1, 2)  # change to (c, h,w) order
          img = img[np.newaxis, :]  # extend to (n, c, h, w)
          f = feature_extractor.predict(img)
          f = np.ravel(f)
          #print f.shape		  
          feature_array.append((f[0],subfolder,filename))
      random.shuffle(feature_array)
      #print len(feature_array)	  
      np.save((os.path.join(workspace_folder,"test.npy")),feature_array[:int(i*test_ratio)])
      np.save((os.path.join(workspace_folder,"train.npy")),feature_array[int(i*(test_ratio)):])

def main(args):
    num_rnn_layers = 1
    net = rnn_unroll(num_rnn_layers,
                     args.num_unroll_steps,
                     args.input_size,
                     args.hidden_size,
                     args.num_classes)

    arg_shapes, out_shapes, aux_shapes = net.infer_shape(data=(args.batch_size, args.num_unroll_steps, args.input_size))
    arg_types, out_types, aux_types = net.infer_type(data=mx.base.mx_real_t)

    arg_arrays = [mx.nd.zeros(shape, mx.gpu(0), dtype=dtype)
                  for shape, dtype in zip(arg_shapes, arg_types)]
    grad_dict = {name : mx.nd.zeros(shape, mx.gpu(0), dtype=dtype)
                 for name, shape, dtype in zip(net.list_arguments(), arg_shapes, arg_types)
                 if name != 'data'}

    executor = net.bind(ctx=mx.gpu(0),
                        args=arg_arrays,
                        args_grad=grad_dict,
                        grad_req='write')

    for i in range(args.num_loops):
        if i == num_cold:
            start = time.time()
        outputs = executor.forward()
        if args.only_forward:
            for o in outputs:
                o.wait_to_read()
            continue
        executor.backward([outputs[0]])
        for name, grad in grad_dict.items():
            grad.wait_to_read()
    dur = time.time() - start
    print('Per Loop Time: %.6f' % (dur / (args.num_loops - num_cold)))

def test_concat_with_type():
    sym = mx.sym.Concat(name="concat", num_args=2)
    ctx_list = [
        {
            "ctx": mx.gpu(0),
            "concat_arg1": (2, 10),
            "concat_arg0": (2, 10),
            "type_dict": {"concat_arg0": np.float64, "concat_arg1": np.float64},
        },
        {
            "ctx": mx.gpu(0),
            "concat_arg1": (2, 10),
            "concat_arg0": (2, 10),
            "type_dict": {"concat_arg0": np.float32, "concat_arg1": np.float32},
        },
        {
            "ctx": mx.gpu(0),
            "concat_arg1": (2, 10),
            "concat_arg0": (2, 10),
            "type_dict": {"concat_arg0": np.float16, "concat_arg1": np.float16},
        },
        {
            "ctx": mx.cpu(0),
            "concat_arg1": (2, 10),
            "concat_arg0": (2, 10),
            "type_dict": {"concat_arg0": np.float64, "concat_arg1": np.float64},
        },
        {
            "ctx": mx.cpu(0),
            "concat_arg1": (2, 10),
            "concat_arg0": (2, 10),
            "type_dict": {"concat_arg0": np.float32, "concat_arg1": np.float32},
        },
    ]
    check_consistency(sym, ctx_list)

def test_elementwisesum_with_type():
    sym = mx.sym.ElementWiseSum(name="ews", num_args=2)
    ctx_list = [
        {
            "ctx": mx.gpu(0),
            "ews_arg1": (2, 10),
            "ews_arg0": (2, 10),
            "type_dict": {"ews_arg0": np.float64, "ews_arg1": np.float64},
        },
        {
            "ctx": mx.gpu(0),
            "ews_arg1": (2, 10),
            "ews_arg0": (2, 10),
            "type_dict": {"ews_arg0": np.float32, "ews_arg1": np.float32},
        },
        {
            "ctx": mx.gpu(0),
            "ews_arg1": (2, 10),
            "ews_arg0": (2, 10),
            "type_dict": {"ews_arg0": np.float16, "ews_arg1": np.float16},
        },
        {
            "ctx": mx.cpu(0),
            "ews_arg1": (2, 10),
            "ews_arg0": (2, 10),
            "type_dict": {"ews_arg0": np.float64, "ews_arg1": np.float64},
        },
        {
            "ctx": mx.cpu(0),
            "ews_arg1": (2, 10),
            "ews_arg0": (2, 10),
            "type_dict": {"ews_arg0": np.float32, "ews_arg1": np.float32},
        },
    ]
    check_consistency(sym, ctx_list)

    def ffbp(self, X, y):
        h = mx.nd.zeros(self.hshape, ctx=mx.gpu(0))  # init hidden state
        rnn_cache = []
        for t in xrange(self.num_unroll_steps):
            h, rnn_cache_t = rnn_step_forward(X, h, self.params['Wx'],
                                self.params['Wh'], self.params['b'])
            rnn_cache.append(rnn_cache_t)
        predict, affine_cache = affine_forward(h, self.params['Wa'], self.params['ba'])
        loss, grad = l2_loss(predict, y)

        daffine, dWa, dba = affine_backward(grad, affine_cache)

        dx = mx.nd.zeros((X.shape[0], X.shape[1]), ctx=mx.gpu(0))
        dWx = mx.nd.zeros((X.shape[1], daffine.shape[1]), ctx=mx.gpu(0))
        dWh = mx.nd.zeros((daffine.shape[1], daffine.shape[1]), ctx=mx.gpu(0))
        db = mx.nd.zeros((daffine.shape[1],), ctx=mx.gpu(0))

        dnext_h_t = daffine
        for t in xrange(self.num_unroll_steps):
            dx_t, dprev_h_t, dWx_t, dWh_t, db_t = rnn_step_backward(dnext_h_t, rnn_cache[t])
            dnext_h_t = dprev_h_t
        
            dx += dx_t
            dWx += dWx_t
            dWh += dWh_t
            db += db_t
        
        dx.wait_to_read()
        dWx.wait_to_read()
        dWh.wait_to_read()
        db.wait_to_read()

def test_convolution_with_type():
    np.random.seed(1234)
    sym1 = mx.sym.Convolution(num_filter=3, kernel=(3,3), name='conv')

    data = mx.sym.Variable('conv_data')
    w = mx.sym.Variable('conv_weight')
    b = mx.sym.Variable('conv_bias')
    w = mx.sym.transpose(w, axes=(0,2,3,1))
    sym2 = mx.sym.transpose(data, axes=(0,2,3,1))
    sym2 = mx.sym.Convolution(sym2, w, b, layout='NHWC', num_filter=3, kernel=(3,3))
    sym2 = mx.sym.transpose(sym2, axes=(0,3,1,2), name='conv')

    sym = [sym1, sym1, sym1, sym1, sym1, sym2, sym2]
    ctx_list = [{'ctx': mx.gpu(0), 'conv_data': (2, 2, 10, 10), 'type_dict': {'conv_data': np.float64}},
                {'ctx': mx.gpu(0), 'conv_data': (2, 2, 10, 10), 'type_dict': {'conv_data': np.float32}},
                {'ctx': mx.gpu(0), 'conv_data': (2, 2, 10, 10), 'type_dict': {'conv_data': np.float16}},
                {'ctx': mx.cpu(0), 'conv_data': (2, 2, 10, 10), 'type_dict': {'conv_data': np.float64}},
                {'ctx': mx.cpu(0), 'conv_data': (2, 2, 10, 10), 'type_dict': {'conv_data': np.float32}},
                # NHWC
                {'ctx': mx.gpu(0), 'conv_data': (2, 2, 10, 10), 'conv_weight': (3, 2, 3, 3),
                 'type_dict': {'conv_data': np.float32, 'conv_weight': np.float32}},
                {'ctx': mx.gpu(0), 'conv_data': (2, 2, 10, 10), 'conv_weight': (3, 2, 3, 3),
                 'type_dict': {'conv_data': np.float16, 'conv_weight': np.float16}}
                ]
    # wider tolerance needed for true-fp16 NCHW test above
    tol = {np.dtype(np.float16): 0.5,
               np.dtype(np.float32): 1e-3,
               np.dtype(np.float64): 1e-5,
               np.dtype(np.uint8): 0,
               np.dtype(np.int32): 0}
    check_consistency(sym, ctx_list, tol=tol)
    # test ability to turn off training on bias
    check_consistency(sym, ctx_list, grad_req={'conv_data': 'write', 'conv_weight': 'write', 'conv_bias': 'null'}, tol=tol)

def main(args):
    # Create data iterators for training and testing sets.
    net = mx.symbol.Variable('data')
    net = mx.symbol.FullyConnected(data=net, num_hidden=args.hidden_size)
    net = mx.symbol.Activation(data=net, act_type="relu")
    net = mx.symbol.FullyConnected(data=net, num_hidden=10)
    net = mx.symbol.SoftmaxOutput(data=net, name='softmax')

    arg_shapes, out_shapes, aux_shapes = net.infer_shape(data=(args.batch_size, 784))
    arg_types, out_types, aux_types = net.infer_type(data=mx.base.mx_real_t)

    arg_arrays = [mx.nd.zeros(shape, mx.gpu(0), dtype=dtype)
                  for shape, dtype in zip(arg_shapes, arg_types)]
    grad_dict = {name : mx.nd.zeros(shape, mx.gpu(0), dtype=dtype)
                 for name, shape, dtype in zip(net.list_arguments(), arg_shapes, arg_types)
                 if name != 'data'}

    executor = net.bind(ctx=mx.gpu(0),
                        args=arg_arrays,
                        args_grad=grad_dict,
                        grad_req='write')

    start = time.time()
    for i in range(num_loops):
        outputs = executor.forward()
        if args.only_forward:
            for o in outputs:
                o.wait_to_read()
            continue
        executor.backward([outputs[0]])
        for name, grad in grad_dict.items():
            grad.wait_to_read()
    dur = time.time() - start
    print('Per Loop Time: %.6f' % (dur / num_loops))

def test_row_sparse_pull():
    kv = init_kv_with_str('row_sparse')
    kv.init('e', mx.nd.ones(shape).tostype('row_sparse'))

    def check_row_sparse_pull(kv, count, ctx=default_context()):
        num_rows = shape[0]
        vals = []
        row_ids = []
        all_row_ids = np.arange(num_rows)
        for i in range(count):
            vals.append(mx.nd.zeros(shape, ctx=ctx).tostype('row_sparse'))
            row_id = np.random.randint(num_rows, size=num_rows)
            row_ids.append(mx.nd.array(row_id, dtype='int64'))
        row_ids_to_pull = row_ids[0] if len(row_ids) == 1 else row_ids
        vals_to_pull = vals[0] if len(vals) == 1 else vals

        kv.row_sparse_pull('e', out=vals_to_pull, row_ids=row_ids_to_pull)
        for val, row_id in zip(vals, row_ids):
            retained = val.asnumpy()
            excluded_row_ids = np.setdiff1d(all_row_ids, row_id.asnumpy())
            for row in range(num_rows):
                expected_val = np.zeros_like(retained[row])
                expected_val += 0 if row in excluded_row_ids else 1
                assert_almost_equal(retained[row], expected_val)

    check_row_sparse_pull(kv, 1, mx.gpu(0))
    check_row_sparse_pull(kv, 4, mx.gpu(0))

def test_upsampling_with_type():
    sym = mx.sym.UpSampling(scale=2, num_filter=2, name='up', sample_type = 'nearest', num_args=1)
    ctx_list = [{'ctx': mx.gpu(0), 'up_arg0': (2, 2, 2, 10), 'type_dict': {'up_arg0': np.float64}},
                {'ctx': mx.gpu(0), 'up_arg0': (2, 2, 2, 10), 'type_dict': {'up_arg0': np.float32}},
                {'ctx': mx.gpu(0), 'up_arg0': (2, 2, 2, 10), 'type_dict': {'up_arg0': np.float16}},
                {'ctx': mx.cpu(0), 'up_arg0': (2, 2, 2, 10), 'type_dict': {'up_arg0': np.float64}},
                {'ctx': mx.cpu(0), 'up_arg0': (2, 2, 2, 10), 'type_dict': {'up_arg0': np.float32}}]
    check_consistency(sym, ctx_list)

def get_context(args):
    if args.gpu is None or args.gpu == '':
        context = [mx.cpu()]
    elif isinstance(args.gpu, int):
        context = [mx.gpu(args.gpu)]
    else:
        context = [mx.gpu(int(i)) for i in args.gpu]
    return context

def test_swapaxis_with_type():
    sym = mx.sym.SwapAxis(name='swap', dim1=1)
    ctx_list = [{'ctx': mx.gpu(0), 'swap_data': (2, 2, 2, 10), 'type_dict': {'swap_data': np.float64}},
                {'ctx': mx.gpu(0), 'swap_data': (2, 2, 2, 10), 'type_dict': {'swap_data': np.float32}},
                {'ctx': mx.gpu(0), 'swap_data': (2, 2, 2, 10), 'type_dict': {'swap_data': np.float16}},
                {'ctx': mx.cpu(0), 'swap_data': (2, 2, 2, 10), 'type_dict': {'swap_data': np.float64}},
                {'ctx': mx.cpu(0), 'swap_data': (2, 2, 2, 10), 'type_dict': {'swap_data': np.float32}}]
    check_consistency(sym, ctx_list)

 def test_wrapper(*args, **kwargs):
     try:
         a = mx.nd.zeros((1,), ctx=mx.gpu(gpu_id))
         ctx = mx.gpu(gpu_id)
     except Exception:
         ctx = mx.cpu()
     with ctx:
         orig_test(*args, **kwargs)

def test_blockgrad_with_type():
    sym = mx.sym.BlockGrad(name='bg')
    ctx_list = [{'ctx': mx.gpu(0), 'bg_data': (2, 2, 2, 10), 'type_dict': {'bg_data': np.float64}},
                {'ctx': mx.gpu(0), 'bg_data': (2, 2, 2, 10), 'type_dict': {'bg_data': np.float32}},
                {'ctx': mx.gpu(0), 'bg_data': (2, 2, 2, 10), 'type_dict': {'bg_data': np.float16}},
                {'ctx': mx.cpu(0), 'bg_data': (2, 2, 2, 10), 'type_dict': {'bg_data': np.float64}},
                {'ctx': mx.cpu(0), 'bg_data': (2, 2, 2, 10), 'type_dict': {'bg_data': np.float32}}]
    check_consistency(sym, ctx_list)

def test_reshape_with_type():
    sym = mx.sym.Reshape(name='reshape', shape=(-1,1,1,0))
    ctx_list = [{'ctx': mx.gpu(0), 'reshape_data': (2, 2, 2, 10), 'type_dict': {'reshape_data': np.float64}},
                {'ctx': mx.gpu(0), 'reshape_data': (2, 2, 2, 10), 'type_dict': {'reshape_data': np.float32}},
                {'ctx': mx.gpu(0), 'reshape_data': (2, 2, 2, 10), 'type_dict': {'reshape_data': np.float16}},
                {'ctx': mx.cpu(0), 'reshape_data': (2, 2, 2, 10), 'type_dict': {'reshape_data': np.float64}},
                {'ctx': mx.cpu(0), 'reshape_data': (2, 2, 2, 10), 'type_dict': {'reshape_data': np.float32}}]
    check_consistency(sym, ctx_list)