!
!     Creates a tridiagonal sparse matrix explicitly in Fortran and solves a linear system with it
!
!     The matrix is provided in two three ways
!       Compressed sparse row: ia(), ja(), and a()
!     Entry triples:  rows(), cols(), and a()
!     Entry triples in a way that supports new nonzero values with the same nonzero structure
!
#include <petsc/finclude/petscksp.h>
program main
  use petscksp
  implicit none

  PetscInt i, n
  PetscCount nz
  PetscBool flg, equal
  PetscErrorCode ierr
  PetscInt, allocatable, dimension(:) :: ia, ja, rows, cols
  PetscScalar, allocatable, dimension(:) :: a, x, b

  Mat J, Jt, Jr
  Vec rhs, solution
  KSP ksp
  PC pc

  PetscCallA(PetscInitialize(ierr))
  n = 3
  PetscCallA(PetscOptionsGetInt(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-n', n, flg, ierr))
  nz = 3*n - 4

  allocate (b(n), x(n))

!     Fill the sparse matrix representation
  allocate (ia(n + 1), ja(nz), a(nz))
  allocate (rows(nz), cols(nz))

  b(1:n) = 1.0

!     PETSc ia() and ja() values begin at 0, not 1, you may need to shift the indices used in your code
  ia(1) = 0
  ia(2) = 1
  do i = 3, n
    ia(i) = ia(i - 1) + 3
  end do
  ia(n + 1) = ia(n) + 1

  ja(1) = 0
  rows(1) = 0; cols(1) = 0
  a(1) = 1.0
  do i = 2, n - 1
    ja(2 + 3*(i - 2)) = i - 2
    rows(2 + 3*(i - 2)) = i - 1; cols(2 + 3*(i - 2)) = i - 2
    a(2 + 3*(i - 2)) = -1.0
    ja(2 + 3*(i - 2) + 1) = i - 1
    rows(2 + 3*(i - 2) + 1) = i - 1; cols(2 + 3*(i - 2) + 1) = i - 1
    a(2 + 3*(i - 2) + 1) = 2.0
    ja(2 + 3*(i - 2) + 2) = i
    rows(2 + 3*(i - 2) + 2) = i - 1; cols(2 + 3*(i - 2) + 2) = i
    a(2 + 3*(i - 2) + 2) = -1.0
  end do
  ja(nz) = n - 1
  rows(nz) = n - 1; cols(nz) = n - 1
  a(nz) = 1.0

  PetscCallA(MatCreateSeqAIJWithArrays(PETSC_COMM_SELF, n, n, ia, ja, a, J, ierr))
  PetscCallA(MatCreateSeqAIJFromTriple(PETSC_COMM_SELF, n, n, rows, cols, a, Jt, nz, PETSC_FALSE, ierr))
  PetscCallA(MatEqual(J, Jt, equal, ierr))
  PetscCheckA(equal .eqv. PETSC_TRUE, PETSC_COMM_SELF, PETSC_ERR_PLIB, 'Matrices J and Jt must be equal')
  PetscCallA(MatDestroy(Jt, ierr))
  PetscCallA(MatCreate(PETSC_COMM_SELF, Jr, ierr))
  PetscCallA(MatSetSizes(Jr, n, n, n, n, ierr))
  PetscCallA(MatSetType(Jr, MATSEQAIJ, ierr))
  PetscCallA(MatSetPreallocationCOO(Jr, nz, rows, cols, ierr))
  PetscCallA(MatSetValuesCOO(Jr, a, INSERT_VALUES, ierr))
  PetscCallA(MatEqual(J, Jr, equal, ierr))
  PetscCheckA(equal .eqv. PETSC_TRUE, PETSC_COMM_SELF, PETSC_ERR_PLIB, 'Matrices J and Jr must be equal')

  PetscCallA(VecCreateSeqWithArray(PETSC_COMM_SELF, 1_PETSC_INT_KIND, n, b, rhs, ierr))
  PetscCallA(VecCreateSeqWithArray(PETSC_COMM_SELF, 1_PETSC_INT_KIND, n, x, solution, ierr))

  PetscCallA(KSPCreate(PETSC_COMM_SELF, ksp, ierr))
  PetscCallA(KSPSetErrorIfNotConverged(ksp, PETSC_TRUE, ierr))
! Default to a direct sparse LU solver for robustness
  PetscCallA(KSPGetPC(ksp, pc, ierr))
  PetscCallA(PCSetType(pc, PCLU, ierr))
  PetscCallA(KSPSetFromOptions(ksp, ierr))
  PetscCallA(KSPSetOperators(ksp, J, J, ierr))

  PetscCallA(KSPSolve(ksp, rhs, solution, ierr))

! Keep the same size and nonzero structure of the matrix but change its numerical entries
  do i = 2, n - 1
    a(2 + 3*(i - 2) + 1) = 4.0
  end do
  PetscCallA(PetscObjectStateIncrease(J, ierr))
  PetscCallA(MatSetValuesCOO(Jr, a, INSERT_VALUES, ierr))
  PetscCallA(MatEqual(J, Jr, equal, ierr))
  PetscCheckA(equal .eqv. PETSC_TRUE, PETSC_COMM_SELF, PETSC_ERR_PLIB, 'Matrices J and Jr should be equal')
  PetscCallA(MatDestroy(Jr, ierr))

  PetscCallA(KSPSolve(ksp, rhs, solution, ierr))

  PetscCallA(KSPDestroy(ksp, ierr))
  PetscCallA(VecDestroy(rhs, ierr))
  PetscCallA(VecDestroy(solution, ierr))
  PetscCallA(MatDestroy(J, ierr))

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

  PetscCallA(PetscFinalize(ierr))
end

!/*TEST
!
!     test:
!       args: -ksp_monitor -ksp_view
!
!TEST*/
