// File: xbai0624[P]dnn/src/ConstructLayer.cpp
// Function: ConstructLayer::ReplaceNeuronZWithBatchNormalization(int)
void ConstructLayer::ReplaceNeuronZWithBatchNormalization(int sample_id)
{
    // replace the Z values (un-normalized) in neurons with normalized Z values 

    // logic:
    // since 1) z, a, \sigma^\prime calculation is all done in neuron 
    //          and layer just fetch the calculated value in neuron and form matrices (images)
    //       2) the batch normalization is done in layer, 
    //       3) so after batch normalization, one need to replace the old Z value in neuron (un-normalized)
    //          with the new Z value (normalized), then the neuron can use the normalized z value to
    //          calculate a and \sigma^\prime

    // it is simpler to use __imageYFull_BN instead of __imageY_BN
    // because:
    //     1) for 2D layer, the drop out happens on kernel, not on image, so the __imageZFull = __imageZ
    //        thus for 2D layer: __imageYFull_BN = __imageY_BN
    //     2) for 1D layer, the drop out happens on neuron, the image is not the same any more,
    //        however, the normalization process happens on each feature(channel), when we updating the Z 
    //        images, the inactive feature(channel) was filled with 0, so after normalization, the inactive 
    //        feature(channel) is still 0, __imageYFull_BN is a super-set of __imageY_BN with inactive channels
    //        filled with 0. So when replacing neuron Z value, if this neuron is inactive, we don't do anything;
    //        if the neuron is active, we replace it with the corresponding normalized z value. This way ensures
    //        that the dimension of neuron and __imageYFull_BN match each other during the replacing process, 
    //        thus saves time for matching dimensions

    auto & image_sample = __imageYFull_BN[sample_id];

    // currently only works for FC, output and cnn layers
    // other types of layers are not necessary

    if(__type == LayerType::fullyConnected || __type == LayerType::output)
    {
        assert(image_sample.OutputImageFromKernel.size() ==1);
        Matrix & y_matrix = image_sample.OutputImageFromKernel[0];
        for(size_t i=0;i<__neuronDim.i;i++){
            for(size_t j=0;j<__neuronDim.j;j++)
            {
                if(! __neurons[0][i][j] -> IsActive()) continue;
                double normalized_z = y_matrix[i][j];
                // fc layer only has one kernel
                __neurons[0][i][j] -> ReplaceZ(sample_id, normalized_z);
            }
        }
    }
    else if( __type == LayerType::cnn)
    {
        assert(image_sample.OutputImageFromKernel.size() == GetNumberOfKernelsCNN());

        for(size_t k=0;k<__neuronDim.k;k++)
        {
            Matrix & y_matrix = image_sample.OutputImageFromKernel[k];
            for(size_t i=0;i<__neuronDim.i;i++) {
                for(size_t j=0;j<__neuronDim.j;j++)
                {
                    if(!__neurons[k][i][j] -> IsActive()) continue;
                    double z_normalized = y_matrix[i][j];
                    __neurons[k][i][j] -> ReplaceZ(sample_id, z_normalized);
                }
            }
        }
    }
}

