// File: /workspace/github_c_makefile_star2_O3/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, 0x58(%rsp)
movl %r8d, 0x24(%rsp)
movq %rdx, %r14
movq %rsi, 0x38(%rsp)
movq %rdi, 0x30(%rsp)
movl %ecx, 0x8(%rsp)
movslq %ecx, %r15
leaq (,%r15,8), %r13
movq %r13, %rdi
callq 0x11b0
movq %rax, %rbx
leaq 0x68(%rsp), %rdi
xorl %esi, %esi
callq 0x10f0
movq %r13, %rdi
callq 0x11b0
movq %rax, %rbp
movq %r13, %rdi
callq 0x11b0
movq %rax, %r12
movq %r13, %rdi
callq 0x11b0
movq %rax, %r13
movq %rbp, %rdi
movl %r15d, %ebp
leaq (,%rbp,8), %rdx
testl %r15d, %r15d
movq %rdi, 0x28(%rsp)
jle 0x1458
xorl %esi, %esi
callq 0x1110
leaq (,%rbp,8), %r15
movq %r12, %rdi
xorl %esi, %esi
movq %r15, %rdx
callq 0x1110
movq %r13, %rdi
xorl %esi, %esi
movq %r15, %rdx
callq 0x1110
movq %r15, %rdx
cvtsi2sdl 0x8(%rsp), %xmm0
movsd %xmm0, 0x50(%rsp)
pushq $0x1
popq %r15
movq %rdx, 0x60(%rsp)
movl %r15d, 0xc(%rsp)
movl 0x8(%rsp), %r15d
testl %r15d, %r15d
jle 0x1525
movq %rbx, %rdi
xorl %esi, %esi
callq 0x1110
movq 0x38(%rsp), %rdi
movq %r14, %rsi
movq %rbx, %rdx
movl %r15d, %ecx
callq *0x30(%rsp)
movsd %xmm0, 0x48(%rsp)
xorl %eax, %eax
movsd 0xbfec(%rip), %xmm2     # 0xd498
movsd 0xbfec(%rip), %xmm3     # 0xd4a0
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 0x14b4
testl %r15d, %r15d
movl 0xc(%rsp), %ecx
jle 0x1541
xorl %eax, %eax
movsd 0xbfbc(%rip), %xmm2     # 0xd4a8
movsd 0xbfbc(%rip), %xmm3     # 0xd4b0
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 0x14f4
jmp 0x1541
movq 0x38(%rsp), %rdi
movq %r14, %rsi
movq %rbx, %rdx
movl %r15d, %ecx
callq *0x30(%rsp)
movsd %xmm0, 0x48(%rsp)
movl 0xc(%rsp), %ecx
movl %ecx, %eax
xorl %edx, %edx
pushq $0x32
popq %rsi
divl %esi
cvtsi2sd %eax, %xmm0
xorpd 0xbb5a(%rip), %xmm0     # 0xd0b0
movl %ecx, %r15d
callq 0x1040
mulsd 0xbf2a(%rip), %xmm0     # 0xd490
movsd %xmm0, 0x18(%rsp)
cvtsi2sd %r15d, %xmm1
movsd 0xbf37(%rip), %xmm0     # 0xd4b0
movsd %xmm1, 0x10(%rsp)
callq 0x10a0
movapd %xmm0, %xmm1
movsd 0xbad8(%rip), %xmm0     # 0xd068
subsd %xmm1, %xmm0
ucomisd 0xbad4(%rip), %xmm0     # 0xd070
jb 0x15a4
sqrtsd %xmm0, %xmm0
jmp 0x15a9
callq 0x1220
movsd %xmm0, 0x40(%rsp)
movsd 0xbee9(%rip), %xmm0     # 0xd4a0
movl 0x8(%rsp), %r15d
movsd 0x10(%rsp), %xmm1
callq 0x10a0
movsd 0xbee9(%rip), %xmm1     # 0xd4b8
movapd %xmm1, %xmm3
movapd %xmm1, %xmm4
testl %r15d, %r15d
movl 0xc(%rsp), %edx
movsd 0xbeb2(%rip), %xmm2     # 0xd498
jle 0x1744
movsd 0x18(%rsp), %xmm3
mulsd 0x40(%rsp), %xmm3
movsd 0xba68(%rip), %xmm1     # 0xd068
subsd %xmm0, %xmm1
divsd %xmm1, %xmm3
xorl %r15d, %r15d
movsd %xmm3, 0x18(%rsp)
movsd (%r12,%r15,8), %xmm1
mulsd %xmm3, %xmm1
movsd (%r13,%r15,8), %xmm0
ucomisd 0xba46(%rip), %xmm0     # 0xd070
jb 0x1632
sqrtsd %xmm0, %xmm0
jmp 0x1655
movsd %xmm1, 0x10(%rsp)
callq 0x1220
movsd 0x10(%rsp), %xmm1
movsd 0x18(%rsp), %xmm3
movsd 0xbe47(%rip), %xmm2     # 0xd498
movl 0xc(%rsp), %edx
addsd 0xbe63(%rip), %xmm0     # 0xd4c0
divsd %xmm0, %xmm1
movsd (%r14,%r15,8), %xmm0
subsd %xmm1, %xmm0
movsd %xmm0, (%r14,%r15,8)
incq %r15
cmpq %r15, %rbp
jne 0x1611
movsd 0xbe37(%rip), %xmm0     # 0xd4b8
movapd %xmm0, %xmm3
movapd %xmm0, %xmm4
movl 0x8(%rsp), %r15d
testl %r15d, %r15d
jle 0x1744
xorl %eax, %eax
movq 0x28(%rsp), %rcx
movsd 0xbdfa(%rip), %xmm3     # 0xd4a0
movsd (%r14,%rax,8), %xmm0
mulsd %xmm2, %xmm0
movsd (%rcx,%rax,8), %xmm1
mulsd %xmm3, %xmm1
addsd %xmm0, %xmm1
movsd %xmm1, (%rcx,%rax,8)
incq %rax
cmpq %rax, %rbp
jne 0x16a6
movsd 0xbde6(%rip), %xmm0     # 0xd4b8
movapd %xmm0, %xmm3
movapd %xmm0, %xmm4
testl %r15d, %r15d
jle 0x1744
xorl %eax, %eax
movsd 0xbdcf(%rip), %xmm3     # 0xd4b8
movsd 0x50(%rsp), %xmm1
movapd 0xb9a9(%rip), %xmm2     # 0xd0a0
movsd (%r14,%rax,8), %xmm0
andpd %xmm2, %xmm0
divsd %xmm1, %xmm0
addsd %xmm0, %xmm3
incq %rax
cmpq %rax, %rbp
jne 0x16f7
testl %r15d, %r15d
jle 0x173c
xorl %eax, %eax
movsd 0xbd98(%rip), %xmm4     # 0xd4b8
movsd (%r12,%rax,8), %xmm0
andpd %xmm2, %xmm0
divsd %xmm1, %xmm0
addsd %xmm0, %xmm4
incq %rax
cmpq %rax, %rbp
jne 0x1720
jmp 0x1744
movsd 0xbd74(%rip), %xmm4     # 0xd4b8
movsd %xmm4, 0x10(%rsp)
movsd %xmm3, 0x18(%rsp)
cmpl $0x1, %edx
jne 0x1771
movq 0xf89c(%rip), %rax      # 0x10ff8
movq (%rax), %rcx
leaq 0xbe9d(%rip), %rdi      # 0xd603
pushq $0x1b
popq %rsi
pushq $0x1
popq %rdx
callq 0x1210
movq 0xf880(%rip), %rax      # 0x10ff8
movq (%rax), %r15
leaq 0x68(%rsp), %rdi
callq 0xb0e7
movq %r15, %rdi
leaq 0xbe90(%rip), %rsi      # 0xd61f
movl 0xc(%rsp), %r15d
movl %r15d, %edx
movsd 0x48(%rsp), %xmm1
movsd 0x18(%rsp), %xmm2
movsd 0x10(%rsp), %xmm3
movapd %xmm3, %xmm4
movb $0x5, %al
callq 0x1180
cmpl 0x24(%rsp), %r15d
movq 0x60(%rsp), %rdx
jge 0x17d5
incl %r15d
movsd 0x10(%rsp), %xmm0
ucomisd 0x58(%rsp), %xmm0
ja 0x146d
cmpl $0x0, 0x8(%rsp)
movq 0x28(%rsp), %r15
jle 0x17ec
movq %r14, %rdi
movq %r15, %rsi
callq 0x11a0
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