// File: /workspace/github_cpp_makefile_star2_O1/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
leaq (%rax,%rax,2), %rax
shlq $0x3, %rax
addq 0x260(%rdi), %rax
movq %rax, 0x8(%rsp)
movl 0xc(%rdi), %eax
cmpl $0x4, %eax
je 0xca8f
cmpl $0x1, %eax
je 0xcb29
testl %eax, %eax
jne 0xcc24
movq 0x8(%rsp), %rax
movq (%rax), %r15
movq 0x8(%rax), %rax
subq %r15, %rax
cmpq $0x18, %rax
jne 0xcc33
cmpq $0x0, 0x398(%r14)
je 0xcc24
xorl %r12d, %r12d
cmpq $0x0, 0x3a0(%r14)
je 0xcb18
leaq (%r12,%r12,2), %r13
xorl %ebp, %ebp
movq 0x40(%r14), %rax
movq (%rax), %rax
movq (%rax,%r13,8), %rax
movq (%rax,%rbp,8), %rdi
callq 0x2097e
testb %al, %al
je 0xcb0c
movq %r15, %rdi
movq %r12, %rsi
callq 0x1ae76
movq (%rax), %rax
movsd (%rax,%rbp,8), %xmm0
movq 0x40(%r14), %rax
movq (%rax), %rax
movq (%rax,%r13,8), %rax
movq (%rax,%rbp,8), %rdi
movl 0x4(%rsp), %esi
callq 0x21e4a
incq %rbp
cmpq 0x3a0(%r14), %rbp
jb 0xcac9
incq %r12
cmpq 0x398(%r14), %r12
jb 0xcab9
jmp 0xcc24
movq 0x8(%rsp), %rcx
movq 0x8(%rcx), %rax
subq (%rcx), %rax
sarq $0x3, %rax
movabsq $-0x5555555555555555, %rbx # imm = 0xAAAAAAAAAAAAAAAB
imulq %rax, %rbx
movq (%r14), %rax
movq %r14, %rdi
callq *0x288(%rax)
cmpq %rax, %rbx
jne 0xcc52
cmpq $0x0, 0x3a8(%r14)
je 0xcc24
movq $0x0, 0x10(%rsp)
cmpq $0x0, 0x398(%r14)
je 0xcc0a
movq 0x10(%rsp), %rax
leaq (%rax,%rax,2), %r13
leaq (,%r13,8), %r15
movq 0x8(%rsp), %rax
addq (%rax), %r15
xorl %r12d, %r12d
cmpq $0x0, 0x3a0(%r14)
je 0xcbfe
leaq (%r12,%r12,2), %rbp
xorl %ebx, %ebx
movq 0x40(%r14), %rax
movq (%rax,%r13,8), %rax
movq (%rax,%rbp,8), %rax
movq (%rax,%rbx,8), %rdi
callq 0x2097e
testb %al, %al
je 0xcbf2
movq %r15, %rdi
movq %r12, %rsi
callq 0x1ae76
movq (%rax), %rax
movsd (%rax,%rbx,8), %xmm0
movq 0x40(%r14), %rax
movq (%rax,%r13,8), %rax
movq (%rax,%rbp,8), %rax
movq (%rax,%rbx,8), %rdi
movl 0x4(%rsp), %esi
callq 0x21e4a
incq %rbx
cmpq 0x3a0(%r14), %rbx
jb 0xcbad
incq %r12
cmpq 0x398(%r14), %r12
jb 0xcb9d
movq 0x10(%rsp), %rcx
incq %rcx
movq %rcx, 0x10(%rsp)
cmpq 0x3a8(%r14), %rcx
jb 0xcb73
addq $0x18, %rsp
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
retq
leaq 0x1b611(%rip), %rdi     # 0x2824b
leaq 0x1ab09(%rip), %rsi     # 0x2774a
leaq 0x1b631(%rip), %rcx     # 0x28279
movl $0x96e, %edx            # imm = 0x96E
callq 0x4160
leaq 0x1b65f(%rip), %rdi     # 0x282b8
leaq 0x1aaea(%rip), %rsi     # 0x2774a
leaq 0x1b612(%rip), %rcx     # 0x28279
movl $0x97c, %edx            # imm = 0x97C
callq 0x4160
nop