// File: /workspace/github_cpp_makefile_star2_O1/WeihuaZhu[P]hpc_prog2/mpi_jacobi.cpp
// Function: distributed_jacobi(int, double*, double*, double*, ompi_communicator_t*, int, double)
void distributed_jacobi(const int n, double* local_A, double* local_b, double* local_x,
                MPI_Comm comm, int max_iter, double l2_termination)
{
    // obtain the necessary parameters required in thie method
    int rank;
    MPI_Comm_rank(comm, &rank);
    int p;
    MPI_Comm_size(comm, &p);
    int q = sqrt(p);
    int dims[2];
    int coords[2];
    int periods[2];
    MPI_Cart_get(comm, 2, dims, periods, coords);

    MPI_Status stat;
    int num_iter = 0; //num_iter counts the numebr of iterations in jacobi iteration
    int sign = 1; //sign defines whether to continue iteration in jacobi iteration
    MPI_Comm n_comm;
    double Ax_b_square; //stores the intermediate value of (Ax - b) ^ 2
    double norm; //norm stores the l2-norm value
    //initialization of the dimention size in each processor:
    int count_dim_0 = block_decompose_by_dim(n, comm, 0);
    int count_dim_1 = block_decompose_by_dim(n, comm, 1);

    //initialization of D, R
    double *local_D = (double*) malloc(sizeof(double) * count_dim_0);
    double *local_R = (double*) malloc(sizeof(double) * (count_dim_0 * count_dim_1));
    for (int i = 0; i < count_dim_0; i++) {
        for (int j = 0; j < count_dim_1; j++) {
            local_R[count_dim_1 * i + j] = local_A[count_dim_1 * i + j];
        }
    }
    if (coords[0] == coords[1]) {
        for (int i = 0; i < count_dim_0; i++) {
            local_D[i] = local_A[count_dim_1 * i + i];
            local_R[count_dim_1 * i + i] = 0;
        }
    }

    //initialization of x
    for (int i = 0; i < count_dim_0; i++) {
        local_x[i] = 0;
    }

    // diagonal processors send D to first column, A and R are in each processor, D is in the
    // first column
    if (coords[0] == coords[1]) {
        MPI_Send(local_D, count_dim_0, MPI_DOUBLE, coords[0] * q, 0, comm);
    }
    if (coords[1] == 0) { //coords[0] != 0 is not required, don't nned to avoid dead lock
        MPI_Recv(local_D, count_dim_0, MPI_DOUBLE, (coords[0] * (q + 1)), 0, comm, &stat);
    }

    //split the first colum as a communicator, and all other processors as another comminicator
    int color = 1;
    if (coords[1] == 0) {
        color = 0;
    }
    MPI_Comm_split(comm, color, rank, &n_comm);

    // local_Ax and local_Rx stores the Ax and Rx values
    double *local_Ax = (double*) malloc(sizeof(double) * count_dim_0);
    double *local_Rx = (double*) malloc(sizeof(double) * count_dim_0);

    //jacobi iteration
    while(sign == 1) {
        //calculate Ax and stores the solution in the first column
        distributed_matrix_vector_mult(n, local_A, local_x, local_Ax, comm);
        //compute the 2nd norm of Ax - b in the first column
        if (coords[1] == 0) {
            Ax_b_square = 0;
            for (int i = 0; i < count_dim_0; i++) {
                local_Ax[i] -= local_b[i];
                Ax_b_square += local_Ax[i] * local_Ax[i];
            }
        }
        // calculate the square of the second norm using reduce on first column processors,
        // and store the solution in the first processor
        MPI_Reduce(&Ax_b_square, &norm, 1, MPI_DOUBLE, MPI_SUM, 0, n_comm);

        //calculate norm and update sign & num_iter in the first processor
        if (rank == 0) {
            norm = sqrt(norm);
            if (norm < l2_termination || num_iter > max_iter) {
                sign = 0;
            }
        }
        num_iter++;

        //Broadcast sign to all the processors to avoid infinite loop
        MPI_Bcast(&sign, 1, MPI_INT, 0, comm);

        //update x in the first column
        if (sign == 1) {
            //calculate Rx, and store it in the first column
            distributed_matrix_vector_mult(n, local_R, local_x, local_Rx, comm);
            //calculate x = (D^-1)(b - Rx) in the first column
            if (coords[1] == 0) {
                for (int i = 0; i < count_dim_0; i++) {
                    local_Rx[i] = local_b[i] - local_Rx[i];
                    local_x[i] = local_Rx[i] / local_D[i];
                }
            }
        }
    }
    // free the useless memory
    MPI_Comm_free(&n_comm);
    free(local_Ax);
    free(local_Rx);
    free(local_D);
    free(local_R);
}

