import mxnet as mx
from ..config import config
from ..symbol.symbol_common import get_sym_train


def conv_only(from_layer, name, num_filter, kernel=(1, 1), pad=(0, 0),
              stride=(1, 1), bias_wd_mult=0.0):
    weight = mx.symbol.Variable(name="{}_weight".format(name),
                                init=mx.init.Normal(0.01), attr={'__lr_mult__': '1.0'})
    bias = mx.symbol.Variable(name="{}_bias".format(name),
                              init=mx.init.Constant(0.0), attr={'__lr_mult__': '2.0', '__wd_mult__': str(bias_wd_mult)})
    conv = mx.symbol.Convolution(data=from_layer, kernel=kernel, pad=pad,
                                 stride=stride, num_filter=num_filter, name="{}".format(name), weight=weight, bias=bias)
    return conv


def conv_act_layer_dw(from_layer, name, num_filter, kernel=(1, 1), pad=(0, 0),
                      stride=(1, 1), act_type="relu", bias_wd_mult=0.0):
    assert kernel[0] == 3
    weight = mx.symbol.Variable(name="{}_weight".format(name),
                                init=mx.init.Normal(0.01), attr={'__lr_mult__': '1.0'})
    bias = mx.symbol.Variable(name="{}_bias".format(name),
                              init=mx.init.Constant(0.0), attr={'__lr_mult__': '2.0', '__wd_mult__': str(bias_wd_mult)})
    conv = mx.symbol.Convolution(data=from_layer, kernel=kernel, pad=pad,
                                 stride=stride, num_filter=num_filter, num_group=num_filter, name="{}".format(name),
                                 weight=weight, bias=bias)
    conv = mx.sym.BatchNorm(data=conv, fix_gamma=False, eps=2e-5, momentum=0.9, name=name + '_bn')
    if len(act_type) > 0:
        relu = mx.symbol.Activation(data=conv, act_type=act_type,
                                    name="{}_{}".format(name, act_type))
    else:
        relu = conv
    return relu


def conv_act_layer(from_layer, name, num_filter, kernel=(1, 1), pad=(0, 0),
                   stride=(1, 1), act_type="relu", bias_wd_mult=0.0, separable=False, filter_in=-1):
    separable = False
    if separable:
        assert kernel[0] == 3
    if not separable:
        weight = mx.symbol.Variable(name="{}_weight".format(name),
                                    init=mx.init.Normal(0.01), attr={'__lr_mult__': '1.0'})
        bias = mx.symbol.Variable(name="{}_bias".format(name),
                                  init=mx.init.Constant(0.0), attr={'__lr_mult__': '2.0', '__wd_mult__': str(bias_wd_mult)})
        conv = mx.symbol.Convolution(data=from_layer, kernel=kernel, pad=pad,
                                     stride=stride, num_filter=num_filter, name="{}".format(name), weight=weight, bias=bias)
        conv = mx.sym.BatchNorm(data=conv, fix_gamma=False, eps=2e-5, momentum=0.9, name=name + '_bn')
    else:
        if filter_in < 0:
            filter_in = num_filter
        conv = mx.symbol.Convolution(data=from_layer, kernel=kernel, pad=pad,
                                     stride=stride, num_filter=filter_in, num_group=filter_in, name="{}_sep".format(name))
        conv = mx.sym.BatchNorm(data=conv, fix_gamma=False, eps=2e-5, momentum=0.9, name=name + '_sep_bn')
        conv = mx.symbol.Activation(data=conv, act_type='relu',
                                    name="{}_sep_bn_relu".format(name))
        conv = mx.symbol.Convolution(data=conv, kernel=(1, 1), pad=(0, 0),
                                     stride=(1, 1), num_filter=num_filter, name="{}".format(name))
        conv = mx.sym.BatchNorm(data=conv, fix_gamma=False, eps=2e-5, momentum=0.9, name=name + '_bn')
    if len(act_type) > 0:
        relu = mx.symbol.Activation(data=conv, act_type=act_type,
                                    name="{}_{}".format(name, act_type))
    else:
        relu = conv
    return relu


