// File: xbai0624[P]dnn/src/ConstructLayer.cpp
// Function: ConstructLayer::BatchNormalization_UpdateZ_2DLayer()
void ConstructLayer::BatchNormalization_UpdateZ_2DLayer()
{
    // 1) get mu and sigma_square
    // part I) active
    Images mu, sigma_square;
    GetMu_BFor2DLayer(__imageZ, mu);
    GetSigmaSquareFor2DLayer(__imageZ, mu, sigma_square);
    // part II) full
    Images mu_full, sigma_square_full;
    GetMu_BFor2DLayer(__imageZFull, mu_full);
    GetSigmaSquareFor2DLayer(__imageZFull, mu_full, sigma_square_full);

    // 2) save mu and sigma square
    // part I) active
    __imageAverage_BN[0] = mu;
    __imageVariance_BN[0] = sigma_square;
    // part II) full
    __imageAverageFull_BN[0] = mu_full;
    __imageVarianceFull_BN[0] = sigma_square_full;

    // 3) get normalized \hat{x} and save it to __imageZ_BN
    size_t batch_size = __imageZ.size();
    double epsilon = 1e-8; // avoid divide by zero issue; a hard-coded parameter from algorithm
    size_t nKernel = __imageZ[0].OutputImageFromKernel.size();
    // for parallelization
    std::vector<std::thread> vth;
    int Range[NTHREAD+1];
    for(int i=0;i<NTHREAD;i++)
        Range[i] = batch_size/NTHREAD * i;
    Range[NTHREAD] = batch_size;

    /*
    for(size_t k=0;k<nKernel;k++)
    {
        // part I) active mu and sigma matrix
        Matrix & average_m_active = mu.OutputImageFromKernel[k];
        Matrix variance_m_active = sigma_square.OutputImageFromKernel[k];
        variance_m_active = variance_m_active + epsilon;
        variance_m_active(sqrt);
        // part II) full mu and sigma matrix
        Matrix & average_m_full = mu_full.OutputImageFromKernel[k];
        Matrix variance_m_full = sigma_square_full.OutputImageFromKernel[k];
        variance_m_full = variance_m_full + epsilon;
        variance_m_full(sqrt);

        // since one kernel corresponds to one pair of (\gamma, \beta)
        // so matrix average_m_active... etc must have dimension(1, 1)
        assert(average_m_active.Dimension().first == 1);
        assert(average_m_active.Dimension().second == 1);
        double _fMu_active = average_m_active[0][0];
        double _fVariance_active = variance_m_active[0][0];
        double _fMu_full = average_m_full[0][0];
        double _fVariance_full = variance_m_full[0][0];

        for(size_t i=0;i<batch_size;i++) // TO-DO: need to parallize
        {
            // part I) for active image
            Matrix m = __imageZ[i].OutputImageFromKernel[k]; // using copy assignment
            m = m - _fMu_active;
            m = m / _fVariance_active;
            __imageZ_BN[i].OutputImageFromKernel.push_back(m);

            // part II) for full image
            Matrix m_full = __imageZFull[i].OutputImageFromKernel[k];
            m_full = m_full - _fMu_full;
            m_full = m_full / _fVariance_full;
            __imageZFull_BN[i].OutputImageFromKernel.push_back(m_full);
        }
    }
    */

    // parallelize
    auto fill_z = [&](size_t start, size_t end)
    {
        for(size_t k=0;k<nKernel;k++)
        {
            // part I) active mu and sigma matrix
            Matrix & average_m_active = mu.OutputImageFromKernel[k];
            Matrix variance_m_active = sigma_square.OutputImageFromKernel[k];
            variance_m_active = variance_m_active + epsilon;
            variance_m_active(sqrt);
            // part II) full mu and sigma matrix
            Matrix & average_m_full = mu_full.OutputImageFromKernel[k];
            Matrix variance_m_full = sigma_square_full.OutputImageFromKernel[k];
            variance_m_full = variance_m_full + epsilon;
            variance_m_full(sqrt);

            // since one kernel corresponds to one pair of (\gamma, \beta)
            // so matrix average_m_active... etc must have dimension(1, 1)
            assert(average_m_active.Dimension().first == 1);
            assert(average_m_active.Dimension().second == 1);
            double _fMu_active = average_m_active[0][0];
            double _fVariance_active = variance_m_active[0][0];
            double _fMu_full = average_m_full[0][0];
            double _fVariance_full = variance_m_full[0][0];

            for(size_t i=start;i<end;i++) // TO-DO: need to parallize
            {
                // part I) for active image
                Matrix m = __imageZ[i].OutputImageFromKernel[k]; // using copy assignment
                m = m - _fMu_active;
                m = m / _fVariance_active;
                __imageZ_BN[i].OutputImageFromKernel.push_back(m);

                // part II) for full image
                Matrix m_full = __imageZFull[i].OutputImageFromKernel[k];
                m_full = m_full - _fMu_full;
                m_full = m_full / _fVariance_full;
                __imageZFull_BN[i].OutputImageFromKernel.push_back(m_full);
            }
        }
    };

#ifdef MULTI_THREAD
    vth.clear();
    for(int n = 0;n<NTHREAD;n++)
        vth.push_back(std::thread(fill_z, Range[n], Range[n+1]));
    for(auto &th: vth)
        th.join();
#else
    fill_z(0, batch_size);
#endif

    // 4) affine transformation to get y
    assert(__gamma_BN.size() == nKernel);
    assert(__beta_BN.size() == nKernel);

    /*
    for(size_t i = 0; i<batch_size;i++) // To-Do: need to parallize
    {
        for(size_t k=0;k<nKernel;k++)
        {
            double gamma = (__gamma_BN[k])[0][0];
            double beta = (__beta_BN[k])[0][0];
            // part I) active 
            Matrix m = __imageZ_BN[i].OutputImageFromKernel[k] * gamma;
            m = m + beta;
            __imageY_BN[i].OutputImageFromKernel.push_back(m);

            // part II) full
            m = __imageZFull_BN[i].OutputImageFromKernel[k] * gamma;
            m = m + beta;
            __imageYFull_BN[i].OutputImageFromKernel.push_back(m);
        }
    }
    */

    // parallelize
    auto fill_y = [&](size_t start, size_t end)
    {
        for(size_t i = start; i<end;i++) // To-Do: need to parallize
        {
            for(size_t k=0;k<nKernel;k++)
            {
                double gamma = (__gamma_BN[k])[0][0];
                double beta = (__beta_BN[k])[0][0];
                // part I) active 
                Matrix m = __imageZ_BN[i].OutputImageFromKernel[k] * gamma;
                m = m + beta;
                __imageY_BN[i].OutputImageFromKernel.push_back(m);

                // part II) full
                m = __imageZFull_BN[i].OutputImageFromKernel[k] * gamma;
                m = m + beta;
                __imageYFull_BN[i].OutputImageFromKernel.push_back(m);
            }
        }
    };
    vth.clear();
    for(int i=0;i<NTHREAD;i++)
        vth.push_back(std::thread(fill_y, Range[i], Range[i+1]));
    for(auto &i: vth) i.join();
}

