#include <petsc/finclude/petsc.h>
#include <slepc/finclude/slepc.h>
#include "macros.h"

!====================================================================================================
   PROGRAM test_nep
!====================================================================================================
!
      USE slepcnep

      IMPLICIT NONE
      NEP            :: nep
      Mat            :: A, B
      PetscMPIInt    :: rank, nsize
      PetscViewer    :: viewer
      PetscInt       :: its
      PetscInt       :: n     = 10
      PetscInt       :: nev   = 5
      PetscInt       :: maxit = 100
      PetscReal      :: tol   = 1.0E-07
      PetscScalar    :: lambda
      PetscErrorCode :: ierr

      INTEGER        :: ShellMat_Ax_iter
      INTEGER        :: ShellMat_Bx_iter
!
!.... initialize counters
      ShellMat_Ax_iter = 0
      ShellMat_Bx_iter = 0
      its              = 0
!
!.... Initialize SLEPc
      PetscCall(SlepcInitialize(PETSC_NULL_CHARACTER, ierr))
!
!.... Option string
      PetscCall(PetscOptionsInsertString(PETSC_NULL_OPTIONS, "-pc_type none -rg_type interval -rg_interval_endpoints 0,10,0,0" , ierr))
!
!.... Print out message
      PetscCall(PetscPrintf(PETSC_COMM_SELF, "==============================\n", ierr))
      PetscCall(PetscPrintf(PETSC_COMM_SELF, " Nonlinear Eigenvalue Problem \n", ierr))
      PetscCall(PetscPrintf(PETSC_COMM_SELF, "==============================\n", ierr))
!
!.... get size and rank of MPI process
      PetscCall(MPI_Comm_size(PETSC_COMM_SELF,nsize,ierr))
      PetscCall(MPI_Comm_rank(PETSC_COMM_SELF,rank ,ierr))
!
!.... View command line table
      PetscCall(PetscViewerASCIIOpen(PETSC_COMM_SELF,'stdout', viewer, ierr))
      PetscCall(PetscOptionsView(PETSC_NULL_OPTIONS, viewer, ierr))

!.... Create matrix-free operators for A and B
      PetscCall(MatCreateShell(PETSC_COMM_SELF,n,n,PETSC_DETERMINE,PETSC_DETERMINE, PETSC_NULL_INTEGER, A, ierr))
      PetscCall(MatCreateShell(PETSC_COMM_SELF,n,n,PETSC_DETERMINE,PETSC_DETERMINE, PETSC_NULL_INTEGER, B, ierr))
      PetscCall(MatShellSetOperation(A, MATOP_MULT, MatMult_A, ierr))
      PetscCall(MatShellSetOperation(B, MATOP_MULT, MatMult_B, ierr))
!
!.... Create nonlinear eigensolver
      PetscCall(NEPCreate(PETSC_COMM_SELF, nep, ierr))
!
!.... Set functions and Jacobians for NEP
      PetscCall(NEPSetFunction(nep, A, A, MyNEPFunction, PETSC_NULL_INTEGER, ierr))
      PetscCall(NEPSetJacobian(nep, B,    MyNEPJacobian, PETSC_NULL_INTEGER, ierr))
!
!.... Set the problem type
      PetscCall(NEPSetProblemType(nep, NEP_GENERAL, ierr))
!
!.... set the solver type
      PetscCall(NEPSetType(nep, NEPNLEIGS, ierr))
!
!.... Number of requested eigenvalues
      PetscCall(NEPSetDimensions(nep, nev, PETSC_DEFAULT_INTEGER, PETSC_DEFAULT_INTEGER, ierr))
!
!.... set stopping criteria
      PetscCall(NEPSetTolerances(nep, tol, maxit, ierr))
!
!.... Set solver parameters at runtime using command line options
      PetscCall(NEPSetFromOptions(nep, ierr))
!
!.... View NEP setup
      PetscCall(NEPView(nep, viewer, ierr))
!
!.... Print out message
      PetscCall(PetscPrintf(PETSC_COMM_SELF, "==============================\n", ierr))
      PetscCall(PetscPrintf(PETSC_COMM_SELF, " Nonlinear Eigenvalue Problem \n", ierr))
      PetscCall(PetscPrintf(PETSC_COMM_SELF, "==============================\n", ierr))
!
!.... Solve the eigenvalue problem
      PetscCall(NEPSolve(nep, ierr))
!
!.... View NEP setup
      PetscCall(NEPView(nep, viewer, ierr))