// Assembly:
pushq %rbp
pushq %r15
pushq %r14
pushq %r13
pushq %r12
pushq %rbx
subq $0xc8, %rsp
movsd %xmm0, 0x90(%rsp)
movl %r9d, %ebp
movq %r8, %r14
movq %rcx, 0x30(%rsp)
movq %rdx, %r12
movq %rsi, 0x48(%rsp)
movl %edi, 0x1c(%rsp)
leaq 0x3c(%rsp), %rsi
movq %r8, %rdi
callq 0x6790
leaq 0x9c(%rsp), %rbx
movq %r14, 0x10(%rsp)
movq %r14, %rdi
movq %rbx, %rsi
callq 0x65e0
xorps %xmm0, %xmm0
cvtsi2sdl (%rbx), %xmm0
xorpd %xmm1, %xmm1
ucomisd %xmm1, %xmm0
movl %ebp, 0x64(%rsp)
jb 0x2820a
sqrtsd %xmm0, %xmm0
jmp 0x2820f
callq 0x63a0
movsd %xmm0, 0x68(%rsp)
leaq 0xc0(%rsp), %rdx
leaq 0xb8(%rsp), %rcx
leaq 0x20(%rsp), %r8
movq 0x10(%rsp), %r15
movq %r15, %rdi
movl $0x2, %esi
callq 0x6420
movl $0x1, 0xc(%rsp)
leaq 0xa0(%rsp), %rbx
leaq 0x58(%rsp), %rcx
leaq 0x50(%rsp), %r14
movq %r15, %rdi
movl $0x2, %esi
movq %rbx, %rdx
movq %r14, %r8
callq 0x6420
movl 0x1c(%rsp), %ebp
movl %ebp, %eax
cltd
idivl (%rbx)
xorl %ecx, %ecx
cmpl (%r14), %edx
setg %cl
movq %rax, 0x80(%rsp)
movq %rcx, 0x78(%rsp)
leal (%rax,%rcx), %ebx
leaq 0xa0(%rsp), %r14
leaq 0x58(%rsp), %rcx
leaq 0x50(%rsp), %r13
movq %r15, %rdi
movl $0x2, %esi
movq %r14, %rdx
movq %r13, %r8
callq 0x6420
movl %ebp, %eax
cltd
idivl 0x4(%r14)
movl %eax, %r14d
xorl %r15d, %r15d
cmpl 0x4(%r13), %edx
setg %r15b
leal (%r14,%r15), %r13d
movslq %ebx, %rbp
leaq (,%rbp,8), %rdi
movq %rdi, 0x88(%rsp)
callq 0x6720
movq %rax, 0x28(%rsp)
movl %r13d, %eax
imull %ebx, %eax
movslq %eax, %rdi
shlq $0x3, %rdi
callq 0x6720
movl %ebx, 0x8(%rsp)
movl %ebx, %ebx
testl %ebp, %ebp
movq 0x48(%rsp), %rdx
movq %rax, 0x40(%rsp)
jle 0x28353
leal (%r14,%r15), %ecx
decl %ecx
leaq 0x8(,%rcx,8), %r14
xorl %r15d, %r15d
movq %rbx, %rbp
testl %r13d, %r13d
jle 0x2834b
movl %r15d, %eax
movq 0x48(%rsp), %rcx
leaq (%rcx,%rax,8), %rsi
movq 0x40(%rsp), %rcx
leaq (%rcx,%rax,8), %rdi
movq %r14, %rdx
callq 0x63e0
movq 0x40(%rsp), %rax
movq 0x48(%rsp), %rdx
addl %r13d, %r15d
decq %rbp
jne 0x2831f
movl 0x20(%rsp), %r14d
movl 0x24(%rsp), %r15d
cmpl %r15d, %r14d
movq 0x28(%rsp), %rsi
jne 0x28395
cmpl $0x0, 0x8(%rsp)
jle 0x28395
incl %r13d
xorl %edi, %edi
xorl %ecx, %ecx
movslq %edi, %rdi
movsd (%rdx,%rdi,8), %xmm0
movsd %xmm0, (%rsi,%rcx,8)
movq $0x0, (%rax,%rdi,8)
incq %rcx
addl %r13d, %edi
cmpq %rcx, %rbx
jne 0x28375
cmpl $0x0, 0x8(%rsp)
movl 0x64(%rsp), %ebp
jle 0x283c5
movq 0x80(%rsp), %rax
movq 0x78(%rsp), %rcx
addl %ecx, %eax
decl %eax
leaq 0x8(,%rax,8), %rdx
movq 0x30(%rsp), %rdi
xorl %esi, %esi
callq 0x6280
cvttsd2si 0x68(%rsp), %r13d
cmpl %r15d, %r14d
jne 0x283f5
imull %r13d, %r14d
movq 0x10bc4(%rip), %rdx     # 0x38fa0
movq 0x28(%rsp), %rdi
movl 0x8(%rsp), %esi
movl %r14d, %ecx
xorl %r8d, %r8d
movq 0x10(%rsp), %r9
callq 0x61c0
cmpl $0x0, 0x24(%rsp)
jne 0x28431
incl %r13d
imull 0x20(%rsp), %r13d
leaq 0xa0(%rsp), %rax
movq %rax, (%rsp)
movq 0x10b88(%rip), %rdx     # 0x38fa0
movq 0x28(%rsp), %rdi
movl 0x8(%rsp), %esi
movl %r13d, %ecx
xorl %r8d, %r8d
movq 0x10(%rsp), %r9
callq 0x63d0
xorl %esi, %esi
cmpl $0x0, 0x24(%rsp)
setne %sil
movl 0x3c(%rsp), %edx
leaq 0x58(%rsp), %rcx
movq 0x10(%rsp), %rdi
callq 0x68d0
movq 0x88(%rsp), %r14
movq %r14, %rdi
callq 0x6720
movq %rax, %r13
movq %r14, %rdi
callq 0x6720
movq %rax, %r14
cmpl $0x1, 0xc(%rsp)
movq 0x30(%rsp), %rdx
jne 0x285e8
xorl %r15d, %r15d
movl 0x1c(%rsp), %edi
movq 0x48(%rsp), %rsi
movq %r13, %rcx
movq 0x10(%rsp), %r8
callq 0x27fcf
cmpl $0x0, 0x24(%rsp)
jne 0x284dd
movq $0x0, 0x50(%rsp)
cmpl $0x0, 0x8(%rsp)
jle 0x284dd
xorpd %xmm0, %xmm0
xorl %eax, %eax
movsd (%r13,%rax,8), %xmm1
subsd (%r12,%rax,8), %xmm1
movsd %xmm1, (%r13,%rax,8)
mulsd %xmm1, %xmm1
addsd %xmm1, %xmm0
incq %rax
cmpq %rax, %rbx
jne 0x284b3
movsd %xmm0, 0x50(%rsp)
movq 0x58(%rsp), %rax
movq %rax, (%rsp)
leaq 0x50(%rsp), %rdi
leaq 0x70(%rsp), %rsi
movl $0x1, %edx
movq 0x10aa4(%rip), %rcx     # 0x38fa0
movq 0x10a65(%rip), %r8      # 0x38f68
xorl %r9d, %r9d
callq 0x6710
cmpl $0x0, 0x3c(%rsp)
jne 0x28558
movsd 0x70(%rsp), %xmm0
ucomisd 0xb08(%rip), %xmm0      # 0x29028
jb 0x28528
sqrtsd %xmm0, %xmm0
jmp 0x2852d
callq 0x63a0
movsd %xmm0, 0x70(%rsp)
movsd 0x90(%rsp), %xmm1
ucomisd %xmm0, %xmm1
seta %al
cmpl %ebp, %r15d
setg %cl
orb %al, %cl
cmpb $0x1, %cl
jne 0x28558
movl $0x0, 0xc(%rsp)
leaq 0xc(%rsp), %rdi
movl $0x1, %esi
movq 0x10a4f(%rip), %rdx     # 0x38fb8
xorl %ecx, %ecx
movq 0x10(%rsp), %r8
callq 0x60d0
cmpl $0x1, 0xc(%rsp)
movq 0x30(%rsp), %rdx
jne 0x285da
movl 0x1c(%rsp), %edi
movq 0x40(%rsp), %rsi
movq 0x30(%rsp), %rdx
movq %r14, %rcx
movq 0x10(%rsp), %r8
callq 0x27fcf
movq 0x28(%rsp), %rcx
movq 0x30(%rsp), %rdx
cmpl $0x0, 0x24(%rsp)
jne 0x285da
cmpl $0x0, 0x8(%rsp)
jle 0x285da
xorl %eax, %eax
movsd (%r12,%rax,8), %xmm0
subsd (%r14,%rax,8), %xmm0
movsd %xmm0, (%r14,%rax,8)
divsd (%rcx,%rax,8), %xmm0
movsd %xmm0, (%rdx,%rax,8)
incq %rax
cmpq %rax, %rbx
jne 0x285b6
incl %r15d
cmpl $0x1, 0xc(%rsp)
je 0x28480
leaq 0x58(%rsp), %rdi
callq 0x68b0
movq %r13, %rdi
callq 0x6630
movq %r14, %rdi
callq 0x6630
movq 0x28(%rsp), %rdi
callq 0x6630
movq 0x40(%rsp), %rdi
callq 0x6630
addq $0xc8, %rsp
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
retq