// File: xbai0624[P]dnn/src/Matrix.cpp
// Function: Matrix::BatchNormalization(std::vector<Matrix, std::allocator<Matrix>>&)
void Matrix::BatchNormalization(std::vector<Matrix> &vM)
{
    // if only one matrix in vM, return. Otherwise the whitening procedure will set every element to 0.
    if(vM.size() <= 1) return; 
    auto dim = vM[0].Dimension();

    double epsilon = 1e-8;
    double batchSize = static_cast<double>(vM.size());

    auto get_mean_and_sigma = [&](std::vector<Matrix> &_vM, size_t i, size_t j, double &mean, double &sigma)
    {
        assert((int)i >= 0 && i < dim.first && (int)j >=0 && j < dim.second);

        mean = 0; sigma = 0;
        for(auto &m: _vM) {
            mean += m[i][j];
        }
        mean /= batchSize;

        for(auto &m: _vM){
            sigma += (m[i][j] - mean) * (m[i][j] - mean);
        }
        sigma /= batchSize;
        if(batchSize > 1.) sigma *= (batchSize/(batchSize - 1.)); // get the un-biased variate
        sigma = sqrt(sigma);
        sigma += epsilon; // avoid divide by zero error
    };

    // batch normalization
#ifndef PARALLEL_NORMALIZATION
    for(size_t i=0;i<dim.first;i++){
        for(size_t j=0;j<dim.second;j++)
        {
            double mean = 0, sigma = 0;
            get_mean_and_sigma(vM, i, j, mean, sigma);
            for(auto &m: vM)
                m[i][j] = (m[i][j] - mean)/sigma;
        }
    }
#else
    //int nElements = dim.first * dim.second;
    //int interval = nElements / PARALLEL_NORMALIZATION;
    //size_t Range[PARALLEL_NORMALIZATION+1];

    // to be implemented as needed in later, not much gain for now
#endif
}

// Assembly:
subq $0x88, %rsp
movq %rdi, 0x80(%rsp)
movq 0x80(%rsp), %rdi
callq 0x1f9c0
cmpq $0x1, %rax
ja 0x41a67
jmp 0x41c39
movq 0x80(%rsp), %rdi
xorl %eax, %eax
movl %eax, %esi
callq 0x1f9e0
movq %rax, %rdi
callq 0x3bae0
movq %rax, 0x70(%rsp)
movq %rdx, 0x78(%rsp)
movabsq $0x3e45798ee2308c3a, %rax # imm = 0x3E45798EE2308C3A
movq %rax, 0x68(%rsp)
movq 0x80(%rsp), %rdi
callq 0x1f9c0
movq %rax, %xmm0
movaps 0x105ae(%rip), %xmm1    # 0x52060
punpckldq %xmm1, %xmm0    # xmm0 = xmm0[0],xmm1[0],xmm0[1],xmm1[1]
movapd 0x105b2(%rip), %xmm1    # 0x52070
subpd %xmm1, %xmm0
movaps %xmm0, %xmm1
unpckhpd %xmm0, %xmm0            # xmm0 = xmm0[1,1]
addsd %xmm1, %xmm0
movsd %xmm0, 0x60(%rsp)
leaq 0x70(%rsp), %rax
movq %rax, 0x48(%rsp)
leaq 0x60(%rsp), %rax
movq %rax, 0x50(%rsp)
leaq 0x68(%rsp), %rax
movq %rax, 0x58(%rsp)
movq $0x0, 0x40(%rsp)
movq 0x40(%rsp), %rax
cmpq 0x70(%rsp), %rax
jae 0x41c39
movq $0x0, 0x38(%rsp)
movq 0x38(%rsp), %rax
cmpq 0x78(%rsp), %rax
jae 0x41c24
xorps %xmm0, %xmm0
movsd %xmm0, 0x30(%rsp)
xorps %xmm0, %xmm0
movsd %xmm0, 0x28(%rsp)
movq 0x80(%rsp), %rsi
movq 0x40(%rsp), %rdx
movq 0x38(%rsp), %rcx
leaq 0x48(%rsp), %rdi
leaq 0x30(%rsp), %r8
leaq 0x28(%rsp), %r9
callq 0x41c50
movq 0x80(%rsp), %rax
movq %rax, 0x20(%rsp)
movq 0x20(%rsp), %rdi
callq 0x1fe00
movq %rax, 0x18(%rsp)
movq 0x20(%rsp), %rdi
callq 0x1fe30
movq %rax, 0x10(%rsp)
leaq 0x18(%rsp), %rdi
leaq 0x10(%rsp), %rsi
callq 0x1fe60
testb $0x1, %al
jne 0x41b9b
jmp 0x41c0f
leaq 0x18(%rsp), %rdi
callq 0x1fea0
movq %rax, 0x8(%rsp)
movq 0x8(%rsp), %rdi
movq 0x40(%rsp), %rsi
callq 0x3be50
movq %rax, %rdi
movq 0x38(%rsp), %rsi
callq 0x1fd60
movsd (%rax), %xmm0
subsd 0x30(%rsp), %xmm0
divsd 0x28(%rsp), %xmm0
movsd %xmm0, (%rsp)
movq 0x8(%rsp), %rdi
movq 0x40(%rsp), %rsi
callq 0x3be50
movq %rax, %rdi
movq 0x38(%rsp), %rsi
callq 0x1fd60
movsd (%rsp), %xmm0
movsd %xmm0, (%rax)
leaq 0x18(%rsp), %rdi
callq 0x1ff20
jmp 0x41b86
jmp 0x41c11
movq 0x38(%rsp), %rax
addq $0x1, %rax
movq %rax, 0x38(%rsp)
jmp 0x41b13
jmp 0x41c26
movq 0x40(%rsp), %rax
addq $0x1, %rax
movq %rax, 0x40(%rsp)
jmp 0x41afa
addq $0x88, %rsp
retq
nopw %cs:(%rax,%rax)