// Assembly:
subq $0x198, %rsp            # imm = 0x198
movq %rdi, 0x190(%rsp)
movq 0x190(%rsp), %rax
movq %rax, 0x8(%rsp)
leaq 0x178(%rsp), %rdi
callq 0x1fda0
leaq 0x160(%rsp), %rdi
callq 0x1fda0
jmp 0x12868
movq 0x8(%rsp), %rdi
movq %rdi, %rsi
addq $0x150, %rsi            # imm = 0x150
movq (%rdi), %rax
movq 0x1b0(%rax), %rax
leaq 0x178(%rsp), %rdx
callq *%rax
jmp 0x1288d
movq 0x8(%rsp), %rdi
movq %rdi, %rsi
addq $0x150, %rsi            # imm = 0x150
movq (%rdi), %rax
movq 0x1c0(%rax), %rax
leaq 0x178(%rsp), %rdx
leaq 0x160(%rsp), %rcx
callq *%rax
jmp 0x128ba
leaq 0x138(%rsp), %rdi
callq 0x1fda0
jmp 0x128c9
leaq 0x120(%rsp), %rdi
callq 0x1fda0
jmp 0x128d8
movq 0x8(%rsp), %rdi
movq %rdi, %rsi
addq $0x1b0, %rsi            # imm = 0x1B0
movq (%rdi), %rax
movq 0x1b0(%rax), %rax
leaq 0x138(%rsp), %rdx
callq *%rax
jmp 0x128fd
movq 0x8(%rsp), %rdi
movq %rdi, %rsi
addq $0x1b0, %rsi            # imm = 0x1B0
movq (%rdi), %rax
movq 0x1c0(%rax), %rax
leaq 0x138(%rsp), %rdx
leaq 0x120(%rsp), %rcx
callq *%rax
jmp 0x1292a
movq 0x8(%rsp), %rdi
addq $0x278, %rdi            # imm = 0x278
xorl %eax, %eax
movl %eax, %esi
callq 0x1fd40
movq %rax, %rdi
leaq 0x178(%rsp), %rsi
callq 0x1fdb0
jmp 0x12951
movq 0x8(%rsp), %rdi
addq $0x2a8, %rdi            # imm = 0x2A8
xorl %eax, %eax
movl %eax, %esi
callq 0x1fd40
movq %rax, %rdi
leaq 0x160(%rsp), %rsi
callq 0x1fdb0
jmp 0x12978
movq 0x8(%rsp), %rdi
addq $0x290, %rdi            # imm = 0x290
xorl %eax, %eax
movl %eax, %esi
callq 0x1fd40
movq %rax, %rdi
leaq 0x138(%rsp), %rsi
callq 0x1fdb0
jmp 0x1299f
movq 0x8(%rsp), %rdi
addq $0x2c0, %rdi            # imm = 0x2C0
xorl %eax, %eax
movl %eax, %esi
callq 0x1fd40
movq %rax, %rdi
leaq 0x120(%rsp), %rsi
callq 0x1fdb0
jmp 0x129c6
movq 0x8(%rsp), %rdi
addq $0x150, %rdi            # imm = 0x150
callq 0x1ff40
movq 0x8(%rsp), %rdi
movq %rax, 0x118(%rsp)
movsd 0x3f64c(%rip), %xmm0    # 0x52038
movsd %xmm0, 0x110(%rsp)
addq $0x150, %rdi            # imm = 0x150
xorl %eax, %eax
movl %eax, %esi
callq 0x1fd40
movq %rax, %rdi
callq 0x1f9c0
movq %rax, 0x108(%rsp)
leaq 0xf0(%rsp), %rdi
callq 0x1fac0
movl $0x0, 0xbc(%rsp)
cmpl $0x8, 0xbc(%rsp)
jge 0x12adb
movq 0x118(%rsp), %rax
shrq $0x3, %rax
movslq 0xbc(%rsp), %rcx
imulq %rcx, %rax
movl %eax, %ecx
movslq 0xbc(%rsp), %rax
movl %ecx, 0xc0(%rsp,%rax,4)
movl 0xbc(%rsp), %eax
addl $0x1, %eax
movl %eax, 0xbc(%rsp)
jmp 0x12a2d
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x158(%rsp)
movl %eax, 0x154(%rsp)
jmp 0x12e75
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x158(%rsp)
movl %eax, 0x154(%rsp)
jmp 0x12e68
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x158(%rsp)
movl %eax, 0x154(%rsp)
jmp 0x12e5b
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x158(%rsp)
movl %eax, 0x154(%rsp)
jmp 0x12e4e
movq 0x8(%rsp), %rax
movq 0x118(%rsp), %rcx
movl %ecx, 0xe0(%rsp)
leaq 0x108(%rsp), %rcx
movq %rcx, 0x80(%rsp)
leaq 0x178(%rsp), %rcx
movq %rcx, 0x88(%rsp)
leaq 0x160(%rsp), %rcx
movq %rcx, 0x90(%rsp)
leaq 0x110(%rsp), %rcx
movq %rcx, 0x98(%rsp)
leaq 0x138(%rsp), %rcx
movq %rcx, 0xa0(%rsp)
leaq 0x120(%rsp), %rcx
movq %rcx, 0xa8(%rsp)
movq %rax, 0xb0(%rsp)
leaq 0xf0(%rsp), %rdi
callq 0x1fbf0
movl $0x0, 0x7c(%rsp)
cmpl $0x8, 0x7c(%rsp)
jge 0x12c07
movslq 0x7c(%rsp), %rax
leaq 0xc0(%rsp,%rax,4), %rdx
leaq 0xc4(%rsp,%rax,4), %rcx
leaq 0x70(%rsp), %rdi
leaq 0x80(%rsp), %rsi
callq 0x12e90
jmp 0x12ba0
leaq 0xf0(%rsp), %rdi
leaq 0x70(%rsp), %rsi
callq 0x1fad0
jmp 0x12bb4
leaq 0x70(%rsp), %rdi
callq 0x1fb00
movl 0x7c(%rsp), %eax
addl $0x1, %eax
movl %eax, 0x7c(%rsp)
jmp 0x12b6c
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x158(%rsp)
movl %eax, 0x154(%rsp)
jmp 0x12e41
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x158(%rsp)
movl %eax, 0x154(%rsp)
leaq 0x70(%rsp), %rdi
callq 0x1fb00
jmp 0x12e41
leaq 0xf0(%rsp), %rax
movq %rax, 0x68(%rsp)
movq 0x68(%rsp), %rdi
callq 0x1fb20
movq %rax, 0x60(%rsp)
movq 0x68(%rsp), %rdi
callq 0x1fb50
movq %rax, 0x58(%rsp)
leaq 0x60(%rsp), %rdi
leaq 0x58(%rsp), %rsi
callq 0x1fb80
testb $0x1, %al
jne 0x12c47
jmp 0x12c70
leaq 0x60(%rsp), %rdi
callq 0x1fbc0
movq %rax, 0x50(%rsp)
movq 0x50(%rsp), %rdi
callq 0x43c0
jmp 0x12c62
jmp 0x12c64
leaq 0x60(%rsp), %rdi
callq 0x1fbd0
jmp 0x12c32
movq 0x8(%rsp), %rdi
addq $0x308, %rdi            # imm = 0x308
callq 0x1f9c0
cmpq 0x108(%rsp), %rax
jne 0x12c8d
jmp 0x12cac
leaq 0x3ff5b(%rip), %rdi     # 0x52bef
leaq 0x3f537(%rip), %rsi     # 0x521d2
movl $0x8fe, %edx            # imm = 0x8FE
leaq 0x3ff65(%rip), %rcx     # 0x52c0c
callq 0x4170
movq 0x8(%rsp), %rdi
addq $0x320, %rdi            # imm = 0x320
callq 0x1f9c0
cmpq 0x108(%rsp), %rax
jne 0x12cc9
jmp 0x12ce8
leaq 0x3ff7e(%rip), %rdi     # 0x52c4e
leaq 0x3f4fb(%rip), %rsi     # 0x521d2
movl $0x8ff, %edx            # imm = 0x8FF
leaq 0x3ff29(%rip), %rcx     # 0x52c0c
callq 0x4170
movq 0x8(%rsp), %rax
leaq 0x108(%rsp), %rcx
movq %rcx, 0x40(%rsp)
movq %rax, 0x48(%rsp)
leaq 0xf0(%rsp), %rdi
callq 0x1fbf0
movl $0x0, 0x3c(%rsp)
cmpl $0x8, 0x3c(%rsp)
jge 0x12d8f
movslq 0x3c(%rsp), %rax
leaq 0xc0(%rsp,%rax,4), %rdx
leaq 0xc4(%rsp,%rax,4), %rcx
leaq 0x30(%rsp), %rdi
leaq 0x40(%rsp), %rsi
callq 0x12f60
jmp 0x12d41
leaq 0xf0(%rsp), %rdi
leaq 0x30(%rsp), %rsi
callq 0x1fad0
jmp 0x12d55
leaq 0x30(%rsp), %rdi
callq 0x1fb00
movl 0x3c(%rsp), %eax
addl $0x1, %eax
movl %eax, 0x3c(%rsp)
jmp 0x12d14
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x158(%rsp)
movl %eax, 0x154(%rsp)
leaq 0x30(%rsp), %rdi
callq 0x1fb00
jmp 0x12e41
leaq 0xf0(%rsp), %rax
movq %rax, 0x28(%rsp)
movq 0x28(%rsp), %rdi
callq 0x1fb20
movq %rax, 0x20(%rsp)
movq 0x28(%rsp), %rdi
callq 0x1fb50
movq %rax, 0x18(%rsp)
leaq 0x20(%rsp), %rdi
leaq 0x18(%rsp), %rsi
callq 0x1fb80
testb $0x1, %al
jne 0x12dcf
jmp 0x12df8
leaq 0x20(%rsp), %rdi
callq 0x1fbc0
movq %rax, 0x10(%rsp)
movq 0x10(%rsp), %rdi
callq 0x43c0
jmp 0x12dea
jmp 0x12dec
leaq 0x20(%rsp), %rdi
callq 0x1fbd0
jmp 0x12dba
leaq 0xf0(%rsp), %rdi
callq 0x1fc10
leaq 0x120(%rsp), %rdi
callq 0x1fde0
leaq 0x138(%rsp), %rdi
callq 0x1fde0
leaq 0x160(%rsp), %rdi
callq 0x1fde0
leaq 0x178(%rsp), %rdi
callq 0x1fde0
addq $0x198, %rsp            # imm = 0x198
retq
leaq 0xf0(%rsp), %rdi
callq 0x1fc10
leaq 0x120(%rsp), %rdi
callq 0x1fde0
leaq 0x138(%rsp), %rdi
callq 0x1fde0
leaq 0x160(%rsp), %rdi
callq 0x1fde0
leaq 0x178(%rsp), %rdi
callq 0x1fde0
movq 0x158(%rsp), %rdi
callq 0x43f0
nop