! 'A' now is stored in the compressed sparse row format (CSR)
! The subroutine cg_sparse calls 'amux' (SPARSKIT) and 'ddot' (SPARTSKIT/BLAS)
! So either install SPARSKIT and compile using
!     $ ifort cgs.f90 -L/opt/sparskit -lsparskit
! Or get/download the individual subroutines 'amux' and 'ddot' and build using
!     $ ifort cgs.f90 amux.f90 ddot.f90
!
! stali@purdue.edu

program main
  implicit none
  
  real(8), allocatable :: a(:), x(:), b(:)
  integer :: n, nnz, max_iteration
  integer, allocatable :: ja(:), ia(:)
  real(8) :: tolerance

  n = 5 ! Size of the system
  nnz = 13 ! Number of non zeros in A

  allocate(a(nnz),ja(nnz),ia(n+1),b(n),x(n))

  ! Matrix A is in the CSR format (defined by three 1D arrays)
  a = (/2.0, -1.0, -1.0, 2.0, -1.0, -1.0, 2.0, -1.0, -1.0, 2.0, -1.0, &
       -1.0, 2.0/)
  ja = (/1, 2, 1, 2, 3, 2, 3, 4, 3, 4, 5, 4, 5/)
  ia = (/1, 3, 6, 9, 12, 14/)

  ! RHS
  b = (/0.0, 0.0, 0.0, 0.0, 100.0/)

  tolerance = 0.000001; max_iteration=20
  call cg_sparse(a,ja,ia,x,b,tolerance,max_iteration)
  print*, x

  deallocate(a,ja,ia,b,x)

contains

!------
  subroutine cg_sparse(a,ja,ia,x,b,tolerance,max_iteration)
    implicit none

    real(8) :: a(:), x(:), b(:), tolerance
    integer :: ja(:), ia(:), max_iteration
    
    real(8) :: r(size(x)), Ax(size(x)),p(size(x)), Ap(size(x))
    real(8) :: alpha, beta, rtr, ddot
    integer :: n, k

    n = size(x)
    x = 0.0

    call amux (n, x, Ax, a, ja, ia)
    r = b-Ax
    p = r
    k = 1
    do
       call amux (n, p, Ap, a, ja, ia)
       rtr = ddot (n, r, 1, r, 1)
       alpha = rtr / ddot (n, p, 1, Ap, 1)
       x = x + alpha * p
       r = r - alpha * Ap
       if (sqrt(sum(r**2)) <= tolerance .or. k == max_iteration) exit
       beta = ddot (n, r, 1, r, 1) / rtr
       p = r + beta * p
       k = k + 1
    end do

    if (k == max_iteration) then
         print*, "CG stops without converging at iteration", k   
    else
         print*, "Converged in", k, "iterations"
    end if
end subroutine cg_sparse
!------

end program main