!
!
!    Fortran kernel for sparse triangular solve in the BAIJ matrix format
! This ONLY works for factorizations in the NATURAL ORDERING, i.e.
! with MatSolve_SeqBAIJ_4_NaturalOrdering()
!
#include <petsc/finclude/petscsys.h>
!

pure subroutine FortranSolveBAIJ4Unroll(n,x,ai,aj,adiag,a,b)
  use, intrinsic :: ISO_C_binding
  implicit none (type, external)
  MatScalar, intent(in) :: a(0:*)
  PetscScalar, intent(inout) :: x(0:*)
  PetscScalar, intent(in) :: b(0:*)
  PetscInt, intent(in) :: n
  PetscInt, intent(in) :: ai(0:*), aj(0:*), adiag(0:*)

  PetscInt :: i,j,jstart,jend
  PetscInt :: idx,ax,jdx
  PetscScalar :: s(0:3)

  PETSC_AssertAlignx(16,a(1))
  PETSC_AssertAlignx(16,x(1))
  PETSC_AssertAlignx(16,b(1))
  PETSC_AssertAlignx(16,ai(1))
  PETSC_AssertAlignx(16,aj(1))
  PETSC_AssertAlignx(16,adiag(1))

  !
  ! Forward Solve
  !
  x(0:3) = b(0:3)
  idx  = 0
  do i=1,n-1
    jstart = ai(i)
    jend   = adiag(i) - 1
    ax     = 16*jstart
    idx    = idx + 4
    s(0:3) = b(idx+0:idx+3)
    do j=jstart,jend
      jdx = 4*aj(j)

      s(0) = s(0)-(a(ax+0)*x(jdx+0)+a(ax+4)*x(jdx+1)+a(ax+ 8)*x(jdx+2)+a(ax+12)*x(jdx+3))
      s(1) = s(1)-(a(ax+1)*x(jdx+0)+a(ax+5)*x(jdx+1)+a(ax+ 9)*x(jdx+2)+a(ax+13)*x(jdx+3))
      s(2) = s(2)-(a(ax+2)*x(jdx+0)+a(ax+6)*x(jdx+1)+a(ax+10)*x(jdx+2)+a(ax+14)*x(jdx+3))
      s(3) = s(3)-(a(ax+3)*x(jdx+0)+a(ax+7)*x(jdx+1)+a(ax+11)*x(jdx+2)+a(ax+15)*x(jdx+3))
      ax = ax + 16
    end do
    x(idx+0:idx+3) = s(0:3)
  end do

  !
  ! Backward solve the upper triangular
  !
  do i=n-1,0,-1
    jstart = adiag(i) + 1
    jend   = ai(i+1) - 1
    ax     = 16*jstart
    s(0:3) = x(idx+0:idx+3)
    do j=jstart,jend
      jdx   = 4*aj(j)
      s(0) = s(0)-(a(ax+0)*x(jdx+0)+a(ax+4)*x(jdx+1)+a(ax+ 8)*x(jdx+2)+a(ax+12)*x(jdx+3))
      s(1) = s(1)-(a(ax+1)*x(jdx+0)+a(ax+5)*x(jdx+1)+a(ax+ 9)*x(jdx+2)+a(ax+13)*x(jdx+3))
      s(2) = s(2)-(a(ax+2)*x(jdx+0)+a(ax+6)*x(jdx+1)+a(ax+10)*x(jdx+2)+a(ax+14)*x(jdx+3))
      s(3) = s(3)-(a(ax+3)*x(jdx+0)+a(ax+7)*x(jdx+1)+a(ax+11)*x(jdx+2)+a(ax+15)*x(jdx+3))
      ax = ax + 16
    end do
    ax      = 16*adiag(i)
    x(idx+0) = a(ax+0)*s(0)+a(ax+4)*s(1)+a(ax+ 8)*s(2)+a(ax+12)*s(3)
    x(idx+1) = a(ax+1)*s(0)+a(ax+5)*s(1)+a(ax+ 9)*s(2)+a(ax+13)*s(3)
    x(idx+2) = a(ax+2)*s(0)+a(ax+6)*s(1)+a(ax+10)*s(2)+a(ax+14)*s(3)
    x(idx+3) = a(ax+3)*s(0)+a(ax+7)*s(1)+a(ax+11)*s(2)+a(ax+15)*s(3)
    idx      = idx - 4
  end do
