// File: /workspace/github_cpp_makefile_star2_O2/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:
pushq %rbp
pushq %r15
pushq %r14
pushq %r13
pushq %r12
pushq %rbx
subq $0x18, %rsp
movq %rdi, %r14
movl %esi, 0x4(%rsp)
movslq %esi, %rax
imulq $0x18, %rax, %rax
addq 0x260(%rdi), %rax
movq %rax, 0x8(%rsp)
movl 0xc(%rdi), %eax
cmpl $0x4, %eax
je 0xbb09
cmpl $0x1, %eax
je 0xbb8e
testl %eax, %eax
jne 0xbc54
movq 0x8(%rsp), %rax
movq (%rax), %r15
movq 0x8(%rax), %rax
subq %r15, %rax
cmpq $0x18, %rax
jne 0xbc63
xorl %r12d, %r12d
cmpq 0x398(%r14), %r12
jae 0xbc54
imulq $0x18, %r12, %r13
xorl %ebx, %ebx
cmpq 0x3a0(%r14), %rbx
jae 0xbb89
movq 0x40(%r14), %rax
movq (%rax), %rax
movq (%rax,%r13), %rax
movq (%rax,%rbx,8), %rdi
callq 0x1c256
testb %al, %al
je 0xbb84
movq %r15, %rdi
movq %r12, %rsi
callq 0x17c46
movq (%rax), %rax
movsd (%rax,%rbx,8), %xmm0
movq 0x40(%r14), %rax
movq (%rax), %rax
movq (%rax,%r13), %rax
movq (%rax,%rbx,8), %rdi
movl 0x4(%rsp), %esi
callq 0x1d55e
incq %rbx
jmp 0xbb38
incq %r12
jmp 0xbb25
movq 0x8(%rsp), %rcx
movq 0x8(%rcx), %rax
subq (%rcx), %rax
pushq $0x18
popq %rcx
cqto
idivq %rcx
movq %rax, %r15
movq (%r14), %rax
movq %r14, %rdi
callq *0x288(%rax)
cmpq %rax, %r15
jne 0xbc82
xorl %eax, %eax
cmpq 0x3a8(%r14), %rax
jae 0xbc54
movq %rax, 0x10(%rsp)
imulq $0x18, %rax, %r13
movq 0x8(%rsp), %rax
movq (%rax), %r15
addq %r13, %r15
xorl %r12d, %r12d
cmpq 0x398(%r14), %r12
jae 0xbc47
imulq $0x18, %r12, %rbp
xorl %ebx, %ebx
cmpq 0x3a0(%r14), %rbx
jae 0xbc42
movq 0x40(%r14), %rax
movq (%rax,%r13), %rax
movq (%rax,%rbp), %rax
movq (%rax,%rbx,8), %rdi
callq 0x1c256
testb %al, %al
je 0xbc3d
movq %r15, %rdi
movq %r12, %rsi
callq 0x17c46
movq (%rax), %rax
movsd (%rax,%rbx,8), %xmm0
movq 0x40(%r14), %rax
movq (%rax,%r13), %rax
movq (%rax,%rbp), %rax
movq (%rax,%rbx,8), %rdi
movl 0x4(%rsp), %esi
callq 0x1d55e
incq %rbx
jmp 0xbbef
incq %r12
jmp 0xbbe0
movq 0x10(%rsp), %rax
incq %rax
jmp 0xbbbc
addq $0x18, %rsp
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
retq
leaq 0x16532(%rip), %rdi     # 0x2219c
leaq 0x15a69(%rip), %rsi     # 0x216da
leaq 0x16552(%rip), %rcx     # 0x221ca
movl $0x96e, %edx            # imm = 0x96E
callq 0x4180
leaq 0x16580(%rip), %rdi     # 0x22209
leaq 0x15a4a(%rip), %rsi     # 0x216da
leaq 0x16533(%rip), %rcx     # 0x221ca
movl $0x97c, %edx            # imm = 0x97C
callq 0x4180
nop