// File: /workspace/github_c_makefile_star2_O1/debbiemarkslab[P]plmc/src/inference.c
// Function: SGDOptimize
void SGDOptimize(gradfun_t grad, void *data, numeric_t *x, const int n,
    const int maxIter, const numeric_t crit) {
    /* Opitimize an objective function by Stochastic Gradient Descent (Adam)
       Arguments:
            grad            gradient of objective
            data            pointer to data
            x               estimated parameters (length n)
            n               number of parameters
            eps             learning rate
            maxIter         maximum number of iterations
            crit            stop when ||grad|| / ||x|| < crit
    */
    // numeric_t ALPHA0 = 0.001;
    // numeric_t ALPHAT = 0.00001;
    numeric_t BETA1 = 0.9;
    numeric_t BETA2 = 0.99;
    numeric_t EPSILON = 1E-8;

    numeric_t *g = (numeric_t *) malloc(n * sizeof(numeric_t));
    numeric_t criterion = crit + 1.0;

    /* Begin profiling */
    struct timeval start;
    gettimeofday(&start, NULL);

    /* Initialize estimates of first and second moments of the gradient */
    numeric_t *meanX = (numeric_t *) malloc(n * sizeof(numeric_t));
    numeric_t *meanG = (numeric_t *) malloc(n * sizeof(numeric_t));
    numeric_t *squareG = (numeric_t *) malloc(n * sizeof(numeric_t));
    for (int i = 0; i < n; i++) meanX[i] = 0;
    for (int i = 0; i < n; i++) meanG[i] = 0;
    for (int i = 0; i < n; i++) squareG[i] = 0;

    /* Optimization loop */
    int t = 1;
    do {
        /* Estimate the gradient */
        for (int i = 0; i < n; i++) g[i] = 0;
        numeric_t f = grad(data, x, g, n);

        /* Update estimates of moments */
        for (int i = 0; i < n; i++)
            meanG[i] = BETA1 * meanG[i] + (1.0 - BETA1) * g[i];
        for (int i = 0; i < n; i++)
            squareG[i] = BETA2 * squareG[i] + (1.0 - BETA2) * g[i] * g[i];

        /* Update Q with Adam learning rates */
        // numeric_t schedule = ALPHA;
        // numeric_t frac = (numeric_t) t / (numeric_t) maxIter;
        // frac = floor(frac * 5) / 5.;
        // numeric_t schedule = exp((1 - frac) * log(ALPHA0) + frac * log(ALPHAT));
        // Anneal strategy #2
        numeric_t schedule = 0.01 * pow(0.5, (t / 50));
        // numeric_t schedule = 0.01;
        numeric_t alpha = schedule
                          * sqrt(1.0 - pow(BETA2, (numeric_t) t)) 
                              / (1.0 - pow(BETA1, (numeric_t) t));
        for (int i = 0; i < n; i++)
            x[i] -= meanG[i] * alpha / (sqrt(squareG[i]) + EPSILON);

        /* Update Polyak average */
        for (int i = 0; i < n; i++)
            meanX[i] = BETA1 * meanX[i] + (1 - BETA1) * x[i];

        /* Stopping criterion: ||grad(params)|| / ||params|| */
        numeric_t paramNorm = 1E-6;
        for (int i = 0; i < n; i++) paramNorm += fabs(x[i]) / (numeric_t) n;
        numeric_t gradNorm = 1E-6;
        for (int i = 0; i < n; i++)
            gradNorm += fabs(meanG[i]) / (numeric_t) n;
        criterion = gradNorm;

        if (t == 1)
            fprintf(stderr, "iter\ttime\tobj\t|x|\t|g|\tcrit\n");
            fprintf(stderr, "%d\t%.1f\t%.1f\t%.1f\t%.1f\t%.1f\n",
            t, ElapsedTime(&start), f, paramNorm, gradNorm, criterion);
        t++;
    } while (t <= maxIter && criterion > crit);
    // for (int i = 0; i < n; i++) x[i] = meanX[i] / ((numeric_t) t - 1);
    for (int i = 0; i < n; i++) x[i] = meanX[i];
    free(meanX);
    free(meanG);
    free(squareG);
    free(g);
}