def ssh_context_module(body, num_filter, filter_in, name):
    conv_dimred = conv_act_layer(body, name + '_conv1',
                                 num_filter, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', separable=True,
                                 filter_in=filter_in)
    conv5x5 = conv_act_layer(conv_dimred, name + '_conv2',
                             num_filter, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='', separable=True)
    conv7x7_1 = conv_act_layer(conv_dimred, name + '_conv3_1',
                               num_filter, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', separable=True)
    conv7x7 = conv_act_layer(conv7x7_1, name + '_conv3_2',
                             num_filter, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='', separable=True)
    return (conv5x5, conv7x7)


def conv_deformable(net, num_filter, num_group=1, act_type='relu', name=''):
    f = num_group * 18
    conv_offset = mx.symbol.Convolution(name=name + '_conv_offset', data=net,
                                        num_filter=f, pad=(1, 1), kernel=(3, 3), stride=(1, 1))
    net = mx.contrib.symbol.DeformableConvolution(name=name + "_conv", data=net, offset=conv_offset,
                                                  num_filter=num_filter, pad=(1, 1), kernel=(3, 3),
                                                  num_deformable_group=num_group, stride=(1, 1), no_bias=False)
    net = mx.sym.BatchNorm(data=net, fix_gamma=False, eps=2e-5, momentum=0.9, name=name + '_bn')
    if len(act_type) > 0:
        net = mx.symbol.Activation(data=net, act_type=act_type, name=name + '_act')
    return net


def ssh_detection_module(body, num_filter, filter_in, name):
    conv3x3 = conv_act_layer(body, name + '_conv1',
                             num_filter, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='', separable=True,
                             filter_in=filter_in)
    conv5x5, conv7x7 = ssh_context_module(body, num_filter // 2, filter_in, name + '_context')
    ret = mx.sym.concat(*[conv3x3, conv5x5, conv7x7], dim=1, name=name + '_concat')
    ret = mx.symbol.Activation(data=ret, act_type='relu', name=name + '_concat_relu')
    if config.USE_DCN == 1:
        ret = conv_deformable(ret, num_filter=num_filter * 2, name=name + '_concat_dcn')
    # elif config.USE_DCN == 2:
    #     ret = conv_deformable2(ret, num_filter=num_filter * 2, name=name + '_concat_dcn')
    return ret


def get_resnet_conv(data, sym):
    all_layers = sym.get_internals()
    isize = 640
    _, out_shape, _ = all_layers.infer_shape(data=(1, 3, isize, isize))
    last_entry = None
    c1 = None
    c2 = None
    c3 = None
    outputs = all_layers.list_outputs()
    count = len(outputs)
    for i in range(count):
        name = outputs[i]
        shape = out_shape[i]
        if not name.endswith('_output'):
            continue
        if len(shape) != 4:
            continue

        if c1 is None and shape[2] == isize // 16:
            cname = last_entry[0]

            c1 = all_layers[cname]
        if c2 is None and shape[2] == isize // 32:
            cname = last_entry[0]

            c2 = all_layers[cname]
        if shape[2] == isize // 32:
            c3 = all_layers[name]

        last_entry = (name, shape)

    c1_filter = -1
    c2_filter = -1
    c3_filter = -1

    F1 = 256
    F2 = 256
    _bwm = 1.0
    if config.NET_MODE == 0:
        c1_lateral = conv_act_layer(c1, 'ssh_m1_red_conv',
                                    F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c2_lateral = conv_act_layer(c2, 'ssh_m2_red_conv',
                                    F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c2_up = mx.symbol.UpSampling(c2_lateral, scale=2, sample_type='nearest', workspace=512, name='ssh_m2_red_up',
                                     num_args=1)

        c2_up = mx.symbol.Crop(*[c2_up, c1_lateral])

        c1 = c1_lateral + c2_up

        c1 = conv_act_layer(c1, 'ssh_m1_conv',
                            F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        m1 = ssh_detection_module(c1, F2, F2, 'ssh_m1_det')
        m2 = ssh_detection_module(c2, F1, c2_filter, 'ssh_m2_det')
        m3 = ssh_detection_module(c3, F1, c3_filter, 'ssh_m3_det')
    elif config.NET_MODE == 1:
        c3_lateral = conv_act_layer(c3, 'ssh_c3_lateral',
                                    F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c3_up = mx.symbol.UpSampling(c3_lateral, scale=2, sample_type='nearest', workspace=512, name='ssh_c3_up', num_args=1)
        c2_lateral = conv_act_layer(c2, 'ssh_c2_lateral',
                                    F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c3_up = mx.symbol.Crop(*[c3_up, c2_lateral])
        c2 = c2_lateral + c3_up
        c2 = conv_act_layer(c2, 'ssh_c2_aggr',
                            F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c1_lateral = conv_act_layer(c1, 'ssh_m1_red_conv',
                                    F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c2_up = mx.symbol.UpSampling(c2, scale=2, sample_type='nearest', workspace=512, name='ssh_m2_red_up', num_args=1)

        c2_up = mx.symbol.Crop(*[c2_up, c1_lateral])

        c1 = c1_lateral + c2_up

        c1 = conv_act_layer(c1, 'ssh_m1_conv',
                            F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        m1 = ssh_detection_module(c1, F2, F2, 'ssh_m1_det')
        m2 = ssh_detection_module(c2, F1, c2_filter, 'ssh_m2_det')
        m3 = ssh_detection_module(c3, F1, c3_filter, 'ssh_m3_det')
    elif config.NET_MODE == 2:
        n1 = ssh_detection_module(c1, F2, F2, 'ssh_n1_det')
        n2 = ssh_detection_module(c2, F1, c2_filter, 'ssh_n2_det')
        n3 = ssh_detection_module(c3, F1, c3_filter, 'ssh_n3_det')
        c3 = conv_act_layer(c3, 'ssh_c3_lateral',
                            F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c3_up = mx.symbol.UpSampling(c3, scale=2, sample_type='nearest', workspace=512, name='ssh_c3_up', num_args=1)
        c2_lateral = conv_act_layer(c2, 'ssh_c2_lateral',
                                    F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c3_up = mx.symbol.Crop(*[c3_up, c2_lateral])
        c2 = c2_lateral + c3_up
        c2 = conv_act_layer(c2, 'ssh_c2_aggr',
                            F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c1_lateral = conv_act_layer(c1, 'ssh_m1_red_conv',
                                    F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c2_up = mx.symbol.UpSampling(c2, scale=2, sample_type='nearest', workspace=512, name='ssh_m2_red_up', num_args=1)

        c2_up = mx.symbol.Crop(*[c2_up, c1_lateral])
        c1 = c1_lateral + c2_up
        c1 = conv_act_layer(c1, 'ssh_c1_aggr',
                            F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)

        m1 = ssh_detection_module(c1, F2, F2, 'ssh_m1_det')
        m2 = ssh_detection_module(c2, F1, c2_filter, 'ssh_m2_det')
        m3 = ssh_detection_module(c3, F1, c3_filter, 'ssh_m3_det')
    elif config.NET_MODE == 3:
        c3 = ssh_detection_module(c3, F2 // 2, c3_filter, 'ssh_c3_lateral')
        c3_up = mx.symbol.UpSampling(c3, scale=2, sample_type='nearest', workspace=512, name='ssh_c3_up', num_args=1)
        c2_lateral = ssh_detection_module(c2, F2 // 2, c2_filter, 'ssh_c2_lateral')
        c3_up = mx.symbol.Crop(*[c3_up, c2_lateral])
        c2 = c2_lateral + c3_up
        c2 = conv_act_layer(c2, 'ssh_c2_aggr',
                            F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)

        c1_lateral = ssh_detection_module(c1, F2 // 2, c1_filter, 'ssh_c1_lateral')
        c2_up = mx.symbol.UpSampling(c2, scale=2, sample_type='nearest', workspace=512, name='ssh_m2_red_up', num_args=1)

        c2_up = mx.symbol.Crop(*[c2_up, c1_lateral])
        c1 = c1_lateral + c2_up
        c1 = conv_act_layer(c1, 'ssh_c1_aggr',
                            F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)

        m1 = ssh_detection_module(c1, F2, F2, 'ssh_m1_det')
        m2 = ssh_detection_module(c2, F1, c2_filter, 'ssh_m2_det')
        m3 = ssh_detection_module(c3, F1, c3_filter, 'ssh_m3_det')
    elif config.NET_MODE == 4:
        c3 = conv_act_layer(c3, 'ssh_c3_lateral',
                            F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c3_up = mx.symbol.UpSampling(c3, scale=2, sample_type='nearest', workspace=512, name='ssh_c3_up', num_args=1)
        c2_lateral = conv_act_layer(c2, 'ssh_c2_lateral',
                                    F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c3_up = mx.symbol.Crop(*[c3_up, c2_lateral])
        c2 = c2_lateral + c3_up
        c2 = conv_act_layer(c2, 'ssh_c2_aggr',
                            F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c1_lateral = conv_act_layer(c1, 'ssh_m1_red_conv',
                                    F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c2_up = mx.symbol.UpSampling(c2, scale=2, sample_type='nearest', workspace=512, name='ssh_m2_red_up', num_args=1)

        c2_up = mx.symbol.Crop(*[c2_up, c1_lateral])
        c1 = c1_lateral + c2_up
        c1 = conv_act_layer(c1, 'ssh_c1_aggr',
                            F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)

        m1 = ssh_detection_module(c1, F2 // 2, F2, 'ssh_m1_det')
        m2 = ssh_detection_module(c2, F1 // 2, c2_filter, 'ssh_m2_det')
        m3 = ssh_detection_module(c3, F1 // 2, c3_filter, 'ssh_m3_det')
    elif config.NET_MODE == 5:
        c3 = conv_act_layer_dw(c3, 'ssh_c3_lateral_m',
                               F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c3 = conv_act_layer(c3, 'ssh_c3_lateral',
                            F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c3_up = mx.symbol.UpSampling(c3, scale=2, sample_type='nearest', workspace=512, name='ssh_c3_up', num_args=1)
        c2 = conv_act_layer_dw(c2, 'ssh_c2_lateral_m',
                               F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c2_lateral = conv_act_layer(c2, 'ssh_c2_lateral',
                                    F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c3_up = mx.symbol.Crop(*[c3_up, c2_lateral])
        c2 = c2_lateral + c3_up
        c2 = conv_act_layer(c2, 'ssh_c2_aggr',
                            F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c1 = conv_act_layer_dw(c1, 'ssh_c1_lateral_m',
                               F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c1_lateral = conv_act_layer(c1, 'ssh_m1_red_conv',
                                    F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
        c2_up = mx.symbol.UpSampling(c2, scale=2, sample_type='nearest', workspace=512, name='ssh_m2_red_up', num_args=1)

        c2_up = mx.symbol.Crop(*[c2_up, c1_lateral])
        c1 = c1_lateral + c2_up
        c1 = conv_act_layer(c1, 'ssh_c1_aggr',
                            F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)

        m1 = ssh_detection_module(c1, F2, F2, 'ssh_m1_det')
        m2 = ssh_detection_module(c2, F1, c2_filter, 'ssh_m2_det')
        m3 = ssh_detection_module(c3, F1, c3_filter, 'ssh_m3_det')

    return {8: m1, 16: m2, 32: m3}, {8: n1, 16: n2, 32: n3}


def get_out(conv_fpn_feat, prefix, stride, landmark=False, lr_mult=1.0):
    A = config.NUM_ANCHORS
    bbox_pred_len = 4
    landmark_pred_len = 10
    if config.USE_BLUR:
        bbox_pred_len = 5
    if config.USE_OCCLUSION:
        landmark_pred_len = 15
    ret_group = []
    num_anchors = config.RPN_ANCHOR_CFG[str(stride)]['NUM_ANCHORS']
    label = mx.symbol.Variable(name='%s_label_stride%d' % (prefix, stride))
    bbox_target = mx.symbol.Variable(name='%s_bbox_target_stride%d' % (prefix, stride))
    bbox_weight = mx.symbol.Variable(name='%s_bbox_weight_stride%d' % (prefix, stride))
    if landmark:
        landmark_target = mx.symbol.Variable(name='%s_landmark_target_stride%d' % (prefix, stride))
        landmark_weight = mx.symbol.Variable(name='%s_landmark_weight_stride%d' % (prefix, stride))
    rpn_relu = conv_fpn_feat[stride]
    maxout_stat = 0
    if config.USE_MAXOUT >= 1 and stride == config.RPN_FEAT_STRIDE[-1]:
        maxout_stat = 1
    if config.USE_MAXOUT >= 2 and stride != config.RPN_FEAT_STRIDE[-1]:
        maxout_stat = 2

    if maxout_stat == 0:
        rpn_cls_score = conv_only(rpn_relu, '%s_rpn_cls_score_stride%d' % (prefix, stride), 2 * num_anchors,
                                  kernel=(1, 1), pad=(0, 0), stride=(1, 1))
    elif maxout_stat == 1:
        cls_list = []
        for a in range(num_anchors):
            rpn_cls_score_bg = conv_only(rpn_relu, '%s_rpn_cls_score_stride%d_anchor%d_bg' % (prefix, stride, a), 3,
                                         kernel=(1, 1), pad=(0, 0), stride=(1, 1))
            rpn_cls_score_bg = mx.sym.max(rpn_cls_score_bg, axis=1, keepdims=True)
            cls_list.append(rpn_cls_score_bg)
            rpn_cls_score_fg = conv_only(rpn_relu, '%s_rpn_cls_score_stride%d_anchor%d_fg' % (prefix, stride, a), 1,
                                         kernel=(1, 1), pad=(0, 0), stride=(1, 1))
            cls_list.append(rpn_cls_score_fg)
        rpn_cls_score = mx.sym.concat(*cls_list, dim=1, name='%s_rpn_cls_score_stride%d' % (prefix, stride))
    else:
        cls_list = []
        for a in range(num_anchors):
            rpn_cls_score_bg = conv_only(rpn_relu, '%s_rpn_cls_score_stride%d_anchor%d_bg' % (prefix, stride, a), 1,
                                         kernel=(1, 1), pad=(0, 0), stride=(1, 1))
            cls_list.append(rpn_cls_score_bg)
            rpn_cls_score_fg = conv_only(rpn_relu, '%s_rpn_cls_score_stride%d_anchor%d_fg' % (prefix, stride, a), 3,
                                         kernel=(1, 1), pad=(0, 0), stride=(1, 1))
            rpn_cls_score_fg = mx.sym.max(rpn_cls_score_fg, axis=1, keepdims=True)
            cls_list.append(rpn_cls_score_fg)
        rpn_cls_score = mx.sym.concat(*cls_list, dim=1, name='%s_rpn_cls_score_stride%d' % (prefix, stride))

    rpn_bbox_pred = conv_only(rpn_relu, '%s_rpn_bbox_pred_stride%d' % (prefix, stride), bbox_pred_len * num_anchors,
                              kernel=(1, 1), pad=(0, 0), stride=(1, 1))

    # prepare rpn data
    if not config.FBN:
        rpn_cls_score_reshape = mx.symbol.Reshape(data=rpn_cls_score,
                                                  shape=(0, 2, -1),
                                                  name="%s_rpn_cls_score_reshape_stride%s" % (prefix, stride))
    else:
        rpn_cls_score_reshape = mx.symbol.Reshape(data=rpn_cls_score,
                                                  shape=(0, 2, -1),
                                                  name="%s_rpn_cls_score_reshape_stride%s_pre" % (prefix, stride))
        rpn_cls_score_reshape = mx.symbol.BatchNorm(rpn_cls_score_reshape, fix_gamma=True, eps=2e-5,
                                                    name="%s_rpn_cls_score_reshape_stride%s" % (prefix, stride))

    rpn_bbox_pred_reshape = mx.symbol.Reshape(data=rpn_bbox_pred,
                                              shape=(0, 0, -1),
                                              name="%s_rpn_bbox_pred_reshape_stride%s" % (prefix, stride))
    if landmark:
        rpn_landmark_pred = conv_only(rpn_relu, '%s_rpn_landmark_pred_stride%d' % (prefix, stride),
                                      landmark_pred_len * num_anchors,
                                      kernel=(1, 1), pad=(0, 0), stride=(1, 1))
        rpn_landmark_pred_reshape = mx.symbol.Reshape(data=rpn_landmark_pred,
                                                      shape=(0, 0, -1),
                                                      name="%s_rpn_landmark_pred_reshape_stride%s" % (prefix, stride))

    if config.TRAIN.RPN_ENABLE_OHEM >= 2:
        label, anchor_weight = mx.sym.Custom(op_type='rpn_fpn_ohem3', stride=int(stride), network=config.network,
                                             dataset=config.dataset, prefix=prefix, cls_score=rpn_cls_score_reshape,
                                             labels=label)

        _bbox_weight = mx.sym.tile(anchor_weight, (1, 1, bbox_pred_len))
        _bbox_weight = _bbox_weight.reshape((0, -1, A * bbox_pred_len)).transpose((0, 2, 1))
        bbox_weight = mx.sym.elemwise_mul(bbox_weight, _bbox_weight, name='%s_bbox_weight_mul_stride%s' % (prefix, stride))

        if landmark:
            _landmark_weight = mx.sym.tile(anchor_weight, (1, 1, landmark_pred_len))
            _landmark_weight = _landmark_weight.reshape((0, -1, A * landmark_pred_len)).transpose((0, 2, 1))
            landmark_weight = mx.sym.elemwise_mul(landmark_weight, _landmark_weight,
                                                  name='%s_landmark_weight_mul_stride%s' % (prefix, stride))
    # cls loss
    rpn_cls_prob = mx.symbol.SoftmaxOutput(data=rpn_cls_score_reshape,
                                           label=label,
                                           multi_output=True,
                                           normalization='valid', use_ignore=True, ignore_label=-1,
                                           grad_scale=lr_mult,
                                           name='%s_rpn_cls_prob_stride%d' % (prefix, stride))
    ret_group.append(rpn_cls_prob)
    ret_group.append(mx.sym.BlockGrad(label))

    # bbox loss
    bbox_diff = rpn_bbox_pred_reshape - bbox_target
    bbox_diff = bbox_diff * bbox_weight
    rpn_bbox_loss_ = mx.symbol.smooth_l1(name='%s_rpn_bbox_loss_stride%d_' % (prefix, stride), scalar=3.0, data=bbox_diff)
    rpn_bbox_loss = mx.sym.MakeLoss(name='%s_rpn_bbox_loss_stride%d' % (prefix, stride), data=rpn_bbox_loss_,
                                    grad_scale=1.0 * lr_mult / (config.TRAIN.RPN_BATCH_SIZE))
    ret_group.append(rpn_bbox_loss)
    ret_group.append(mx.sym.BlockGrad(bbox_weight))

    # landmark loss
    if landmark:
        landmark_diff = rpn_landmark_pred_reshape - landmark_target
        landmark_diff = landmark_diff * landmark_weight
        rpn_landmark_loss_ = mx.symbol.smooth_l1(name='%s_rpn_landmark_loss_stride%d_' % (prefix, stride), scalar=3.0,
                                                 data=landmark_diff)
        rpn_landmark_loss = mx.sym.MakeLoss(name='%s_rpn_landmark_loss_stride%d' % (prefix, stride), data=rpn_landmark_loss_,
                                            grad_scale=0.5 * lr_mult / (config.TRAIN.RPN_BATCH_SIZE))
        ret_group.append(rpn_landmark_loss)
        ret_group.append(mx.sym.BlockGrad(landmark_weight))
    return ret_group


def get_resnet_train(sym):
    return get_sym_train(sym)
