from __future__ import print_function

caffe_root = '/home/toanhoi/caffe/'

import sys

sys.path.insert(0, caffe_root + 'python')

import caffe

import math

from caffe import layers as L

from caffe.proto import caffe_pb2

def bn_relu_conv_bn_relu(bottom, nout, dropout,split):

    if split == 'train':

        use_global_stats = False

    else:

        use_global_stats=True

    batch_norm1 = L.BatchNorm(bottom, batch_norm_param=dict(use_global_stats=use_global_stats), in_place=False,

                              param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),

                                     dict(lr_mult=0, decay_mult=0)])

    scale1 = L.Scale(batch_norm1, bias_term=True, in_place=True,filler=dict(value=1), bias_filler=dict(value=0))

    relu1 = L.ReLU(scale1, in_place=True)

    conv1 = L.Convolution(relu1, kernel_size=[1, 1, 1], pad=[0, 0, 0], stride=[1,1,1],

                          param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,

                          num_output=nout * 4, axis=1, weight_filler=dict(type='msra'),

                          bias_filler=dict(type='constant'))

    batch_norm2 = L.BatchNorm(conv1, batch_norm_param=dict(use_global_stats=use_global_stats), in_place=False,

                              param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),

                                     dict(lr_mult=0, decay_mult=0)])

    scale2 = L.Scale(batch_norm2, bias_term=True, in_place=True,filler=dict(value=1), bias_filler=dict(value=0))

    relu2 = L.ReLU(scale2, in_place=True)

    conv2 = L.Convolution(relu2, param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,

                          axis=1, num_output=nout, pad=[1, 1, 1], kernel_size=[3, 3, 3], stride=[1,1,1],

                          weight_filler=dict(type='msra'), bias_filler=dict(type='constant'))

    if dropout > 0:

        conv2 = L.Dropout(conv2, dropout_ratio=dropout)

    return conv2

def add_layer(bottom, num_filter, dropout,split):

    conv = bn_relu_conv_bn_relu(bottom, nout=num_filter, dropout=dropout,split=split)

    concate = L.Concat(bottom, conv, axis=1)

    return concate

def transition(bottom, num_filter, split):

    if split == 'train':

        use_global_stats = False

    else:

        use_global_stats=True

    batch_norm1 = L.BatchNorm(bottom, batch_norm_param=dict(use_global_stats=use_global_stats), in_place=False,

                              param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),

                                     dict(lr_mult=0, decay_mult=0)])

    scale1 = L.Scale(batch_norm1, bias_term=True, in_place=True,filler=dict(value=1), bias_filler=dict(value=0))

    relu1 = L.ReLU(scale1, in_place=True)

    conv1 = L.Convolution(relu1, param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,

                          axis=1, num_output=num_filter, pad=[0, 0, 0], kernel_size=[1, 1, 1],stride=[1,1,1],

                          weight_filler=dict(type='msra'), bias_filler=dict(type='constant'))

    batch_norm2 = L.BatchNorm(conv1, batch_norm_param=dict(use_global_stats=use_global_stats), in_place=False,

                               param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),

                                  dict(lr_mult=0, decay_mult=0)])

    scale2 = L.Scale(batch_norm2, bias_term=True, in_place=True, filler=dict(value=1), bias_filler=dict(value=0))

    relu2 = L.ReLU(scale2, in_place=True)

    conv_down = L.Convolution(relu2, param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,

                          axis=1, num_output=num_filter, pad=[0, 0, 0], kernel_size=[2, 2, 2], stride=2,

                          weight_filler=dict(type='msra'), bias_filler=dict(type='constant'))

    #pooling = L.Pooling(conv1, type="Pooling", pool=P.Pooling.MAX, kernel_size=2, stride=2, engine=1)

    return conv_down

# first_output -- #channels before entering the first dense block, set it to be comparable to growth_rate

# growth_rate -- growth rate

# dropout -- set to 0 to disable dropout, non-zero number to set dropout rate