// Assembly:
pushq %rbp
pushq %r15
pushq %r14
pushq %r13
pushq %r12
pushq %rbx
subq $0x78, %rsp
movsd %xmm0, 0x40(%rsp)
movl %r8d, 0x2c(%rsp)
movl %ecx, %ebp
movq %rdx, %r14
movq %rsi, 0x58(%rsp)
movq %rdi, 0x50(%rsp)
movslq %ecx, %rbx
leaq (,%rbx,8), %r13
movq %r13, %rdi
callq 0x11a0
movq %rax, 0x30(%rsp)
leaq 0x68(%rsp), %rdi
xorl %esi, %esi
callq 0x10f0
movq %r13, %rdi
callq 0x11a0
movq %rax, %r15
movq %r13, %rdi
callq 0x11a0
movq %rax, %r12
movq %r13, %rdi
callq 0x11a0
movq %rax, %r13
movl %ebp, 0x8(%rsp)
movl %ebp, %ebp
testl %ebx, %ebx
jle 0x1420
leaq (,%rbp,8), %rdx
movq %r15, %rdi
xorl %esi, %esi
callq 0x1110
cmpl $0x0, 0x8(%rsp)
jle 0x1439
leaq (,%rbp,8), %rdx
movq %r12, %rdi
xorl %esi, %esi
callq 0x1110
cmpl $0x0, 0x8(%rsp)
movq 0x30(%rsp), %rbx
jle 0x1457
leaq (,%rbp,8), %rdx
movq %r13, %rdi
xorl %esi, %esi
callq 0x1110
cvtsi2sdl 0x8(%rsp), %xmm0
movsd %xmm0, 0x48(%rsp)
leaq (,%rbp,8), %rax
movq %rax, 0x38(%rsp)
pushq $0x1
popq %rax
movl %eax, 0xc(%rsp)
cmpl $0x0, 0x8(%rsp)
jle 0x148d
movq %rbx, %rdi
xorl %esi, %esi
movq 0x38(%rsp), %rdx
callq 0x1110
movq 0x58(%rsp), %rdi
movq %r14, %rsi
movq %rbx, %rdx
movl 0x8(%rsp), %ecx
callq *0x50(%rsp)
movsd %xmm0, 0x60(%rsp)
cmpl $0x0, 0x8(%rsp)
movsd 0xcff5(%rip), %xmm2     # 0xe4a8
movsd 0xcff5(%rip), %xmm3     # 0xe4b0
jle 0x14e4
xorl %eax, %eax
movsd (%rbx,%rax,8), %xmm0
mulsd %xmm2, %xmm0
movsd (%r12,%rax,8), %xmm1
mulsd %xmm3, %xmm1
addsd %xmm0, %xmm1
movsd %xmm1, (%r12,%rax,8)
incq %rax
cmpq %rax, %rbp
jne 0x14bf
cmpl $0x0, 0x8(%rsp)
movsd 0xcfc7(%rip), %xmm2     # 0xe4b8
movsd 0xcfc7(%rip), %xmm3     # 0xe4c0
jle 0x152c
xorl %eax, %eax
movsd (%rbx,%rax,8), %xmm0
movapd %xmm0, %xmm1
mulsd %xmm2, %xmm1
mulsd %xmm0, %xmm1
movsd (%r13,%rax,8), %xmm0
mulsd %xmm3, %xmm0
addsd %xmm1, %xmm0
movsd %xmm0, (%r13,%rax,8)
incq %rax
cmpq %rax, %rbp
jne 0x14fd
movl 0xc(%rsp), %ebx
movl %ebx, %eax
xorl %edx, %edx
pushq $0x32
popq %rcx
divl %ecx
cvtsi2sd %eax, %xmm0
xorpd 0xcb6b(%rip), %xmm0     # 0xe0b0
callq 0x1040
mulsd 0xcf4e(%rip), %xmm0     # 0xe4a0
movsd %xmm0, 0x10(%rsp)
cvtsi2sd %ebx, %xmm1
movsd 0xcf5c(%rip), %xmm0     # 0xe4c0
movsd %xmm1, 0x18(%rsp)
callq 0x10a0
movapd %xmm0, %xmm1
movsd 0xcaed(%rip), %xmm0     # 0xe068
subsd %xmm1, %xmm0
ucomisd 0xcae9(%rip), %xmm0     # 0xe070
jb 0x158f
sqrtsd %xmm0, %xmm0
jmp 0x1594
callq 0x1210
movsd %xmm0, 0x20(%rsp)
movsd 0xcf0e(%rip), %xmm0     # 0xe4b0
movsd 0x18(%rsp), %xmm1
callq 0x10a0
cmpl $0x0, 0x8(%rsp)
movsd 0xceee(%rip), %xmm2     # 0xe4a8
jle 0x1644
movsd 0x10(%rsp), %xmm3
mulsd 0x20(%rsp), %xmm3
movsd 0xca94(%rip), %xmm1     # 0xe068
subsd %xmm0, %xmm1
divsd %xmm1, %xmm3
movsd %xmm3, 0x10(%rsp)
xorl %ebx, %ebx
movsd (%r12,%rbx,8), %xmm1
mulsd 0x10(%rsp), %xmm1
movsd (%r13,%rbx,8), %xmm0
ucomisd 0xca71(%rip), %xmm0     # 0xe070
jb 0x1607
sqrtsd %xmm0, %xmm0
jmp 0x1620
movsd %xmm1, 0x18(%rsp)
callq 0x1210
movsd 0x18(%rsp), %xmm1
movsd 0xce88(%rip), %xmm2     # 0xe4a8
addsd 0xcea0(%rip), %xmm0     # 0xe4c8
divsd %xmm0, %xmm1
movsd (%r14,%rbx,8), %xmm0
subsd %xmm1, %xmm0
movsd %xmm0, (%r14,%rbx,8)
incq %rbx
cmpq %rbx, %rbp
jne 0x15e4
cmpl $0x0, 0x8(%rsp)
movsd 0xce5f(%rip), %xmm3     # 0xe4b0
jle 0x167b
xorl %eax, %eax
movsd (%r14,%rax,8), %xmm0
mulsd %xmm2, %xmm0
movsd (%r15,%rax,8), %xmm1
mulsd %xmm3, %xmm1
addsd %xmm0, %xmm1
movsd %xmm1, (%r15,%rax,8)
incq %rax
cmpq %rax, %rbp
jne 0x1655
movsd 0xce4d(%rip), %xmm2     # 0xe4d0
movapd %xmm2, %xmm4
cmpl $0x0, 0x8(%rsp)
movsd 0x48(%rsp), %xmm1
movapd 0xca06(%rip), %xmm3     # 0xe0a0
jle 0x16bc
xorl %eax, %eax
movapd %xmm2, %xmm4
movsd (%r14,%rax,8), %xmm0
andpd %xmm3, %xmm0
divsd %xmm1, %xmm0
addsd %xmm0, %xmm4
incq %rax
cmpq %rax, %rbp
jne 0x16a2
movapd %xmm2, %xmm5
cmpl $0x0, 0x8(%rsp)
movq 0x30(%rsp), %rbx
jle 0x16ec
xorl %eax, %eax
movapd %xmm2, %xmm5
movsd (%r12,%rax,8), %xmm0
andpd %xmm3, %xmm0
divsd %xmm1, %xmm0
addsd %xmm0, %xmm5
incq %rax
cmpq %rax, %rbp
jne 0x16d2
movsd %xmm5, 0x10(%rsp)
movsd %xmm4, 0x18(%rsp)
cmpl $0x1, 0xc(%rsp)
jne 0x171b
movq 0x108f2(%rip), %rax     # 0x11ff8
movq (%rax), %rcx
leaq 0xcf03(%rip), %rdi      # 0xe613
pushq $0x1b
popq %rsi
pushq $0x1
popq %rdx
callq 0x1200
movq 0x108d6(%rip), %rax     # 0x11ff8
movq (%rax), %rax
movq %rax, 0x20(%rsp)
leaq 0x68(%rsp), %rdi
callq 0xc0f9
movq 0x20(%rsp), %rdi
leaq 0xceef(%rip), %rsi      # 0xe62f
movl 0xc(%rsp), %edx
movsd 0x60(%rsp), %xmm1
movsd 0x18(%rsp), %xmm2
movsd 0x10(%rsp), %xmm3
movapd %xmm3, %xmm4
movb $0x5, %al
callq 0x1170
movl 0xc(%rsp), %eax
cmpl 0x2c(%rsp), %eax
jge 0x177f
incl %eax
movsd 0x10(%rsp), %xmm0
ucomisd 0x40(%rsp), %xmm0
ja 0x1473
cmpl $0x0, 0x8(%rsp)
jle 0x1798
shlq $0x3, %rbp
movq %r14, %rdi
movq %r15, %rsi
movq %rbp, %rdx
callq 0x1190
movq %r15, %rdi
callq 0x1050
movq %r12, %rdi
callq 0x1050
movq %r13, %rdi
callq 0x1050
movq %rbx, %rdi
addq $0x78, %rsp
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
jmp 0x1050