// File: /workspace/github_cpp_makefile_star2_O2/xbai0624[P]dnn/src/ConstructLayer.cpp
// Function: ConstructLayer::AdamOptimizer(Matrix const&, int)
Matrix ConstructLayer::AdamOptimizer(const Matrix &dw, int kernel_index)
{
    //cout<<"----------------------------------------------------------"<<endl;
    // Adam optimizer, refer to: https://arxiv.org/abs/1412.6980
    // sanity check
    assert(Momentum_1st_order.size() == Momentum_2nd_order.size());
    if(Momentum_1st_order.size() == 0)
    {
        // initiaize with 0's
        if(__type == LayerType::fullyConnected || __type == LayerType::output || __type == LayerType::input)
        {
            assert(kernel_index == 0);
            Momentum_1st_order.resize(1, Matrix(dw.Dimension(), 0));
            Momentum_2nd_order.resize(1, Matrix(dw.Dimension(), 0));
        }
        else if(__type == LayerType::cnn || __type == LayerType::pooling)
        {
            assert(__n_kernels_cnn > 0);
            Momentum_1st_order.resize(__n_kernels_cnn, Matrix(dw.Dimension(), 0));
            Momentum_2nd_order.resize(__n_kernels_cnn, Matrix(dw.Dimension(), 0));
        }
    }

    // g_t
    Matrix g_t = dw;
    //cout<<"g_t: "<<endl<<g_t<<endl;

    // these parameters are from paper: https://arxiv.org/abs/1412.6980
    //     according to the ref: learning rate is best to be 0.001
    double beta1 = 0.9, beta2 = 0.999, epsilon = 1e-8;
    // get m_t
    Matrix & m_t_1 = Momentum_1st_order[kernel_index];
    //cout<<"m_t_1: "<<endl<<m_t_1<<endl;
    //cout<<" 1 - beta1: "<<1.-beta1<<endl;
    Matrix m_tmp = g_t * (1. - beta1);
    //cout<<"m_tmp: "<<endl<<m_tmp<<endl;
    Matrix m_t = (m_t_1 * beta1) + m_tmp;
    //Matrix debug_tmp = m_t_1 * beta1;
    //cout<<"m_t_1*beta1: "<<endl<<(debug_tmp)<<endl;
    //cout<<"m_t: = (m_t_1 * beta1) + g_t * (1 - beta1): "<<endl<<m_t<<endl;
    // save m_t
    Momentum_1st_order[kernel_index] = m_t; // update 1st order momentum
    // get v_t
    Matrix & v_t_1 = Momentum_2nd_order[kernel_index];
    //cout<<"v_t_1: "<<endl<<v_t_1<<endl;
    m_tmp = (g_t^g_t) * (1. - beta2);
    //cout<<"g_t^2 * (1. - beta2(0.999))"<<endl<<m_tmp<<endl;
    Matrix v_t = (v_t_1 *beta2) + m_tmp;
    //debug_tmp = v_t_1 *beta2;
    //cout<<"v_t_1 * beta2:"<<endl<<debug_tmp<<endl;
    //cout<<"vt = v_t_1*beta2 + g_t^2*(1-beta2): "<<endl<<v_t<<endl;
    // save v_t
    Momentum_2nd_order[kernel_index] = v_t; // update 2nd order momentum

    // get bias corrected m_t and v-t
    __beta1_to_power_t *= beta1;
    __beta2_to_power_t *= beta2;
    Matrix bias_corrected_1st_order_moment = m_t / (1. - __beta1_to_power_t);
    //cout<<"bias corrected m_t: /"<<(1. - __beta1_to_power_t)<<endl<<bias_corrected_1st_order_moment<<endl;
    Matrix bias_corrected_2nd_order_moment = v_t / (1. - __beta2_to_power_t);
    //cout<<"bias corrected v_t: /"<<(1. - __beta2_to_power_t)<<endl<<bias_corrected_2nd_order_moment<<endl;

    // get d_theta_t 
    Matrix m_epsilon(g_t.Dimension(), epsilon);
    bias_corrected_2nd_order_moment(__AdamSquareRoot);
    //cout<<"sqrt(^v_t): "<<endl<<bias_corrected_2nd_order_moment<<endl;
    Matrix square_root_of_vt = bias_corrected_2nd_order_moment;
    Matrix denominator = square_root_of_vt + m_epsilon;
    //cout<<"m_epsilon: "<<endl<<m_epsilon<<endl;
    //cout<<"sqrt(^v_t): + epsilon"<<endl<<denominator<<endl;
    Matrix ratio = bias_corrected_1st_order_moment / denominator;
    //cout<<"^m_t/(sqrt(^v_t) + epsilon):"<<endl<<ratio<<endl;

    // according to the reference, alpha best to be 0.001, however, 
    // based on my test, it is too slow, alpha = 0.06 works best
    double alpha = __learningRate;
    Matrix dw_Adam = ratio * alpha;
    //Matrix dw_Adam = ratio * __learningRate;
    //cout<<"learning rate: "<<__learningRate<<endl;
    //cout<<"dw_Adam: "<<endl<<dw_Adam<<endl;
    //getchar();

    return dw_Adam;
}