!
!.... Get the number of converged eigenpairs
      PetscCall(NEPGetConverged(nep, nev, ierr))
!
!.... Destroy the eigensolver
      PetscCall(NEPDestroy(nep, ierr))
!
!.... Destroy matrices
      PetscCall(MatDestroy(A, ierr))
      PetscCall(MatDestroy(B, ierr))
!
!.... Destroy the viewer (this has already been destroyed after PetscOptionsView, but if it was created again, then it would be destroyed here)
      PetscCall(PetscViewerDestroy(viewer, ierr))
!
!.... Finalize SLEPc
      PetscCall(SlepcFinalize(ierr))

   CONTAINS
!
!====================================================================================================
   SUBROUTINE MyNEPFunction(nep, slepc_lambda, T, P, ctx, ierr)
!====================================================================================================
!
      NEP            :: nep
      PetscScalar    :: slepc_lambda
      Mat            :: T, P
      PetscInt       :: ctx
      PetscErrorCode :: ierr

      WRITE(*,FMT='(a)',ADVANCE='NO') 'F'
      lambda = slepc_lambda

   END SUBROUTINE MyNEPFunction
!
!====================================================================================================
   SUBROUTINE MyNEPJacobian(nep, slepc_lambda, T, P, ctx, ierr)
!====================================================================================================

      NEP            :: nep
      PetscScalar    :: slepc_lambda
      Mat            :: T, P
      PetscInt       :: ctx
      PetscErrorCode :: ierr

      WRITE(*,FMT='(a)',ADVANCE='NO') 'J'
      lambda = slepc_lambda

   END SUBROUTINE MyNEPJacobian
!
!====================================================================================================
   SUBROUTINE MatMult_A(A, x, y, ierr)
!====================================================================================================
!
!.... declare passed variables
      Mat            :: A
      Vec            :: x, y
      PetscErrorCode :: ierr

      PetscScalar    :: xArray(10)
      PetscInt       :: indices(10)
      PetscScalar    :: yArray(10)
      PetscScalar    :: A_matrix(10,10)
!
!.... Sample 10x10 matrix
      DATA A_matrix / &
         2, 1, 0, 0, 0, 0, 0, 0, 0, 1, &
         1, 3, 1, 0, 0, 0, 0, 0, 0, 0, &
         0, 1, 4, 1, 0, 0, 0, 0, 0, 0, &
         0, 0, 1, 5, 1, 0, 0, 0, 0, 0, &
         0, 0, 0, 1, 6, 1, 0, 0, 0, 0, &
         0, 0, 0, 0, 1, 7, 1, 0, 0, 0, &
         0, 0, 0, 0, 0, 1, 8, 1, 0, 0, &
         0, 0, 0, 0, 0, 0, 1, 9, 1, 0, &
         0, 0, 0, 0, 0, 0, 0, 1, 10, 1, &
         1, 0, 0, 0, 0, 0, 0, 0, 1, 2 /
      DATA indices /0,1,2,3,4,5,6,7,8,9/
!
!.... heartbeat and counter
      WRITE(*,FMT='(a)',ADVANCE='NO') 'A'
      ShellMat_Ax_iter = ShellMat_Ax_iter + 1
!
!.... Get local size of vector x
      CALL VecGetLocalSize(x, n, ierr)
!
!.... If x is not of size 10, return an error
      IF (n .NE. 10) STOP "Error: x is not of size 10"
!
!.... Get array of values from vector x using read-only access
      PetscCall(VecGetValues(x, n, indices, xArray, ierr))
!
!.... Compute multiplication for each row of A and set y values
      yArray = matmul(A_matrix,xArray) - lambda*xArray
!
!.... Assemble the resulting vector y
      PetscCall(VecSetValues(y, n, indices, yArray, INSERT_VALUES, ierr))

      PetscCall(VecAssemblyBegin(y, ierr))
      PetscCall(VecAssemblyEnd  (y, ierr))

   END SUBROUTINE MatMult_A
!
!====================================================================================================
   SUBROUTINE MatMult_B(B, x, y, ierr)
!====================================================================================================
!
      Mat            :: B
      Vec            :: x
      Vec            :: y
      PetscErrorCode :: ierr

      WRITE(*,FMT='(a)',ADVANCE='NO') 'B'
      ShellMat_Bx_iter = ShellMat_Bx_iter + 1

      PetscCall(VecCopy(x,y,ierr))
      PetscCall(VecScale(y, -1.0, ierr))

   END SUBROUTINE MatMult_B

   END PROGRAM test_nep