// File: 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:
subq $0xe8, %rsp
movl %edi, 0xe4(%rsp)
movq %rsi, 0xd8(%rsp)
movq %rdx, 0xd0(%rsp)
movq %rcx, 0xc8(%rsp)
movq %r8, 0xc0(%rsp)
movl %r9d, 0xbc(%rsp)
movsd %xmm0, 0xb0(%rsp)
movq 0xc0(%rsp), %rdi
leaq 0xac(%rsp), %rsi
callq 0xf050
movq 0xc0(%rsp), %rdi
leaq 0xa8(%rsp), %rsi
callq 0xece0
movl 0xa8(%rsp), %edi
callq 0x119d0
cvttsd2si %xmm0, %eax
movl %eax, 0xa4(%rsp)
movq 0xc0(%rsp), %rdi
leaq 0x9c(%rsp), %rdx
leaq 0x8c(%rsp), %rcx
leaq 0x94(%rsp), %r8
movl $0x2, %esi
callq 0xe8f0
movl $0x0, 0x6c(%rsp)
movl $0x1, 0x68(%rsp)
movl 0xe4(%rsp), %edi
movq 0xc0(%rsp), %rsi
xorl %edx, %edx
callq 0x4e120
movl %eax, 0x4c(%rsp)
movl 0xe4(%rsp), %edi
movq 0xc0(%rsp), %rsi
movl $0x1, %edx
callq 0x4e120
movl %eax, 0x48(%rsp)
movslq 0x4c(%rsp), %rdi
shlq $0x3, %rdi
callq 0xef20
movq %rax, 0x40(%rsp)
movl 0x4c(%rsp), %eax
imull 0x48(%rsp), %eax
movslq %eax, %rdi
shlq $0x3, %rdi
callq 0xef20
movq %rax, 0x38(%rsp)
movl $0x0, 0x34(%rsp)
movl 0x34(%rsp), %eax
cmpl 0x4c(%rsp), %eax
jge 0x4db25
movl $0x0, 0x30(%rsp)
movl 0x30(%rsp), %eax
cmpl 0x48(%rsp), %eax
jge 0x4db16
movq 0xd8(%rsp), %rax
movl 0x48(%rsp), %ecx
imull 0x34(%rsp), %ecx
addl 0x30(%rsp), %ecx
movslq %ecx, %rcx
movsd (%rax,%rcx,8), %xmm0
movq 0x38(%rsp), %rax
movl 0x48(%rsp), %ecx
imull 0x34(%rsp), %ecx
addl 0x30(%rsp), %ecx
movslq %ecx, %rcx
movsd %xmm0, (%rax,%rcx,8)
movl 0x30(%rsp), %eax
addl $0x1, %eax
movl %eax, 0x30(%rsp)
jmp 0x4dac8
jmp 0x4db18
movl 0x34(%rsp), %eax
addl $0x1, %eax
movl %eax, 0x34(%rsp)
jmp 0x4dab6
movl 0x94(%rsp), %eax
cmpl 0x98(%rsp), %eax
jne 0x4db9f
movl $0x0, 0x2c(%rsp)
movl 0x2c(%rsp), %eax
cmpl 0x4c(%rsp), %eax
jge 0x4db9d
movq 0xd8(%rsp), %rax
movl 0x48(%rsp), %ecx
imull 0x2c(%rsp), %ecx
addl 0x2c(%rsp), %ecx
movslq %ecx, %rcx
movsd (%rax,%rcx,8), %xmm0
movq 0x40(%rsp), %rax
movslq 0x2c(%rsp), %rcx
movsd %xmm0, (%rax,%rcx,8)
movq 0x38(%rsp), %rax
movl 0x48(%rsp), %ecx
imull 0x2c(%rsp), %ecx
addl 0x2c(%rsp), %ecx
movslq %ecx, %rcx
xorps %xmm0, %xmm0
movsd %xmm0, (%rax,%rcx,8)
movl 0x2c(%rsp), %eax
addl $0x1, %eax
movl %eax, 0x2c(%rsp)
jmp 0x4db3d
jmp 0x4db9f
movl $0x0, 0x28(%rsp)
movl 0x28(%rsp), %eax
cmpl 0x4c(%rsp), %eax
jge 0x4dbd3
movq 0xc8(%rsp), %rax
movslq 0x28(%rsp), %rcx
xorps %xmm0, %xmm0
movsd %xmm0, (%rax,%rcx,8)
movl 0x28(%rsp), %eax
addl $0x1, %eax
movl %eax, 0x28(%rsp)
jmp 0x4dba7
movl 0x94(%rsp), %eax
cmpl 0x98(%rsp), %eax
jne 0x4dc12
movq 0x40(%rsp), %rdi
movl 0x4c(%rsp), %esi
movl 0x94(%rsp), %ecx
imull 0xa4(%rsp), %ecx
movq 0xc0(%rsp), %r9
movq 0x1c39e(%rip), %rdx     # 0x69fa8
xorl %r8d, %r8d
callq 0xe330
cmpl $0x0, 0x98(%rsp)
jne 0x4dc59
movq 0x40(%rsp), %rdi
movl 0x4c(%rsp), %esi
movl 0x94(%rsp), %ecx
movl 0xa4(%rsp), %eax
addl $0x1, %eax
imull %eax, %ecx
movq 0xc0(%rsp), %r9
movq 0x1c360(%rip), %rdx     # 0x69fa8
xorl %r8d, %r8d
leaq 0x70(%rsp), %rax
movq %rax, (%rsp)
callq 0xe860
movl $0x1, 0x24(%rsp)
cmpl $0x0, 0x98(%rsp)
jne 0x4dc73
movl $0x0, 0x24(%rsp)
movq 0xc0(%rsp), %rdi
movl 0x24(%rsp), %esi
movl 0xac(%rsp), %edx
leaq 0x60(%rsp), %rcx
callq 0xf260
movslq 0x4c(%rsp), %rdi
shlq $0x3, %rdi
callq 0xef20
movq %rax, 0x18(%rsp)
movslq 0x4c(%rsp), %rdi
shlq $0x3, %rdi
callq 0xef20
movq %rax, 0x10(%rsp)
cmpl $0x1, 0x68(%rsp)
jne 0x4dee5
movl 0xe4(%rsp), %edi
movq 0xd8(%rsp), %rsi
movq 0xc8(%rsp), %rdx
movq 0x18(%rsp), %rcx
movq 0xc0(%rsp), %r8
callq 0x4d780
cmpl $0x0, 0x98(%rsp)
jne 0x4dd7e
xorps %xmm0, %xmm0
movsd %xmm0, 0x58(%rsp)
movl $0x0, 0xc(%rsp)
movl 0xc(%rsp), %eax
cmpl 0x4c(%rsp), %eax
jge 0x4dd7c
movq 0xd0(%rsp), %rax
movslq 0xc(%rsp), %rcx
movsd (%rax,%rcx,8), %xmm1
movq 0x18(%rsp), %rax
movslq 0xc(%rsp), %rcx
movsd (%rax,%rcx,8), %xmm0
subsd %xmm1, %xmm0
movsd %xmm0, (%rax,%rcx,8)
movq 0x18(%rsp), %rax
movslq 0xc(%rsp), %rcx
movsd (%rax,%rcx,8), %xmm0
movq 0x18(%rsp), %rax
movslq 0xc(%rsp), %rcx
movsd (%rax,%rcx,8), %xmm2
movsd 0x58(%rsp), %xmm1
mulsd %xmm2, %xmm0
addsd %xmm1, %xmm0
movsd %xmm0, 0x58(%rsp)
movl 0xc(%rsp), %eax
addl $0x1, %eax
movl %eax, 0xc(%rsp)
jmp 0x4dd09
jmp 0x4dd7e
movq 0x60(%rsp), %rax
leaq 0x58(%rsp), %rdi
leaq 0x50(%rsp), %rsi
movl $0x1, %edx
movq 0x1c20f(%rip), %rcx     # 0x69fa8
movq 0x1c1d0(%rip), %r8      # 0x69f70
xorl %r9d, %r9d
movq %rax, (%rsp)
callq 0xef00
cmpl $0x0, 0xac(%rsp)
jne 0x4ddf3
movsd 0x50(%rsp), %xmm0
callq 0xe7c0
movsd %xmm0, 0x50(%rsp)
movsd 0x50(%rsp), %xmm1
movsd 0xb0(%rsp), %xmm0
ucomisd %xmm1, %xmm0
ja 0x4dde9
movl 0x6c(%rsp), %eax
cmpl 0xbc(%rsp), %eax
jle 0x4ddf1
movl $0x0, 0x68(%rsp)
jmp 0x4ddf3
movl 0x6c(%rsp), %eax
addl $0x1, %eax
movl %eax, 0x6c(%rsp)
movq 0xc0(%rsp), %r8
leaq 0x68(%rsp), %rdi
movl $0x1, %esi
movq 0x1c1a9(%rip), %rdx     # 0x69fc0
xorl %ecx, %ecx
callq 0xe160
cmpl $0x1, 0x68(%rsp)
jne 0x4dee0
movl 0xe4(%rsp), %edi
movq 0x38(%rsp), %rsi
movq 0xc8(%rsp), %rdx
movq 0x10(%rsp), %rcx
movq 0xc0(%rsp), %r8
callq 0x4d780
cmpl $0x0, 0x98(%rsp)
jne 0x4dede
movl $0x0, 0x8(%rsp)
movl 0x8(%rsp), %eax
cmpl 0x4c(%rsp), %eax
jge 0x4dedc
movq 0xd0(%rsp), %rax
movslq 0x8(%rsp), %rcx
movsd (%rax,%rcx,8), %xmm0
movq 0x10(%rsp), %rax
movslq 0x8(%rsp), %rcx
subsd (%rax,%rcx,8), %xmm0
movq 0x10(%rsp), %rax
movslq 0x8(%rsp), %rcx
movsd %xmm0, (%rax,%rcx,8)
movq 0x10(%rsp), %rax
movslq 0x8(%rsp), %rcx
movsd (%rax,%rcx,8), %xmm0
movq 0x40(%rsp), %rax
movslq 0x8(%rsp), %rcx
divsd (%rax,%rcx,8), %xmm0
movq 0xc8(%rsp), %rax
movslq 0x8(%rsp), %rcx
movsd %xmm0, (%rax,%rcx,8)
movl 0x8(%rsp), %eax
addl $0x1, %eax
movl %eax, 0x8(%rsp)
jmp 0x4de65
jmp 0x4dede
jmp 0x4dee0
jmp 0x4dcb6
leaq 0x60(%rsp), %rdi
callq 0xf240
movq 0x18(%rsp), %rdi
callq 0xed30
movq 0x10(%rsp), %rdi
callq 0xed30
movq 0x40(%rsp), %rdi
callq 0xed30
movq 0x38(%rsp), %rdi
callq 0xed30
addq $0xe8, %rsp
retq
nop