!
!     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
!
program main
#include <petsc/finclude/petscksp.h>
  use petscksp
  implicit none

  PetscInt i, n, one
  PetscCount nz
  PetscBool flg, equal
  PetscErrorCode ierr
  PetscInt, ALLOCATABLE :: ia(:)
  PetscInt, ALLOCATABLE :: ja(:)
  PetscScalar, ALLOCATABLE :: a(:)
  PetscScalar, ALLOCATABLE :: x(:)
  PetscScalar, ALLOCATABLE :: b(:)

  PetscInt, ALLOCATABLE :: rows(:)
  PetscInt, ALLOCATABLE :: cols(:)

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

  PetscCallA(PetscInitialize(ierr))
  one = 1
  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))

  do i = 1, n
    b(i) = 1.0
  end do

!     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, one, n, b, rhs, ierr))
  PetscCallA(VecCreateSeqWithArray(PETSC_COMM_SELF, one, 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*/