// Assembly:
subq $0x78, %rsp
movq %rdi, 0x70(%rsp)
movl %esi, 0x6c(%rsp)
movq 0x70(%rsp), %rdi
movq %rdi, 0x10(%rsp)
addq $0x260, %rdi            # imm = 0x260
movslq 0x6c(%rsp), %rsi
callq 0x1fd40
movq %rax, %rcx
movq 0x10(%rsp), %rax
movq %rcx, 0x60(%rsp)
cmpl $0x0, 0xc(%rax)
je 0x1316a
movq 0x10(%rsp), %rax
cmpl $0x4, 0xc(%rax)
jne 0x132b9
movq 0x60(%rsp), %rdi
callq 0x1f9c0
cmpq $0x1, %rax
jne 0x1317c
jmp 0x1319b
leaq 0x3fb50(%rip), %rdi     # 0x52cd3
leaq 0x3f048(%rip), %rsi     # 0x521d2
movl $0x96e, %edx            # imm = 0x96E
leaq 0x3fb6b(%rip), %rcx     # 0x52d01
callq 0x4170
movq 0x60(%rsp), %rdi
xorl %eax, %eax
movl %eax, %esi
callq 0x1f9e0
movq %rax, 0x58(%rsp)
movq $0x0, 0x50(%rsp)
movq 0x10(%rsp), %rcx
movq 0x50(%rsp), %rax
cmpq 0x398(%rcx), %rax
jae 0x132b4
movq $0x0, 0x48(%rsp)
movq 0x10(%rsp), %rcx
movq 0x48(%rsp), %rax
cmpq 0x3a0(%rcx), %rax
jae 0x1329f
movq 0x10(%rsp), %rdi
addq $0x40, %rdi
xorl %eax, %eax
movl %eax, %esi
callq 0x1f650
movq %rax, %rdi
movq 0x50(%rsp), %rsi
callq 0x1f5c0
movq %rax, %rdi
movq 0x48(%rsp), %rsi
callq 0x1f5f0
movq (%rax), %rdi
callq 0x47fe0
testb $0x1, %al
jne 0x13228
jmp 0x1328c
movq 0x58(%rsp), %rdi
movq 0x50(%rsp), %rsi
callq 0x3be50
movq %rax, %rdi
movq 0x48(%rsp), %rsi
callq 0x1fd60
movq 0x10(%rsp), %rdi
movsd (%rax), %xmm0
movsd %xmm0, 0x40(%rsp)
addq $0x40, %rdi
xorl %eax, %eax
movl %eax, %esi
callq 0x1f650
movq %rax, %rdi
movq 0x50(%rsp), %rsi
callq 0x1f5c0
movq %rax, %rdi
movq 0x48(%rsp), %rsi
callq 0x1f5f0
movq (%rax), %rdi
movl 0x6c(%rsp), %esi
movsd 0x40(%rsp), %xmm0
callq 0x4a160
movq 0x48(%rsp), %rax
addq $0x1, %rax
movq %rax, 0x48(%rsp)
jmp 0x131d7
jmp 0x132a1
movq 0x50(%rsp), %rax
addq $0x1, %rax
movq %rax, 0x50(%rsp)
jmp 0x131b7
jmp 0x13468
movq 0x10(%rsp), %rax
cmpl $0x1, 0xc(%rax)
jne 0x13466
movq 0x60(%rsp), %rdi
callq 0x1f9c0
movq 0x10(%rsp), %rdi
movq %rax, 0x8(%rsp)
movq (%rdi), %rax
callq *0x288(%rax)
movq %rax, %rcx
movq 0x8(%rsp), %rax
cmpq %rcx, %rax
jne 0x132f4
jmp 0x13313
leaq 0x3fa45(%rip), %rdi     # 0x52d40
leaq 0x3eed0(%rip), %rsi     # 0x521d2
movl $0x97c, %edx            # imm = 0x97C
leaq 0x3f9f3(%rip), %rcx     # 0x52d01
callq 0x4170
movq $0x0, 0x38(%rsp)
movq 0x10(%rsp), %rcx
movq 0x38(%rsp), %rax
cmpq 0x3a8(%rcx), %rax
jae 0x13464
movq 0x60(%rsp), %rdi
movq 0x38(%rsp), %rsi
callq 0x1f9e0
movq %rax, 0x30(%rsp)
movq $0x0, 0x28(%rsp)
movq 0x10(%rsp), %rcx
movq 0x28(%rsp), %rax
cmpq 0x398(%rcx), %rax
jae 0x1344f
movq $0x0, 0x20(%rsp)
movq 0x10(%rsp), %rcx
movq 0x20(%rsp), %rax
cmpq 0x3a0(%rcx), %rax
jae 0x1343a
movq 0x10(%rsp), %rdi
addq $0x40, %rdi
movq 0x38(%rsp), %rsi
callq 0x1f650
movq %rax, %rdi
movq 0x28(%rsp), %rsi
callq 0x1f5c0
movq %rax, %rdi
movq 0x20(%rsp), %rsi
callq 0x1f5f0
movq (%rax), %rdi
callq 0x47fe0
testb $0x1, %al
jne 0x133c2
jmp 0x13427
movq 0x30(%rsp), %rdi
movq 0x28(%rsp), %rsi
callq 0x3be50
movq %rax, %rdi
movq 0x20(%rsp), %rsi
callq 0x1fd60
movq 0x10(%rsp), %rdi
movsd (%rax), %xmm0
movsd %xmm0, 0x18(%rsp)
addq $0x40, %rdi
movq 0x38(%rsp), %rsi
callq 0x1f650
movq %rax, %rdi
movq 0x28(%rsp), %rsi
callq 0x1f5c0
movq %rax, %rdi
movq 0x20(%rsp), %rsi
callq 0x1f5f0
movq (%rax), %rdi
movl 0x6c(%rsp), %esi
movsd 0x18(%rsp), %xmm0
callq 0x4a160
movq 0x20(%rsp), %rax
addq $0x1, %rax
movq %rax, 0x20(%rsp)
jmp 0x13370
jmp 0x1343c
movq 0x28(%rsp), %rax
addq $0x1, %rax
movq %rax, 0x28(%rsp)
jmp 0x13350
jmp 0x13451
movq 0x38(%rsp), %rax
addq $0x1, %rax
movq %rax, 0x38(%rsp)
jmp 0x1331c
jmp 0x13466
jmp 0x13468
addq $0x78, %rsp
retq
nopl (%rax)