// File: xbai0624[P]dnn/src/ConstructLayer.cpp
// Function: ConstructLayer::BatchNormalization_UpdateZ_1DLayer()
void ConstructLayer::BatchNormalization_UpdateZ_1DLayer()
{
    // 1) get mu and sigma_square
    // part I) active
    Images mu, sigma_square;
    GetMu_BFor1DLayer(__imageZ, mu); // use active z image
    GetSigmaSquareFor1DLayer(__imageZ, mu, sigma_square);
    // part II) full
    Images mu_full, sigma_square_full;
    GetMu_BFor1DLayer(__imageZFull, mu_full);
    GetSigmaSquareFor1DLayer(__imageZFull, mu_full, sigma_square_full);

    // 2) save mu and sigma_square ---- Do I need to calculate both active and full? This is time consuming.
    // 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();
    assert(__imageZ[0].OutputImageFromKernel.size() == 1); // make sure 1D layer
    double epsilon = 1e-8; // avoid divide by zero issue; a hard-coded parameter from algorithm

    // part I) active mu and sigma matrix
    Matrix & average_m_active = mu.OutputImageFromKernel[0];
    Matrix variance_m_active = sigma_square.OutputImageFromKernel[0];
    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[0];
    Matrix variance_m_full = sigma_square_full.OutputImageFromKernel[0];
    variance_m_full = variance_m_full + epsilon;
    variance_m_full(sqrt);
    /*
    for(size_t i=0;i<batch_size;i++) // TO-DO: need to parallize
    {
        // part I) for active image
        Matrix m = __imageZ[i].OutputImageFromKernel[0]; // using copy assigment
        m = m - average_m_active;
        m = m / variance_m_active;
        Images tmp;
        tmp.OutputImageFromKernel.push_back(m);
        __imageZ_BN[i]=tmp;


        // part II) for full image
        Matrix m_full = __imageZFull[i].OutputImageFromKernel[0]; // using copy assigment
        m_full = m_full - average_m_full;
        m_full = m_full / variance_m_full;
        Images tmp_full;
        tmp_full.OutputImageFromKernel.push_back(m_full);
        __imageZFull_BN[i]=tmp_full;
    }
    */
    // parallelize
    auto fill_z = [&](size_t start, size_t end) 
    {
        for(size_t i=start;i<end;i++)
        {
            // part I) for active image
            Matrix m = __imageZ[i].OutputImageFromKernel[0]; // using copy assigment
            m = m - average_m_active;
            m = m / variance_m_active;
            Images tmp;
            tmp.OutputImageFromKernel.push_back(m);
            __imageZ_BN[i]=tmp;


            // part II) for full image
            Matrix m_full = __imageZFull[i].OutputImageFromKernel[0]; // using copy assigment
            m_full = m_full - average_m_full;
            m_full = m_full / variance_m_full;
            Images tmp_full;
            tmp_full.OutputImageFromKernel.push_back(m_full);
            __imageZFull_BN[i]=tmp_full;
        }
    };

#ifdef MULTI_THREAD
    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 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() == 1); // for 1D layer, size = 1; for 2D layer, size = nKernel
    assert(__beta_BN.size() == 1); 
    /*
    for(size_t i=0;i<batch_size;i++) // TO-DO: need to parallize
    {
        // part I) active
        Matrix m = __imageZ_BN[i].OutputImageFromKernel[0] ^ __gamma_BN[0];
        m = m + __beta_BN[0];
        Images tmp;
        tmp.OutputImageFromKernel.push_back(m);
        __imageY_BN[i]=tmp;

        // part II) full
        m = __imageZFull_BN[i].OutputImageFromKernel[0] ^ __gamma_BN[0];
        m = m + __beta_BN[0];
        tmp.OutputImageFromKernel.clear();
        tmp.OutputImageFromKernel.push_back(m);
        __imageYFull_BN[i]=tmp;
    }
    */
    // parallelize
    auto fill_y = [&](size_t start, size_t end)
    {
        for(size_t i=start;i<end;i++)
        {
            // part I) active
            Matrix m = __imageZ_BN[i].OutputImageFromKernel[0] ^ __gamma_BN[0];
            m = m + __beta_BN[0];
            Images tmp;
            tmp.OutputImageFromKernel.push_back(m);
            __imageY_BN[i]=tmp;

            // part II) full
            m = __imageZFull_BN[i].OutputImageFromKernel[0] ^ __gamma_BN[0];
            m = m + __beta_BN[0];
            tmp.OutputImageFromKernel.clear();
            tmp.OutputImageFromKernel.push_back(m);
            __imageYFull_BN[i]=tmp;
        }
    };

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

