module sor_iterate
  use globals

  implicit none

contains

  subroutine sor_heatsink
    real(8) :: sor, residual, steperror, my_steperror, sorfrac
    real(8) :: airhalo
    integer :: i, j, k
    integer :: numiter
    integer :: Chunk_size, Slice_size

    Chunk_size = Grid_size / mpi_cpu_count
    Slice_size = (Grid_size+2)**2

    numiter = 0
    steperror = (Grid_size**3)*1e10  ! start while loop

    do while( numiter < maxiter .and. steperror > errorstop )
       steperror = 0d0

       ! final sor fraction is not optimal for 1st iterations
       if ( numiter < Grid_size ) then
          sorfrac = 1d0 + (sorfracfinal-1d0)*real(numiter)/real(Grid_size)
       else
          sorfrac = sorfracfinal
       end if

       ! relax red nodes

       do i = 1, Grid_size
          do j = 1, Chunk_size
             do k = 1+mod(i+j,2), Grid_size, 2
                ! skip air cells
                if (heatsink_design(i,j,k) == Air) then
                   cycle
                end if
                ! set air boundaries
                airhalo = temp3d(i,j,k) - C_al_factor * (temp3d(i,j,k) - Air_temp_K) * Delta
                if (heatsink_design(i-1,j,k) == Air) then
                   temp3d(i-1,j,k) = airhalo
                end if
                if (heatsink_design(i+1,j,k) == Air) then
                   temp3d(i+1,j,k) = airhalo
                end if
                if (heatsink_design(i,j-1,k) == Air) then
                   temp3d(i,j-1,k) = airhalo
                end if
                if (heatsink_design(i,j+1,k) == Air) then
                   temp3d(i,j+1,k) = airhalo
                end if
                if (heatsink_design(i,j,k-1) == Air) then
                   temp3d(i,j,k-1) = airhalo
                end if
                if (heatsink_design(i,j,k+1) == Air) then
                   temp3d(i,j,k+1) = airhalo
                end if
                ! lower boundary
                if (k == 1) then
                   if (heatsink_design(i,j,k) == Aluminium_touching_cpu) then
                      ! CPU input
                      temp3d(i,j,k-1) = temp3d(i,j,k) + C_cpu * Delta
                   else
                      ! zero-flow boundary
                      temp3d(i,j,k-1) = temp3d(i,j,k)
                   end if
                end if
                ! relax
                sor = onesixth*&
                     & (temp3d(i-1,j,k)+temp3d(i,j-1,k)+temp3d(i,j,k-1)+&
                     &  temp3d(i+1,j,k)+temp3d(i,j+1,k)+temp3d(i,j,k+1))
                residual = sor - temp3d(i,j,k)
                temp3d(i,j,k) = temp3d(i,j,k) + sorfrac * residual

                steperror = steperror + abs(residual)
             end do
          end do
       end do

       ! pass red ghost points

       if (mod(mpi_cpu_id,2)==0) then
          ! even ranks
          if (mpi_cpu_id > 0) then
             ! send to left
             call MPI_Send(&
                  & temp3d(:,1,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id-1, 0, &
                  & MPI_COMM_WORLD, mpi_error_code)
          end if
          if (mpi_cpu_id < mpi_cpu_count-1) then
             ! send to right
             call MPI_Send(&
                  & temp3d(:,Chunk_size,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id+1, 0, &
                  & MPI_COMM_WORLD, mpi_error_code)
          end if
          if (mpi_cpu_id < mpi_cpu_count-1) then
             ! receive from right
             call MPI_Recv(&
                  & temp3d(:,Chunk_size+1,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id+1, 0, &
                  & MPI_COMM_WORLD, mpistatus, mpi_error_code)
          end if
          if (mpi_cpu_id > 0) then
             ! receive from left
             call MPI_Recv(&
                  & temp3d(:,0,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id-1, 0, &
                  & MPI_COMM_WORLD, mpistatus, mpi_error_code)
          end if
       end if

       if (mod(mpi_cpu_id,2)==1) then
          ! odd ranks
          if (mpi_cpu_id < mpi_cpu_count-1) then
             ! receive from right
             call MPI_Recv(&
                  & temp3d(:,Chunk_size+1,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id+1, 0, &
                  & MPI_COMM_WORLD, mpistatus, mpi_error_code)
          end if
          if (mpi_cpu_id > 0) then
             ! receive from left
             call MPI_Recv(&
                  & temp3d(:,0,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id-1, 0, &
                  & MPI_COMM_WORLD, mpistatus, mpi_error_code)
          end if
          if (mpi_cpu_id > 0) then
             ! send to left
             call MPI_Send(&
                  & temp3d(:,1,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id-1, 0, &
                  & MPI_COMM_WORLD, mpi_error_code)
          end if
          if (mpi_cpu_id < mpi_cpu_count-1) then
             ! send to right
             call MPI_Send(&
                  & temp3d(:,Chunk_size,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id+1, 0, &
                  & MPI_COMM_WORLD, mpi_error_code)
          end if
       end if

       ! relax black nodes

       do i = 1, Grid_size
          do j = 1, Chunk_size
             do k = 1+mod(i+j+1,2), Grid_size, 2
                ! skip air cells
                if (heatsink_design(i,j,k) == Air) then
                   cycle
                end if
                ! set air boundaries
                airhalo = temp3d(i,j,k) - C_al_factor * (temp3d(i,j,k) - Air_temp_K) * Delta
                if (heatsink_design(i-1,j,k) == Air) then
                   temp3d(i-1,j,k) = airhalo
                end if
                if (heatsink_design(i+1,j,k) == Air) then
                   temp3d(i+1,j,k) = airhalo
                end if
                if (heatsink_design(i,j-1,k) == Air) then
                   temp3d(i,j-1,k) = airhalo
                end if
                if (heatsink_design(i,j+1,k) == Air) then
                   temp3d(i,j+1,k) = airhalo
                end if
                if (heatsink_design(i,j,k-1) == Air) then
                   temp3d(i,j,k-1) = airhalo
                end if
                if (heatsink_design(i,j,k+1) == Air) then
                   temp3d(i,j,k+1) = airhalo
                end if
                ! lower boundary
                if (k == 1) then
                   if (heatsink_design(i,j,k) == Aluminium_touching_cpu) then
                      ! CPU input
                      temp3d(i,j,k-1) = temp3d(i,j,k) + C_cpu * Delta
                   else
                      ! zero-flow boundary
                      temp3d(i,j,k-1) = temp3d(i,j,k)
                   end if
                end if
                ! relax
                sor = onesixth*&
                     & (temp3d(i-1,j,k)+temp3d(i,j-1,k)+temp3d(i,j,k-1)+&
                     &  temp3d(i+1,j,k)+temp3d(i,j+1,k)+temp3d(i,j,k+1))
                residual = sor - temp3d(i,j,k)
                temp3d(i,j,k) = temp3d(i,j,k) + sorfrac * residual

                steperror = steperror + abs(residual)
             end do
          end do
       end do

       ! pass black ghost points

       if (mod(mpi_cpu_id,2)==0) then
          ! even ranks
          if (mpi_cpu_id > 0) then
             ! send to left
             call MPI_Send(&
                  & temp3d(:,1,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id-1, 0, &
                  & MPI_COMM_WORLD, mpi_error_code)
          end if
          if (mpi_cpu_id < mpi_cpu_count-1) then
             ! send to right
             call MPI_Send(&
                  & temp3d(:,Chunk_size,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id+1, 0, &
                  & MPI_COMM_WORLD, mpi_error_code)
          end if
          if (mpi_cpu_id < mpi_cpu_count-1) then
             ! receive from right
             call MPI_Recv(&
                  & temp3d(:,Chunk_size+1,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id+1, 0, &
                  & MPI_COMM_WORLD, mpistatus, mpi_error_code)
          end if
          if (mpi_cpu_id > 0) then
             ! receive from left
             call MPI_Recv(&
                  & temp3d(:,0,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id-1, 0, &
                  & MPI_COMM_WORLD, mpistatus, mpi_error_code)
          end if
       end if

       if (mod(mpi_cpu_id,2)==1) then
          ! odd ranks
          if (mpi_cpu_id < mpi_cpu_count-1) then
             ! receive from right
             call MPI_Recv(&
                  & temp3d(:,Chunk_size+1,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id+1, 0, &
                  & MPI_COMM_WORLD, mpistatus, mpi_error_code)
          end if
          if (mpi_cpu_id > 0) then
             ! receive from left
             call MPI_Recv(&
                  & temp3d(:,0,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id-1, 0, &
                  & MPI_COMM_WORLD, mpistatus, mpi_error_code)
          end if
          if (mpi_cpu_id > 0) then
             ! send to left
             call MPI_Send(&
                  & temp3d(:,1,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id-1, 0, &
                  & MPI_COMM_WORLD, mpi_error_code)
          end if
          if (mpi_cpu_id < mpi_cpu_count-1) then
             ! send to right
             call MPI_Send(&
                  & temp3d(:,Chunk_size,:), &
                  & Slice_size, MPI_DOUBLE_PRECISION, mpi_cpu_id+1, 0, &
                  & MPI_COMM_WORLD, mpi_error_code)
          end if
       end if

       my_steperror = steperror
       call MPI_Allreduce(my_steperror, steperror, 1, &
            & MPI_DOUBLE_PRECISION, MPI_SUM, MPI_COMM_WORLD, mpi_error_code)

       if (mpi_cpu_id == 0) then
          if ( (mod(numiter,100) == 0) .or. (steperror <= errorstop) ) then
             write(*,'(a,i8,a,es12.4)'),'iteration=',numiter,'  error=',steperror
          end if
       end if
       
       numiter = numiter +1 

    end do

  end subroutine sor_heatsink


  subroutine serial_sor_heatsink
    real(8) :: sor, residual, steperror, sorfrac
    real(8) :: airhalo
    integer :: i, j, k
    integer :: numiter

    numiter = 0
    steperror = (Grid_size**3)*1e10  ! start while loop

    do while( numiter < maxiter .and. steperror > errorstop )
       steperror = 0d0

       ! final sor fraction is not optimal for 1st iterations
       if ( numiter < Grid_size ) then
          sorfrac = 1d0 + (sorfracfinal-1d0)*real(numiter)/real(Grid_size)
       else
          sorfrac = sorfracfinal
       end if

       do i = 1, Grid_size
          do j = 1, Grid_size / mpi_cpu_count
             do k = 1, Grid_size
                ! skip air cells
                if (heatsink_design(i,j,k) == Air) then
                   cycle
                end if
                ! set air boundaries
                airhalo = temp3d(i,j,k) - C_al_factor * (temp3d(i,j,k) - Air_temp_K) * Delta
                if (heatsink_design(i-1,j,k) == Air) then
                   temp3d(i-1,j,k) = airhalo
                end if
                if (heatsink_design(i+1,j,k) == Air) then
                   temp3d(i+1,j,k) = airhalo
                end if
                if (heatsink_design(i,j-1,k) == Air) then
                   temp3d(i,j-1,k) = airhalo
                end if
                if (heatsink_design(i,j+1,k) == Air) then
                   temp3d(i,j+1,k) = airhalo
                end if
                if (heatsink_design(i,j,k-1) == Air) then
                   temp3d(i,j,k-1) = airhalo
                end if
                if (heatsink_design(i,j,k+1) == Air) then
                   temp3d(i,j,k+1) = airhalo
                end if
                ! lower boundary
                if (k == 1) then
                   if (heatsink_design(i,j,k) == Aluminium_touching_cpu) then
                      ! CPU input
                      temp3d(i,j,k-1) = temp3d(i,j,k) + C_cpu * Delta
                   else
                      ! zero-flow boundary
                      temp3d(i,j,k-1) = temp3d(i,j,k)
                   end if
                end if
                ! relax
                sor = onesixth*&
                     & (temp3d(i-1,j,k)+temp3d(i,j-1,k)+temp3d(i,j,k-1)+&
                     &  temp3d(i+1,j,k)+temp3d(i,j+1,k)+temp3d(i,j,k+1))
                residual = sor - temp3d(i,j,k)
                temp3d(i,j,k) = temp3d(i,j,k) + sorfrac * residual

                steperror = steperror + abs(residual)
             end do
          end do
       end do

       if ( mod(numiter,100) == 0 ) then
          write(*,'(a,i8,a,es12.4)'),'iteration=',numiter,'  error=',steperror
       end if

       numiter = numiter +1 

    end do

  end subroutine serial_sor_heatsink

end module sor_iterate