// File: /workspace/github_cpp_makefile_star2_O2/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 $0xa8, %rsp
movsd %xmm0, 0x70(%rsp)
movl %r9d, 0x4c(%rsp)
movq %r8, %r14
movq %rcx, %r15
movq %rdx, %r12
movq %rsi, 0x40(%rsp)
movl %edi, 0x1c(%rsp)
leaq 0x34(%rsp), %rsi
movq %r8, %rdi
callq 0x67d0
leaq 0x7c(%rsp), %rbx
movq %r14, 0x10(%rsp)
movq %r14, %rdi
movq %rbx, %rsi
callq 0x6600
cvtsi2sdl (%rbx), %xmm0
xorpd %xmm1, %xmm1
ucomisd %xmm1, %xmm0
jb 0x1de41
sqrtsd %xmm0, %xmm0
jmp 0x1de46
callq 0x63b0
movsd %xmm0, 0x50(%rsp)
pushq $0x2
popq %rsi
leaq 0x88(%rsp), %rdx
leaq 0x80(%rsp), %rcx
leaq 0x20(%rsp), %r8
movq 0x10(%rsp), %rbx
movq %rbx, %rdi
callq 0x6430
movl $0x1, 0xc(%rsp)
xorl %r13d, %r13d
movl 0x1c(%rsp), %ebp
movl %ebp, %edi
movq %rbx, %rsi
xorl %edx, %edx
callq 0x1e2fc
movl %eax, %r14d
pushq $0x1
popq %rdx
movl %ebp, %edi
movq %rbx, %rsi
callq 0x1e2fc
movl %eax, %ebp
movl %r14d, 0x30(%rsp)
movslq %r14d, %r14
leaq (,%r14,8), %rdi
movq %rdi, 0x28(%rsp)
callq 0x6750
movq %rax, %rbx
movl %ebp, %eax
imull %r14d, %eax
movslq %eax, %rdi
shlq $0x3, %rdi
callq 0x6750
testl %ebp, %ebp
movl $0x0, %ecx
cmovgl %ebp, %ecx
movq %rax, %rdx
testl %r14d, %r14d
cmovlel %r13d, %r14d
movslq %ebp, %rax
shlq $0x3, %rax
movq %rdx, 0x38(%rsp)
movq 0x40(%rsp), %rsi
cmpq %r14, %r13
je 0x1df18
xorl %edi, %edi
cmpq %rdi, %rcx
je 0x1df0d
movsd (%rsi,%rdi,8), %xmm0
movsd %xmm0, (%rdx,%rdi,8)
incq %rdi
jmp 0x1def9
incq %r13
addq %rax, %rsi
addq %rax, %rdx
jmp 0x1def2
movl 0x20(%rsp), %ecx
movl 0x24(%rsp), %eax
cmpl %eax, %ecx
pushq $0x1
popq %r13
jne 0x1df56
incl %ebp
xorl %edx, %edx
xorl %esi, %esi
cmpq %rsi, %r14
je 0x1df56
movslq %edx, %rdx
movq 0x40(%rsp), %rdi
movsd (%rdi,%rdx,8), %xmm0
movq 0x38(%rsp), %rdi
andq $0x0, (%rdi,%rdx,8)
movsd %xmm0, (%rbx,%rsi,8)
incq %rsi
addl %ebp, %edx
jmp 0x1df2e
cvttsd2si 0x50(%rsp), %ebp
xorl %edx, %edx
cmpq %rdx, %r14
je 0x1df6d
andq $0x0, (%r15,%rdx,8)
incq %rdx
jmp 0x1df5e
cmpl %eax, %ecx
jne 0x1df93
imull %ebp, %ecx
movq 0x1302d(%rip), %rdx     # 0x30fa8
movq %rbx, %rdi
movl 0x30(%rsp), %esi
xorl %r8d, %r8d
movq 0x10(%rsp), %r9
callq 0x61b0
movl 0x24(%rsp), %eax
movl %r13d, %esi
testl %eax, %eax
jne 0x1dfd5
incl %ebp
imull 0x20(%rsp), %ebp
leaq 0x90(%rsp), %rax
movq %rax, (%rsp)
movq 0x12ff4(%rip), %rdx     # 0x30fa8
movq %rbx, %rdi
movl 0x30(%rsp), %esi
movl %ebp, %ecx
xorl %r8d, %r8d
movq 0x10(%rsp), %r9
callq 0x63e0
xorl %esi, %esi
cmpl $0x0, 0x24(%rsp)
setne %sil
movl 0x34(%rsp), %edx
leaq 0x68(%rsp), %rcx
movq 0x10(%rsp), %rdi
callq 0x6920
movq 0x28(%rsp), %rbp
movq %rbp, %rdi
callq 0x6750
movq %rax, %r13
movq %rbp, %rdi
callq 0x6750
movq %rax, %rbp
movq $0x0, 0x28(%rsp)
pushq $0x1
popq %rax
cmpl $0x1, %eax
jne 0x1e169
movl 0x1c(%rsp), %edi
movq 0x40(%rsp), %rsi
movq %r15, %rdx
movq %r13, %rcx
movq 0x10(%rsp), %r8
callq 0x1dc22
cmpl $0x0, 0x24(%rsp)
je 0x1e081
movq 0x68(%rsp), %rax
movq %rax, (%rsp)
leaq 0x60(%rsp), %rdi
leaq 0x58(%rsp), %rsi
pushq $0x1
popq %rdx
movq 0x12f53(%rip), %rcx     # 0x30fa8
movq 0x12f1c(%rip), %r8      # 0x30f78
xorl %r9d, %r9d
callq 0x6740
cmpl $0x0, 0x34(%rsp)
jne 0x1e0e9
movsd 0x58(%rsp), %xmm0
ucomisd 0xfaf(%rip), %xmm0      # 0x1f028
jb 0x1e0b9
sqrtsd %xmm0, %xmm0
jmp 0x1e0be
andq $0x0, 0x60(%rsp)
xorpd %xmm0, %xmm0
xorl %eax, %eax
cmpq %rax, %r14
je 0x1e038
movsd (%r13,%rax,8), %xmm1
subsd (%r12,%rax,8), %xmm1
movsd %xmm1, (%r13,%rax,8)
mulsd %xmm1, %xmm1
addsd %xmm1, %xmm0
movsd %xmm0, 0x60(%rsp)
incq %rax
jmp 0x1e08d
callq 0x63b0
movsd %xmm0, 0x58(%rsp)
movsd 0x70(%rsp), %xmm1
ucomisd %xmm0, %xmm1
seta %al
movq 0x28(%rsp), %rcx
cmpl 0x4c(%rsp), %ecx
setg %cl
orb %al, %cl
cmpb $0x1, %cl
jne 0x1e0e9
andl $0x0, 0xc(%rsp)
leaq 0xc(%rsp), %rdi
pushq $0x1
popq %rsi
movq 0x12ec8(%rip), %rdx     # 0x30fc0
xorl %ecx, %ecx
movq 0x10(%rsp), %r8
callq 0x60d0
cmpl $0x1, 0xc(%rsp)
jne 0x1e12b
movl 0x1c(%rsp), %edi
movq 0x38(%rsp), %rsi
movq %r15, %rdx
movq %rbp, %rcx
movq 0x10(%rsp), %r8
callq 0x1dc22
cmpl $0x0, 0x24(%rsp)
je 0x1e140
movq 0x28(%rsp), %rax
incl %eax
movq %rax, 0x28(%rsp)
movl 0xc(%rsp), %eax
jmp 0x1e00f
xorl %eax, %eax
cmpq %rax, %r14
je 0x1e12b
movsd (%r12,%rax,8), %xmm0
subsd (%rbp,%rax,8), %xmm0
movsd %xmm0, (%rbp,%rax,8)
divsd (%rbx,%rax,8), %xmm0
movsd %xmm0, (%r15,%rax,8)
incq %rax
jmp 0x1e142
leaq 0x68(%rsp), %rdi
callq 0x6900
movq %r13, %rdi
callq 0x6650
movq %rbp, %rdi
callq 0x6650
movq %rbx, %rdi
callq 0x6650
movq 0x38(%rsp), %rdi
callq 0x6650
addq $0xa8, %rsp
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
retq