// Assembly:
subq $0x1f8, %rsp            # imm = 0x1F8
movq %rdi, 0x1f0(%rsp)
movq 0x1f0(%rsp), %rax
movq %rax, 0x18(%rsp)
leaq 0x1d8(%rsp), %rdi
callq 0x1fda0
leaq 0x1c0(%rsp), %rdi
callq 0x1fda0
jmp 0x11e78
movq 0x18(%rsp), %rdi
movq %rdi, %rsi
addq $0x150, %rsi            # imm = 0x150
movq (%rdi), %rax
movq 0x1a8(%rax), %rax
leaq 0x1d8(%rsp), %rdx
callq *%rax
jmp 0x11e9d
movq 0x18(%rsp), %rdi
movq %rdi, %rsi
addq $0x150, %rsi            # imm = 0x150
movq (%rdi), %rax
movq 0x1b8(%rax), %rax
leaq 0x1d8(%rsp), %rdx
leaq 0x1c0(%rsp), %rcx
callq *%rax
jmp 0x11eca
leaq 0x198(%rsp), %rdi
callq 0x1fda0
jmp 0x11ed9
leaq 0x180(%rsp), %rdi
callq 0x1fda0
jmp 0x11ee8
movq 0x18(%rsp), %rdi
movq %rdi, %rsi
addq $0x1b0, %rsi            # imm = 0x1B0
movq (%rdi), %rax
movq 0x1a8(%rax), %rax
leaq 0x198(%rsp), %rdx
callq *%rax
jmp 0x11f0d
movq 0x18(%rsp), %rdi
movq %rdi, %rsi
addq $0x1b0, %rsi            # imm = 0x1B0
movq (%rdi), %rax
movq 0x1b8(%rax), %rax
leaq 0x198(%rsp), %rdx
leaq 0x180(%rsp), %rcx
callq *%rax
jmp 0x11f3a
movq 0x18(%rsp), %rdi
addq $0x278, %rdi            # imm = 0x278
xorl %eax, %eax
movl %eax, %esi
callq 0x1fd40
movq %rax, %rdi
leaq 0x1d8(%rsp), %rsi
callq 0x1fdb0
jmp 0x11f61
movq 0x18(%rsp), %rdi
addq $0x2a8, %rdi            # imm = 0x2A8
xorl %eax, %eax
movl %eax, %esi
callq 0x1fd40
movq %rax, %rdi
leaq 0x1c0(%rsp), %rsi
callq 0x1fdb0
jmp 0x11f88
movq 0x18(%rsp), %rdi
addq $0x290, %rdi            # imm = 0x290
xorl %eax, %eax
movl %eax, %esi
callq 0x1fd40
movq %rax, %rdi
leaq 0x198(%rsp), %rsi
callq 0x1fdb0
jmp 0x11faf
movq 0x18(%rsp), %rdi
addq $0x2c0, %rdi            # imm = 0x2C0
xorl %eax, %eax
movl %eax, %esi
callq 0x1fd40
movq %rax, %rdi
leaq 0x180(%rsp), %rsi
callq 0x1fdb0
jmp 0x11fd6
movq 0x18(%rsp), %rdi
addq $0x150, %rdi            # imm = 0x150
callq 0x1ff40
movq 0x18(%rsp), %rdi
movq %rax, 0x178(%rsp)
addq $0x150, %rdi            # imm = 0x150
xorl %eax, %eax
movl %eax, %esi
callq 0x1fd40
movq %rax, %rdi
callq 0x1f9c0
cmpq $0x1, %rax
jne 0x12017
jmp 0x1209a
leaq 0x40b34(%rip), %rdi     # 0x52b52
leaq 0x401ad(%rip), %rsi     # 0x521d2
movl $0x801, %edx            # imm = 0x801
leaq 0x40b4f(%rip), %rcx     # 0x52b80
callq 0x4170
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x1b8(%rsp)
movl %eax, 0x1b4(%rsp)
jmp 0x12669
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x1b8(%rsp)
movl %eax, 0x1b4(%rsp)
jmp 0x1265c
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x1b8(%rsp)
movl %eax, 0x1b4(%rsp)
jmp 0x1264f
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x1b8(%rsp)
movl %eax, 0x1b4(%rsp)
jmp 0x12642
movabsq $0x3e45798ee2308c3a, %rax # imm = 0x3E45798EE2308C3A
movq %rax, 0x170(%rsp)
xorl %eax, %eax
movl %eax, %esi
movq %rsi, 0x10(%rsp)
leaq 0x1d8(%rsp), %rdi
callq 0x1f9e0
movq 0x10(%rsp), %rsi
movq %rax, 0x168(%rsp)
leaq 0x1c0(%rsp), %rdi
callq 0x1f9e0
movq %rax, %rsi
leaq 0x150(%rsp), %rdi
callq 0x3b810
jmp 0x120ee
movsd 0x170(%rsp), %xmm0
leaq 0x138(%rsp), %rdi
leaq 0x150(%rsp), %rsi
callq 0x3ca80
jmp 0x1210e
leaq 0x150(%rsp), %rdi
leaq 0x138(%rsp), %rsi
callq 0x1fa20
jmp 0x12125
leaq 0x138(%rsp), %rdi
callq 0x3bb80
movq 0x58e97(%rip), %rsi     # 0x6afd0
leaq 0x150(%rsp), %rdi
callq 0x3dc20
jmp 0x12148
xorl %eax, %eax
movl %eax, %esi
movq %rsi, 0x8(%rsp)
leaq 0x198(%rsp), %rdi
callq 0x1f9e0
movq 0x8(%rsp), %rsi
movq %rax, 0x130(%rsp)
leaq 0x180(%rsp), %rdi
callq 0x1f9e0
movq %rax, %rsi
leaq 0x118(%rsp), %rdi
callq 0x3b810
jmp 0x1218a
movsd 0x170(%rsp), %xmm0
leaq 0x100(%rsp), %rdi
leaq 0x118(%rsp), %rsi
callq 0x3ca80
jmp 0x121aa
leaq 0x118(%rsp), %rdi
leaq 0x100(%rsp), %rsi
callq 0x1fa20
jmp 0x121c1
leaq 0x100(%rsp), %rdi
callq 0x3bb80
movq 0x58dfb(%rip), %rsi     # 0x6afd0
leaq 0x118(%rsp), %rdi
callq 0x3dc20
jmp 0x121e4
movq 0x18(%rsp), %rax
movq %rax, 0xd8(%rsp)
movq 0x168(%rsp), %rax
movq %rax, 0xe0(%rsp)
leaq 0x150(%rsp), %rax
movq %rax, 0xe8(%rsp)
movq 0x130(%rsp), %rax
movq %rax, 0xf0(%rsp)
leaq 0x118(%rsp), %rax
movq %rax, 0xf8(%rsp)
leaq 0xc0(%rsp), %rdi
callq 0x1fac0
movl $0x0, 0x8c(%rsp)
cmpl $0x8, 0x8c(%rsp)
jge 0x12311
movq 0x178(%rsp), %rax
shrq $0x3, %rax
movslq 0x8c(%rsp), %rcx
imulq %rcx, %rax
movl %eax, %ecx
movslq 0x8c(%rsp), %rax
movl %ecx, 0x90(%rsp,%rax,4)
movl 0x8c(%rsp), %eax
addl $0x1, %eax
movl %eax, 0x8c(%rsp)
jmp 0x12249
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x1b8(%rsp)
movl %eax, 0x1b4(%rsp)
jmp 0x12635
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x1b8(%rsp)
movl %eax, 0x1b4(%rsp)
leaq 0x138(%rsp), %rdi
callq 0x3bb80
jmp 0x12635
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x1b8(%rsp)
movl %eax, 0x1b4(%rsp)
jmp 0x12628
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x1b8(%rsp)
movl %eax, 0x1b4(%rsp)
leaq 0x100(%rsp), %rdi
callq 0x3bb80
jmp 0x12628
movq 0x178(%rsp), %rax
movl %eax, 0xb0(%rsp)
movq $0x0, 0x80(%rsp)
cmpq $0x8, 0x80(%rsp)
jae 0x123d7
movq 0x80(%rsp), %rax
leaq 0x90(%rsp,%rax,4), %rdx
leaq 0x94(%rsp,%rax,4), %rcx
leaq 0x78(%rsp), %rdi
leaq 0xd8(%rsp), %rsi
callq 0x12690
jmp 0x12367
leaq 0xc0(%rsp), %rdi
leaq 0x78(%rsp), %rsi
callq 0x1fad0
jmp 0x1237b
leaq 0x78(%rsp), %rdi
callq 0x1fb00
movq 0x80(%rsp), %rax
addq $0x1, %rax
movq %rax, 0x80(%rsp)
jmp 0x1232c
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x1b8(%rsp)
movl %eax, 0x1b4(%rsp)
jmp 0x1261b
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x1b8(%rsp)
movl %eax, 0x1b4(%rsp)
leaq 0x78(%rsp), %rdi
callq 0x1fb00
jmp 0x1261b
leaq 0xc0(%rsp), %rax
movq %rax, 0x70(%rsp)
movq 0x70(%rsp), %rdi
callq 0x1fb20
movq %rax, 0x68(%rsp)
movq 0x70(%rsp), %rdi
callq 0x1fb50
movq %rax, 0x60(%rsp)
leaq 0x68(%rsp), %rdi
leaq 0x60(%rsp), %rsi
callq 0x1fb80
testb $0x1, %al
jne 0x12417
jmp 0x12440
leaq 0x68(%rsp), %rdi
callq 0x1fbc0
movq %rax, 0x58(%rsp)
movq 0x58(%rsp), %rdi
callq 0x43c0
jmp 0x12432
jmp 0x12434
leaq 0x68(%rsp), %rdi
callq 0x1fbd0
jmp 0x12402
movq 0x18(%rsp), %rdi
addq $0x308, %rdi            # imm = 0x308
callq 0x1f9c0
cmpq $0x1, %rax
jne 0x12459
jmp 0x12478
leaq 0x40762(%rip), %rdi     # 0x52bc2
leaq 0x3fd6b(%rip), %rsi     # 0x521d2
movl $0x84c, %edx            # imm = 0x84C
leaq 0x4070d(%rip), %rcx     # 0x52b80
callq 0x4170
movq 0x18(%rsp), %rdi
addq $0x320, %rdi            # imm = 0x320
callq 0x1f9c0
cmpq $0x1, %rax
jne 0x12491
jmp 0x124b0
leaq 0x40741(%rip), %rdi     # 0x52bd9
leaq 0x3fd33(%rip), %rsi     # 0x521d2
movl $0x84d, %edx            # imm = 0x84D
leaq 0x406d5(%rip), %rcx     # 0x52b80
callq 0x4170
movq 0x18(%rsp), %rax
movq %rax, 0x50(%rsp)
leaq 0xc0(%rsp), %rdi
callq 0x1fbf0
movq $0x0, 0x48(%rsp)
cmpq $0x8, 0x48(%rsp)
jae 0x1254f
movq 0x48(%rsp), %rax
leaq 0x90(%rsp,%rax,4), %rdx
leaq 0x94(%rsp,%rax,4), %rcx
leaq 0x40(%rsp), %rdi
leaq 0x50(%rsp), %rsi
callq 0x12760
jmp 0x124fe
leaq 0xc0(%rsp), %rdi
leaq 0x40(%rsp), %rsi
callq 0x1fad0
jmp 0x12512
leaq 0x40(%rsp), %rdi
callq 0x1fb00
movq 0x48(%rsp), %rax
addq $0x1, %rax
movq %rax, 0x48(%rsp)
jmp 0x124d0
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x1b8(%rsp)
movl %eax, 0x1b4(%rsp)
leaq 0x40(%rsp), %rdi
callq 0x1fb00
jmp 0x1261b
leaq 0xc0(%rsp), %rax
movq %rax, 0x38(%rsp)
movq 0x38(%rsp), %rdi
callq 0x1fb20
movq %rax, 0x30(%rsp)
movq 0x38(%rsp), %rdi
callq 0x1fb50
movq %rax, 0x28(%rsp)
leaq 0x30(%rsp), %rdi
leaq 0x28(%rsp), %rsi
callq 0x1fb80
testb $0x1, %al
jne 0x1258f
jmp 0x125b8
leaq 0x30(%rsp), %rdi
callq 0x1fbc0
movq %rax, 0x20(%rsp)
movq 0x20(%rsp), %rdi
callq 0x43c0
jmp 0x125aa
jmp 0x125ac
leaq 0x30(%rsp), %rdi
callq 0x1fbd0
jmp 0x1257a
leaq 0xc0(%rsp), %rdi
callq 0x1fc10
leaq 0x118(%rsp), %rdi
callq 0x3bb80
leaq 0x150(%rsp), %rdi
callq 0x3bb80
leaq 0x180(%rsp), %rdi
callq 0x1fde0
leaq 0x198(%rsp), %rdi
callq 0x1fde0
leaq 0x1c0(%rsp), %rdi
callq 0x1fde0
leaq 0x1d8(%rsp), %rdi
callq 0x1fde0
addq $0x1f8, %rsp            # imm = 0x1F8
retq
leaq 0xc0(%rsp), %rdi
callq 0x1fc10
leaq 0x118(%rsp), %rdi
callq 0x3bb80
leaq 0x150(%rsp), %rdi
callq 0x3bb80
leaq 0x180(%rsp), %rdi
callq 0x1fde0
leaq 0x198(%rsp), %rdi
callq 0x1fde0
leaq 0x1c0(%rsp), %rdi
callq 0x1fde0
leaq 0x1d8(%rsp), %rdi
callq 0x1fde0
movq 0x1b8(%rsp), %rdi
callq 0x43f0
nopw %cs:(%rax,%rax)