// Assembly:
pushq %rbp
pushq %r15
pushq %r14
pushq %r13
pushq %r12
pushq %rbx
subq $0x108, %rsp            # imm = 0x108
movl %ecx, %ebp
movq %rdx, %r13
movq %rsi, %r14
movq %rdi, 0x50(%rsp)
movq 0x448(%rsi), %rax
movq 0x450(%rsi), %rcx
movq %rcx, %rdx
subq %rax, %rdx
movq 0x468(%rsi), %rsi
subq 0x460(%r14), %rsi
cmpq %rsi, %rdx
jne 0xf2ac
leaq 0x448(%r14), %rbx
leaq 0x460(%r14), %r15
cmpq %rax, %rcx
jne 0xf00f
movl 0xc(%r14), %eax
cmpl $0x4, %eax
ja 0xf00f
pushq $0x19
popq %rcx
btl %eax, %ecx
jae 0xef92
testl %ebp, %ebp
jne 0xf2cb
movq %rbx, %r12
movq %r13, %rdi
callq 0x17a6a
leaq 0x8(%rsp), %rbx
xorpd %xmm0, %xmm0
movq %rbx, %rdi
movq %rax, %rsi
callq 0x1787a
pushq $0x1
popq %rsi
movq %r12, %rdi
movq %rbx, %rdx
callq 0x11ccc
leaq 0x8(%rsp), %rdi
callq 0x17a9a
movq %r13, %rdi
callq 0x17a6a
leaq 0x8(%rsp), %rbx
xorpd %xmm0, %xmm0
movq %rbx, %rdi
movq %rax, %rsi
callq 0x1787a
pushq $0x1
popq %rsi
movq %r15, %rdi
movq %rbx, %rdx
callq 0x11ccc
jmp 0xf002
movq %rbx, %r12
movq 0x18(%r14), %rbx
testq %rbx, %rbx
je 0xf2ea
movq %r13, %rdi
callq 0x17a6a
leaq 0x8(%rsp), %rdi
xorpd %xmm0, %xmm0
movq %rax, %rsi
callq 0x1787a
movq %r12, %rdi
movq %rbx, %rsi
leaq 0x8(%rsp), %rdx
callq 0x11ccc
leaq 0x8(%rsp), %rdi
callq 0x17a9a
movq 0x18(%r14), %rbx
movq %r13, %rdi
callq 0x17a6a
leaq 0x8(%rsp), %rdi
xorpd %xmm0, %xmm0
movq %rax, %rsi
callq 0x1787a
movq %r15, %rdi
movq %rbx, %rsi
leaq 0x8(%rsp), %rdx
callq 0x11ccc
leaq 0x8(%rsp), %rdi
callq 0x17a9a
movq %r12, %rbx
movq %r14, 0x58(%rsp)
movq %r15, %r12
leaq 0x8(%rsp), %r15
movq %r15, %rdi
movq %r13, %rsi
callq 0x1796e
movq %rbx, %r14
movq (%rbx), %rbx
leaq 0x60(%rsp), %rdi
movsd 0x12526(%rip), %xmm0    # 0x21560
movq %r15, %rsi
callq 0x182dc
movslq %ebp, %rax
imulq $0x18, %rax, %r13
addq %r13, %rbx
leaq 0x20(%rsp), %rdi
movsd 0x1250f(%rip), %xmm0    # 0x21568
movq %rbx, %rsi
callq 0x182dc
leaq 0x90(%rsp), %rdi
leaq 0x20(%rsp), %rsi
leaq 0x60(%rsp), %rdx
callq 0x17cc4
movq %r12, %rbx
leaq 0x20(%rsp), %rdi
callq 0x17a9a
movq (%r14), %rdi
addq %r13, %rdi
leaq 0x90(%rsp), %rsi
callq 0x125b4
movq (%rbx), %r12
addq %r13, %r12
leaq 0x38(%rsp), %rdi
leaq 0x8(%rsp), %rdx
movq %rdx, %rsi
callq 0x17f64
movq 0x58(%rsp), %r14
leaq 0x20(%rsp), %rdi
leaq 0x38(%rsp), %rsi
movsd 0x124a9(%rip), %xmm0    # 0x21570
callq 0x182dc
leaq 0x60(%rsp), %rdi
leaq 0x20(%rsp), %rsi
callq 0x125b4
leaq 0x20(%rsp), %rdi
callq 0x17a9a
leaq 0x38(%rsp), %rdi
callq 0x17a9a
leaq 0x38(%rsp), %rdi
movsd 0x1247c(%rip), %xmm0    # 0x21578
movq %r12, %rsi
callq 0x182dc
leaq 0x20(%rsp), %rdi
leaq 0x38(%rsp), %rsi
leaq 0x60(%rsp), %rdx
callq 0x17cc4
leaq 0x38(%rsp), %rdi
callq 0x17a9a
addq (%rbx), %r13
leaq 0x20(%rsp), %rsi
movq %r13, %rdi
callq 0x125b4
movupd 0x478(%r14), %xmm1
mulpd 0x11f0d(%rip), %xmm1    # 0x21050
movupd %xmm1, 0x478(%r14)
movsd 0x123bc(%rip), %xmm0    # 0x21510
subsd %xmm1, %xmm0
leaq 0x38(%rsp), %rdi
leaq 0x90(%rsp), %rsi
callq 0x18348
movsd 0x1239e(%rip), %xmm0    # 0x21510
subsd 0x480(%r14), %xmm0
leaq 0x78(%rsp), %rdi
leaq 0x20(%rsp), %rsi
callq 0x18348
leaq 0x8(%rsp), %rdi
callq 0x17a6a
leaq 0xf0(%rsp), %rdi
movsd 0x12384(%rip), %xmm0    # 0x21528
movq %rax, %rsi
callq 0x1787a
leaq 0x308(%rip), %rsi       # 0xf4bb
leaq 0x78(%rsp), %rdi
callq 0x188a2
leaq 0xd8(%rsp), %rdi
leaq 0x78(%rsp), %rsi
callq 0x1796e
leaq 0xc0(%rsp), %rdi
leaq 0xd8(%rsp), %rsi
leaq 0xf0(%rsp), %rdx
callq 0x17cc4
leaq 0xa8(%rsp), %rdi
leaq 0x38(%rsp), %rsi
leaq 0xc0(%rsp), %rdx
callq 0x18018
movsd 0xc0(%r14), %xmm0
leaq 0xa8(%rsp), %rsi
movq 0x50(%rsp), %rbx
movq %rbx, %rdi
callq 0x182dc
leaq 0xa8(%rsp), %rdi
callq 0x17a9a
leaq 0xc0(%rsp), %rdi
callq 0x17a9a
leaq 0xd8(%rsp), %rdi
callq 0x17a9a
leaq 0xf0(%rsp), %rdi
callq 0x17a9a
leaq 0x78(%rsp), %rdi
callq 0x17a9a
leaq 0x38(%rsp), %rdi
callq 0x17a9a
leaq 0x20(%rsp), %rdi
callq 0x17a9a
leaq 0x90(%rsp), %rdi
callq 0x17a9a
leaq 0x60(%rsp), %rdi
callq 0x17a9a
leaq 0x8(%rsp), %rdi
callq 0x17a9a
movq %rbx, %rax
addq $0x108, %rsp            # imm = 0x108
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
retq
leaq 0x13c21(%rip), %rdi     # 0x22ed4
leaq 0x12420(%rip), %rsi     # 0x216da
leaq 0x13c4a(%rip), %rcx     # 0x22f0b
movl $0xe99, %edx            # imm = 0xE99
callq 0x4180
leaq 0x13c73(%rip), %rdi     # 0x22f45
leaq 0x12401(%rip), %rsi     # 0x216da
leaq 0x13c2b(%rip), %rcx     # 0x22f0b
movl $0xe9f, %edx            # imm = 0xE9F
callq 0x4180
leaq 0x13c66(%rip), %rdi     # 0x22f57
leaq 0x123e2(%rip), %rsi     # 0x216da
leaq 0x13c0c(%rip), %rcx     # 0x22f0b
movl $0xea5, %edx            # imm = 0xEA5
callq 0x4180
jmp 0xf3a6
jmp 0xf3a6
jmp 0xf3a6
jmp 0xf3a6
movq %rax, %rbx
leaq 0xa8(%rsp), %rdi
callq 0x17a9a
jmp 0xf332
movq %rax, %rbx
leaq 0xc0(%rsp), %rdi
callq 0x17a9a
jmp 0xf344
movq %rax, %rbx
leaq 0xd8(%rsp), %rdi
callq 0x17a9a
jmp 0xf358
jmp 0xf355
movq %rax, %rbx
leaq 0xf0(%rsp), %rdi
callq 0x17a9a
jmp 0xf3ae
movq %rax, %rbx
jmp 0xf3b8
jmp 0xf36e
movq %rax, %rbx
jmp 0xf3c2
jmp 0xf386
jmp 0xf392
movq %rax, %rbx
leaq 0x20(%rsp), %rdi
callq 0x17a9a
jmp 0xf389
movq %rax, %rbx
leaq 0x38(%rsp), %rdi
jmp 0xf3c7
jmp 0xf392
movq %rax, %rbx
jmp 0xf3cc
movq %rax, %rbx
leaq 0x20(%rsp), %rdi
jmp 0xf3d4
movq %rax, %rbx
jmp 0xf3d9
movq %rax, %rbx
jmp 0xf3e3
movq %rax, %rbx
leaq 0x78(%rsp), %rdi
callq 0x17a9a
leaq 0x38(%rsp), %rdi
callq 0x17a9a
leaq 0x20(%rsp), %rdi
callq 0x17a9a
leaq 0x90(%rsp), %rdi
callq 0x17a9a
leaq 0x60(%rsp), %rdi
callq 0x17a9a
leaq 0x8(%rsp), %rdi
callq 0x17a9a
movq %rbx, %rdi
callq 0x43b0
nop