// File: /workspace/github_c_makefile_star2_O2/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 $0x68, %rsp
movsd %xmm0, 0x30(%rsp)
movl %r8d, 0x1c(%rsp)
movl %ecx, %ebp
movq %rdx, %r14
movq %rsi, 0x40(%rsp)
movq %rdi, 0x38(%rsp)
movslq %ecx, %rax
movq %rax, 0x8(%rsp)
leaq (,%rax,8), %r15
movq %r15, %rdi
callq 0x11a0
movq %rax, %rbx
leaq 0x58(%rsp), %rdi
xorl %esi, %esi
callq 0x10f0
movq %r15, %rdi
callq 0x11a0
movq %rax, 0x20(%rsp)
movq %r15, %rdi
callq 0x11a0
movq %rax, %r12
movq %r15, %rdi
callq 0x11a0
movq %rax, %r13
xorl %eax, %eax
cmpl $0x0, 0x8(%rsp)
movl $0x0, %r15d
movl %ebp, 0x18(%rsp)
cmovgl %ebp, %r15d
movq 0x20(%rsp), %rcx
cmpq %rax, %r15
je 0x8cd6
andq $0x0, (%rcx,%rax,8)
incq %rax
jmp 0x8cc7
xorl %eax, %eax
cmpq %rax, %r15
je 0x8ce7
andq $0x0, (%r12,%rax,8)
incq %rax
jmp 0x8cd8
xorl %eax, %eax
cmpq %rax, %r15
je 0x8cf9
andq $0x0, (%r13,%rax,8)
incq %rax
jmp 0x8ce9
cvtsi2sdl 0x18(%rsp), %xmm0
pushq $0x1
popq %rax
movsd %xmm0, 0x28(%rsp)
movl %eax, 0x4(%rsp)
xorl %eax, %eax
cmpq %rax, %r15
je 0x8d1d
andq $0x0, (%rbx,%rax,8)
incq %rax
jmp 0x8d0e
movq 0x40(%rsp), %rdi
movq %r14, %rsi
movq %rbx, %rdx
movl 0x18(%rsp), %ecx
callq *0x38(%rsp)
movsd 0x36e8(%rip), %xmm4     # 0xc420
movsd 0x36e8(%rip), %xmm3     # 0xc428
movsd 0x36f0(%rip), %xmm2     # 0xc438
movsd %xmm0, 0x50(%rsp)
xorl %eax, %eax
cmpq %rax, %r15
je 0x8d77
movsd (%rbx,%rax,8), %xmm0
mulsd %xmm4, %xmm0
movsd (%r12,%rax,8), %xmm1
mulsd %xmm3, %xmm1
addsd %xmm0, %xmm1
movsd %xmm1, (%r12,%rax,8)
incq %rax
jmp 0x8d50
xorl %eax, %eax
movsd 0x36af(%rip), %xmm3     # 0xc430
cmpq %rax, %r15
je 0x8db2
movsd (%rbx,%rax,8), %xmm0
movapd %xmm0, %xmm1
mulsd %xmm3, %xmm1
mulsd %xmm0, %xmm1
movsd (%r13,%rax,8), %xmm0
mulsd %xmm2, %xmm0
addsd %xmm1, %xmm0
movsd %xmm0, (%r13,%rax,8)
incq %rax
jmp 0x8d81
movl 0x4(%rsp), %ebp
movl %ebp, %eax
xorl %edx, %edx
pushq $0x32
popq %rcx
divl %ecx
cvtsi2sd %eax, %xmm0
xorpd 0x32d5(%rip), %xmm0     # 0xc0a0
callq 0x1040
mulsd 0x3640(%rip), %xmm0     # 0xc418
movsd %xmm0, 0x10(%rsp)
cvtsi2sd %ebp, %xmm1
movsd 0x364e(%rip), %xmm0     # 0xc438
movsd %xmm1, 0x8(%rsp)
callq 0x10a0
movapd %xmm0, %xmm1
movsd 0x3267(%rip), %xmm0     # 0xc068
subsd %xmm1, %xmm0
ucomisd 0x3263(%rip), %xmm0     # 0xc070
jb 0x8e15
sqrtsd %xmm0, %xmm0
jmp 0x8e1a
callq 0x1210
movsd 0x3606(%rip), %xmm1     # 0xc428
movsd 0x10(%rsp), %xmm2
mulsd %xmm0, %xmm2
movsd %xmm2, 0x10(%rsp)
movapd %xmm1, %xmm0
movsd 0x8(%rsp), %xmm1
callq 0x10a0
movsd 0x321f(%rip), %xmm1     # 0xc068
subsd %xmm0, %xmm1
movsd 0x10(%rsp), %xmm0
divsd %xmm1, %xmm0
movapd %xmm0, %xmm1
xorl %ebp, %ebp
movsd 0x28(%rsp), %xmm2
movsd 0x35b5(%rip), %xmm4     # 0xc420
movsd %xmm0, 0x10(%rsp)
cmpq %rbp, %r15
je 0x8edd
movsd (%r12,%rbp,8), %xmm3
mulsd %xmm1, %xmm3
movsd (%r13,%rbp,8), %xmm0
ucomisd 0x31e1(%rip), %xmm0     # 0xc070
jb 0x8e97
sqrtsd %xmm0, %xmm0
jmp 0x8ebc
movsd %xmm3, 0x8(%rsp)
callq 0x1210
movsd 0x8(%rsp), %xmm3
movsd 0x10(%rsp), %xmm1
movsd 0x356a(%rip), %xmm4     # 0xc420
movsd 0x28(%rsp), %xmm2
addsd 0x357c(%rip), %xmm0     # 0xc440
divsd %xmm0, %xmm3
movsd (%r14,%rbp,8), %xmm0
subsd %xmm3, %xmm0
movsd %xmm0, (%r14,%rbp,8)
incq %rbp
jmp 0x8e71
xorl %eax, %eax
movq 0x20(%rsp), %rbp
movsd 0x353c(%rip), %xmm3     # 0xc428
cmpq %rax, %r15
je 0x8f14
movsd (%r14,%rax,8), %xmm0
mulsd %xmm4, %xmm0
movsd (%rbp,%rax,8), %xmm1
mulsd %xmm3, %xmm1
addsd %xmm0, %xmm1
movsd %xmm1, (%rbp,%rax,8)
incq %rax
jmp 0x8eec
xorl %eax, %eax
movsd 0x352a(%rip), %xmm3     # 0xc448
movapd 0x318a(%rip), %xmm1     # 0xc0b0
cmpq %rax, %r15
je 0x8f42
movsd (%r14,%rax,8), %xmm0
andpd %xmm1, %xmm0
divsd %xmm2, %xmm0
addsd %xmm0, %xmm3
incq %rax
jmp 0x8f26
xorl %eax, %eax
movsd 0x34fc(%rip), %xmm4     # 0xc448
cmpq %rax, %r15
je 0x8f68
movsd (%r12,%rax,8), %xmm0
andpd %xmm1, %xmm0
divsd %xmm2, %xmm0
addsd %xmm0, %xmm4
incq %rax
jmp 0x8f4c
movsd %xmm4, 0x8(%rsp)
movsd %xmm3, 0x10(%rsp)
cmpl $0x1, 0x4(%rsp)
jne 0x8f97
movq 0x7076(%rip), %rax      # 0xfff8
movq (%rax), %rcx
leaq 0x35eb(%rip), %rdi      # 0xc577
pushq $0x1b
popq %rsi
pushq $0x1
popq %rdx
callq 0x1200
movq 0x705a(%rip), %rax      # 0xfff8
movq (%rax), %rax
movq %rax, 0x48(%rsp)
leaq 0x58(%rsp), %rdi
callq 0x9da0
movq 0x48(%rsp), %rdi
leaq 0x35d7(%rip), %rsi      # 0xc593
movl 0x4(%rsp), %edx
movsd 0x50(%rsp), %xmm1
movsd 0x10(%rsp), %xmm2
movsd 0x8(%rsp), %xmm3
movapd %xmm3, %xmm4
movb $0x5, %al
callq 0x1170
movl 0x4(%rsp), %eax
cmpl 0x1c(%rsp), %eax
jge 0x8ffb
incl %eax
movsd 0x8(%rsp), %xmm1
ucomisd 0x30(%rsp), %xmm1
ja 0x8d08
xorl %eax, %eax
cmpq %rax, %r15
je 0x9013
movsd (%rbp,%rax,8), %xmm0
movsd %xmm0, (%r14,%rax,8)
incq %rax
jmp 0x8ffd
movq %rbp, %rdi
callq 0x1050
movq %r12, %rdi
callq 0x1050
movq %r13, %rdi
callq 0x1050
movq %rbx, %rdi
addq $0x68, %rsp
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
jmp 0x1050