def densenet(split, batch_size=4, first_output=32, growth_rate=16, dropout=0.2):

    source_train_path = './train_list.txt'

    source_test_path = './test_list.txt'

    patch_size = [64, 64, 64]

    n = caffe.NetSpec()

    num_classes = 4

    reduction= 0.5

    N=[4,4,4,4]

    if split == 'train':

        n.data, n.label = L.HDF5Data(name="data", batch_size=batch_size, source=source_train_path, ntop=2, shuffle=True,

                                     transform_param=dict(crop_size_l=patch_size[0], crop_size_h=patch_size[1],

                                                          crop_size_w=patch_size[2]), include={'phase': caffe.TRAIN})

    elif split == 'val':

        n.data, n.label = L.HDF5Data(name="data", batch_size=batch_size, source=source_test_path, ntop=2, shuffle=True,

                                     transform_param=dict(crop_size_l=patch_size[0], crop_size_h=patch_size[1],

                                                          crop_size_w=patch_size[2]),

                                     include={'phase': caffe.TEST})

    else:

        n.data = L.Input(name="data", ntop=1, input_param={'shape': {'dim': [1, 2, patch_size[0], patch_size[1], patch_size[2]]}})

    nchannels = first_output

    # Fist layers

    n.conv1a = L.Convolution(n.data, param=[dict(lr_mult=1, decay_mult=1), dict(lr_mult=2, decay_mult=0)],

                             axis=1, num_output=nchannels, pad=[1,1,1], kernel_size=[3, 3, 3], stride=[1,1,1],

                             weight_filler=dict(type='msra'), bias_filler=dict(type='constant',value=-0.1))

    if split == 'train':

        use_global_stats = False

    else:

        use_global_stats=True

    n.bnorm1a = L.BatchNorm(n.conv1a, batch_norm_param=dict(use_global_stats=use_global_stats), param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),

                                     dict(lr_mult=0, decay_mult=0)],  in_place=False)

    n.scale1a = L.Scale(n.bnorm1a, in_place=True, bias_term=True,filler=dict(value=1), bias_filler=dict(value=0))

    n.relu1a = L.ReLU(n.bnorm1a, in_place=True)

    # conv 1b, after BN set bias_term=false

    n.conv1b = L.Convolution(n.relu1a, param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,

                             axis=1, num_output=nchannels, pad=[1, 1, 1], kernel_size=[3, 3, 3], stride=[1,1,1],

                             weight_filler=dict(type='msra'), bias_filler=dict(type='constant'))

    n.bnorm1b = L.BatchNorm(n.conv1b, batch_norm_param=dict(use_global_stats=use_global_stats), param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),

                                     dict(lr_mult=0, decay_mult=0)], in_place=False)

    n.scale1b = L.Scale(n.bnorm1b, in_place=True, bias_term=True, filler=dict(value=1), bias_filler=dict(value=0))

    n.relu1b = L.ReLU(n.bnorm1b, in_place=True)

    n.conv1c = L.Convolution(n.relu1b, param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,

                             axis=1, num_output=nchannels, pad=[1, 1, 1], kernel_size=[3, 3, 3],stride=[1,1,1],

                             weight_filler=dict(type='msra'), bias_filler=dict(type='constant'))

    print (nchannels)

    # model = L.Pooling(n.conv1c, type="Pooling", pool=P.Pooling.MAX, kernel_size=2, stride=2, engine=1)

    n.bnorm1c = L.BatchNorm(n.conv1c, batch_norm_param=dict(use_global_stats=use_global_stats), param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),

                                     dict(lr_mult=0, decay_mult=0)], in_place=False)

    n.scale1c = L.Scale(n.bnorm1c, in_place=True, bias_term=True, filler=dict(value=1), bias_filler=dict(value=0))

    n.relu1c = L.ReLU(n.bnorm1c, in_place=True)

    model = L.Convolution(n.relu1c, param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,

                             axis=1, num_output=nchannels, pad=[0, 0, 0], kernel_size=[2, 2, 2], stride=2,

                             weight_filler=dict(type='msra'), bias_filler=dict(type='constant'))

    n.__setattr__("Conv_down_1", model)

    # ===============Dense block 2=====================

    for i in range(N[0]):

        if (i == 0):

            concat = add_layer(model, growth_rate, dropout,split)

            n.__setattr__("Concat_%d" % (i + 1), concat)

            nchannels += growth_rate

            continue

        concat = add_layer(concat, growth_rate, dropout,split)

        n.__setattr__("Concat_%d" % (i + 1), concat)

        nchannels += growth_rate

    # ===============End dense block 2=================

    print (nchannels)

    # ===============Deconvolution layer 2==============

    model_deconv_x2 = L.Deconvolution(concat, param=[dict(lr_mult=0.1, decay_mult=1)],

                                      convolution_param=dict(kernel_size=[4,4,4], stride=[2,2,2], num_output=num_classes,

                                                             pad=[1, 1, 1], group=num_classes,

                                                             weight_filler=dict(type='bilinear_3D'),

                                                             bias_term=False))

    n.__setattr__("Deconvolution_%d" % (N[0] + 1), model_deconv_x2)

    # ===============End Deconvolution layer 2==============

    # ===============Transition layer 2=================

    model = transition(concat, int(math.floor(nchannels * reduction)), split)

    n.__setattr__("Conv_down_%d" % (N[0] + 1), model)

    nchannels = int(math.floor(nchannels * reduction))

    # ===============End Transition layer2==============

    # ===============Dense block 3=====================

    for i in range(N[1]):

        if (i == 0):

            concat = add_layer(model, growth_rate, dropout, split)

            n.__setattr__("Concat_%d" % (N[1] + i + 2), concat)

            nchannels += growth_rate

            continue

        concat = add_layer(concat, growth_rate, dropout, split)

        n.__setattr__("Concat_%d" % (N[1] + i + 2), concat)

        nchannels += growth_rate

    # ===============End dense block 3=================

    print (nchannels)

    # ===============Deconvolution layer 3==============

    model_deconv_x4 = L.Deconvolution(concat, param=[dict(lr_mult=0.1, decay_mult=1)],

                                      convolution_param=dict(kernel_size=[6,6,6], stride=[4,4,4], num_output=num_classes,

                                                             pad=[1, 1, 1], group=num_classes,

                                                             weight_filler=dict(type='bilinear_3D'),

                                                             bias_term=False))

    n.__setattr__("Deconvolution_%d" % (N[0] + N[1] + 2), model_deconv_x4)

    # ==============Transition layer 3=================

    model = transition(concat, int(math.floor(nchannels * reduction)), split)

    n.__setattr__("Conv_down_%d" % (N[0] + N[1] + 2), model)

    # ===============End Transition layer3==============

    nchannels = int(math.floor(nchannels * reduction))

    # ===============Dense block 4=====================

    for i in range(N[2]):

        if (i == 0):

            concat = add_layer(model, growth_rate, dropout, split)

            n.__setattr__("Concat_%d" % (N[0] + N[1] + i + 3), concat)

            nchannels += growth_rate

            continue

        concat = add_layer(concat, growth_rate, dropout, split)

        n.__setattr__("Concat_%d" % (N[0] + N[1] + i + 3), concat)

        nchannels += growth_rate

    # ===============End dense block 4=================

    # ===============Transition layer 4=================

    print(nchannels)

    # ===============Deconvolution layer 4==============

    model_deconv_x8 = L.Deconvolution(concat, param=[dict(lr_mult=0.1, decay_mult=1)],

                                      convolution_param=dict(kernel_size=[10,10,10], stride=[8,8,8], num_output=num_classes,

                                                             pad=[1, 1, 1], group=num_classes,

                                                             weight_filler=dict(type='bilinear_3D'),

                                                             bias_term=False))

    n.__setattr__("Deconvolution_%d" % (N[0] + N[1] + N[2] + 3), model_deconv_x8)

    # ===============End Deconvolution layer 4==============

    # ===============Transition layer 4=================

    model = transition(concat, int(math.floor(nchannels * reduction)), split)

    n.__setattr__("Conv_down_%d" % (N[0] + N[1] + N[2] + 3), model)

    nchannels = int(math.floor(nchannels * reduction))

    # ===============End Transition layer3==============

    # ===============Dense block 5=====================

    for i in range(N[3]):

        if (i == 0):

            concat = add_layer(model, growth_rate, dropout, split)

            n.__setattr__("Concat_%d" % (N[0] + N[1] + N[2] + N[3] + i + 3), concat)

            nchannels += growth_rate

            continue

        concat = add_layer(concat, growth_rate, dropout, split)

        n.__setattr__("Concat_%d" % (N[0] + N[1] + N[2] + N[3] + i + 3), concat)

        nchannels += growth_rate

    # ===============End dense block 5=================

    print(nchannels)

    # ===============Deconvolution layer 5==============

    model_deconv_x16 = L.Deconvolution(concat, param=[dict(lr_mult=0.1, decay_mult=1)],

                                       convolution_param=dict(kernel_size=[18, 18, 18], stride=[16,16,16],

                                                              num_output=num_classes,

                                                              pad=[1, 1, 1], group=num_classes,

                                                              weight_filler=dict(type='bilinear_3D'),

                                                              bias_term=False))

    n.__setattr__("Deconvolution_%d" % (N[0] + N[1] + N[2] + N[3] + 4), model_deconv_x16)

    model = L.Concat(n.conv1c,model_deconv_x2, model_deconv_x4, model_deconv_x8, model_deconv_x16,

                     axis=1)

    n.bnorm_concat= L.BatchNorm(model, batch_norm_param=dict(use_global_stats=use_global_stats), param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),

                                     dict(lr_mult=0, decay_mult=0)], in_place=False)

    n.scale_concat = L.Scale(n.bnorm_concat, in_place=True, bias_term=True, filler=dict(value=1), bias_filler=dict(value=0))

    n.relu_concat = L.ReLU(n.scale_concat, in_place=True)

    model_conv_concate = L.Convolution(n.relu_concat,

                                       param=[dict(lr_mult=1, decay_mult=1), dict(lr_mult=2, decay_mult=0)],

                                       axis=1, num_output=num_classes, pad=[0, 0, 0], kernel_size=[1, 1, 1],

                                       weight_filler=dict(type='msra'))

    if (split == 'train'):

        n.loss = L.SoftmaxWithLoss(model_conv_concate, n.label)

    elif (split == 'val'):

        n.loss = L.SoftmaxWithLoss(model_conv_concate, n.label)

    else:

        n.softmax = L.Softmax(model_conv_concate, ntop=1, in_place=False)

    return n.to_proto()

def make_net():

    with open('train_3d_denseseg.prototxt', 'w') as f:

        print(str(densenet('train', batch_size=4)), file=f)

    with open('test_3d_denseseg.prototxt', 'w') as f:

        print(str(densenet('val', batch_size=4)), file=f)

    with open('deploy_3d_denseseg.prototxt', 'w') as f:

        print(str(densenet('deploy', batch_size=0)), file=f)

def make_solver():

    s = caffe_pb2.SolverParameter()

    s.random_seed = 0xCAFFE

    s.train_net = 'train_3d_denseseg.prototxt'

    s.max_iter = 200000

    s.type = 'Adam'

    s.display = 20

    s.base_lr = 0.0002

    #s.power=0.9

    s.momentum = 0.97

    s.weight_decay = 0.0005

    s.average_loss=20

    s.iter_size = 1

    s.lr_policy='step'

    s.stepsize=50000

    s.gamma = 0.1

    s.snapshot_prefix ='./snapshot/3d_denseseg_iseg'

    s.snapshot = 2000

    s.solver_mode = caffe_pb2.SolverParameter.GPU

    solver_path = 'solver.prototxt'

    with open(solver_path, 'w') as f:

        f.write(str(s))

if __name__ == '__main__':

    make_net()

    make_solver()