#include <petsc/finclude/petscdmda.h>
#include <petsc/finclude/petscksp.h>
module ex45fmodule
  use petscmat
  use petscksp
  use petscdmda
  use petscvec
  implicit none

contains

  subroutine ComputeInitialGuess(ksp, b, ctx, ierr)
    PetscErrorCode, intent(out) :: ierr
    KSP ksp
    PetscInt ctx(*)
    Vec b
    PetscScalar, parameter :: h = 0.0

    PetscCall(VecSet(b, h, ierr))
  end subroutine

  subroutine ComputeRHS(ksp, b, unused, ierr)
    PetscErrorCode, intent(out) :: ierr
    KSP ksp
    Vec b
    integer unused(*)
    PetscScalar h, Hx, Hy
    PetscInt mx, my
    DM dm

    PetscCall(KSPGetDM(ksp, dm, ierr))
    PetscCall(DMDAGetInfo(dm, PETSC_NULL_INTEGER, mx, my, PETSC_NULL_INTEGER, PETSC_NULL_INTEGER, PETSC_NULL_INTEGER, PETSC_NULL_INTEGER, PETSC_NULL_INTEGER, PETSC_NULL_INTEGER, PETSC_NULL_DMBOUNDARYTYPE, PETSC_NULL_DMBOUNDARYTYPE, PETSC_NULL_DMBOUNDARYTYPE, PETSC_NULL_DMDASTENCILTYPE, ierr))

    Hx = 1.0/real(mx - 1)
    Hy = 1.0/real(my - 1)
    h = Hx*Hy
    PetscCall(VecSet(b, h, ierr))
  end subroutine

  subroutine ComputeMatrix(ksp, A, B, unused, ierr)
    PetscErrorCode, intent(out) :: ierr
    KSP ksp
    Mat A, B
    integer unused(*)
    DM dm

    PetscInt i, j, mx, my, xm, ym, xs, ys
    PetscScalar v(5), Hx, Hy
    PetscScalar HxdHy, HydHx
    MatStencil row(1), col(5)

    PetscCall(KSPGetDM(ksp, dm, ierr))
    PetscCall(DMDAGetInfo(dm, PETSC_NULL_INTEGER, mx, my, PETSC_NULL_INTEGER, PETSC_NULL_INTEGER, PETSC_NULL_INTEGER, PETSC_NULL_INTEGER, PETSC_NULL_INTEGER, PETSC_NULL_INTEGER, PETSC_NULL_DMBOUNDARYTYPE, PETSC_NULL_DMBOUNDARYTYPE, PETSC_NULL_DMBOUNDARYTYPE, PETSC_NULL_DMDASTENCILTYPE, ierr))

    Hx = 1.0/real(mx - 1)
    Hy = 1.0/real(my - 1)
    HxdHy = Hx/Hy
    HydHx = Hy/Hx
    PetscCall(DMDAGetCorners(dm, xs, ys, PETSC_NULL_INTEGER, xm, ym, PETSC_NULL_INTEGER, ierr))
    do j = ys, ys + ym - 1
      do i = xs, xs + xm - 1
        row(1)%i = i
        row(1)%j = j
        if (i == 0 .or. j == 0 .or. i == mx - 1 .or. j == my - 1) then
          v(1) = 2.0*(HxdHy + HydHx)
          PetscCall(MatSetValuesStencil(B, 1_PETSC_INT_KIND, row, 1_PETSC_INT_KIND, row, v, INSERT_VALUES, ierr))
        else
          v(1) = -HxdHy
          col(1)%i = i
          col(1)%j = j - 1
          v(2) = -HydHx
          col(2)%i = i - 1
          col(2)%j = j
          v(3) = 2.0*(HxdHy + HydHx)
          col(3)%i = i
          col(3)%j = j
          v(4) = -HydHx
          col(4)%i = i + 1
          col(4)%j = j
          v(5) = -HxdHy
          col(5)%i = i
          col(5)%j = j + 1
          PetscCall(MatSetValuesStencil(B, 1_PETSC_INT_KIND, row, 5_PETSC_INT_KIND, col, v, INSERT_VALUES, ierr))
        end if
      end do
    end do
    PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY, ierr))
    PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY, ierr))
    if (A /= B) then
      PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY, ierr))
      PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY, ierr))
    end if
  end subroutine
end module ex45fmodule

program main
  use petscdmda
  use petscksp
  use ex45fmodule
  implicit none

  PetscInt is, js, iw, jw
  PetscErrorCode ierr
  KSP ksp
  DM dm

  PetscCallA(PetscInitialize(ierr))
  PetscCallA(KSPCreate(PETSC_COMM_WORLD, ksp, ierr))
  PetscCallA(DMDACreate2D(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, 3_PETSC_INT_KIND, 3_PETSC_INT_KIND, PETSC_DECIDE, PETSC_DECIDE, 1_PETSC_INT_KIND, 1_PETSC_INT_KIND, PETSC_NULL_INTEGER_ARRAY, PETSC_NULL_INTEGER_ARRAY, dm, ierr))
  PetscCallA(DMSetFromOptions(dm, ierr))
  PetscCallA(DMSetUp(dm, ierr))
  PetscCallA(KSPSetDM(ksp, dm, ierr))
  PetscCallA(KSPSetComputeInitialGuess(ksp, ComputeInitialGuess, 0, ierr))
  PetscCallA(KSPSetComputeRHS(ksp, ComputeRHS, 0, ierr))
  PetscCallA(KSPSetComputeOperators(ksp, ComputeMatrix, 0, ierr))
  PetscCallA(DMDAGetCorners(dm, is, js, PETSC_NULL_INTEGER, iw, jw, PETSC_NULL_INTEGER, ierr))
  PetscCallA(KSPSetFromOptions(ksp, ierr))
  PetscCallA(KSPSetUp(ksp, ierr))
  PetscCallA(KSPSolve(ksp, PETSC_NULL_VEC, PETSC_NULL_VEC, ierr))
  PetscCallA(KSPDestroy(ksp, ierr))
  PetscCallA(DMDestroy(dm, ierr))
  PetscCallA(PetscFinalize(ierr))
end

!/*TEST
!
!   test:
!      nsize: 4
!      args: -ksp_monitor_short -da_refine 5 -pc_type mg -pc_mg_levels 5 -mg_levels_ksp_type chebyshev -mg_levels_ksp_max_it 2 -mg_levels_pc_type jacobi -ksp_pc_side right
!
!TEST*/