end subroutine FortranSolveBAIJ4Unroll

!   version that does not call BLAS 2 operation for each row block
!
pure subroutine FortranSolveBAIJ4(n,x,ai,aj,adiag,a,b,w)
  use, intrinsic :: ISO_C_binding
  implicit none
  MatScalar, intent(in) :: a(0:*)
  PetscScalar, intent(inout) :: x(0:*),w(0:*)
  PetscScalar, intent(in) :: b(0:*)
  PetscInt, intent(in) :: n
  PetscInt, intent(in) :: ai(0:*), aj(0:*), adiag(0:*)

  PetscInt :: ii,jj,i,j
  PetscInt :: jstart,jend,idx,ax,jdx,kdx,nn
  PetscScalar :: s(0:3)

  PETSC_AssertAlignx(16,a(1))
  PETSC_AssertAlignx(16,w(1))
  PETSC_AssertAlignx(16,x(1))
  PETSC_AssertAlignx(16,b(1))
  PETSC_AssertAlignx(16,ai(1))
  PETSC_AssertAlignx(16,aj(1))
  PETSC_AssertAlignx(16,adiag(1))
  !
  !     Forward Solve
  !
  x(0:3) = b(0:3)
  idx  = 0
  do i=1,n-1
    !
    ! Pack required part of vector into work array
    !
    kdx    = 0
    jstart = ai(i)
    jend   = adiag(i) - 1

    if (jend - jstart >= 500) error stop 'Overflowing vector FortranSolveBAIJ4()'

    do j=jstart,jend
      jdx       = 4*aj(j)
      w(kdx:kdx+3) = x(jdx:jdx+3)
      kdx       = kdx + 4
    end do

    ax     = 16*jstart
    idx    = idx + 4
    s(0:3) = b(idx:idx+3)
    !
    !    s = s - a(ax:)*w
    !
    nn = 4*(jend - jstart + 1) - 1
    do ii=0,3
      do jj=0,nn
        s(ii) = s(ii) - a(ax+4*jj+ii)*w(jj)
      end do
    end do

    x(idx:idx+3) = s(0:3)
  end do
  !
  ! Backward solve the upper triangular
  !
  do i=n-1,0,-1
     jstart = adiag(i) + 1
     jend   = ai(i+1) - 1
     ax     = 16*jstart
     s(0:3) = x(idx:idx+3)
     !
     !   Pack each chunk of vector needed
     !
     kdx = 0
     if (jend - jstart >= 500) error stop 'Overflowing vector FortranSolveBAIJ4()'

     do j=jstart,jend
       jdx = 4*aj(j)
       w(kdx:kdx+3) = x(jdx:jdx+3)
       kdx = kdx + 4
     end do
     nn = 4*(jend - jstart + 1) - 1
     do ii=0,3
       do jj=0,nn
         s(ii) = s(ii) - a(ax+4*jj+ii)*w(jj)
       end do
     end do

     ax      = 16*adiag(i)
     x(idx)  = a(ax+0)*s(0)+a(ax+4)*s(1)+a(ax+ 8)*s(2)+a(ax+12)*s(3)
     x(idx+1)= a(ax+1)*s(0)+a(ax+5)*s(1)+a(ax+ 9)*s(2)+a(ax+13)*s(3)
     x(idx+2)= a(ax+2)*s(0)+a(ax+6)*s(1)+a(ax+10)*s(2)+a(ax+14)*s(3)
     x(idx+3)= a(ax+3)*s(0)+a(ax+7)*s(1)+a(ax+11)*s(2)+a(ax+15)*s(3)
     idx     = idx - 4
  end do
end subroutine FortranSolveBAIJ4