!***********************************************************************
!  Integrated Water Flow Model (IWFM)
!  Copyright (C) 2005  
!  State of California, Department of Water Resources 
!
!  This program is free software; you can redistribute it and/or
!  modify it under the terms of the GNU General Public License
!  as published by the Free Software Foundation; either version 2
!  of the License, or (at your option) any later version.
!
!  This program is distributed in the hope that it will be useful,
!  but WITHOUT ANY WARRANTY; without even the implied warranty of
!  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!  GNU General Public License for more details.
!  (http://www.gnu.org/copyleft/gpl.html)
!
!  You should have received a copy of the GNU General Public License
!  along with this program; if not, write to the Free Software
!  Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
!
!  For tecnical support, e-mail: IWFMtechsupport@water.ca.gov 
!
!***********************************************************************
!   Name:    MatrixSolver.f90						      
!   Version: 1.1						       	
!   File dependences: pgmres.f, messagelogger.mod
!***********************************************************************

MODULE MatrixSolver
   USE MessageLogger
   IMPLICIT NONE
 !Default; indicates this variable is not set yet
  REAL(8), PRIVATE, SAVE::STOPC_MATRIX = -1.0  ,&  
                	  RELAX        = -1.0

  INTEGER, PRIVATE, SAVE::MSOLVE       = -1 , &
                	  MXITER       = -1 , &
                	  ITS          =  0		

  LOGICAL, PRIVATE, SAVE::ZERORESET   = .TRUE.	

  ! -------------------------------------------------------------
  ! --- PRIVATE METHODS
  ! -------------------------------------------------------------

	
  PRIVATE::CONVERTTOCRSFORMAT
  ! -------------------------------------------------------------
  ! --- PUBLIC METHODS
  ! -------------------------------------------------------------

  PUBLIC::SOLVE              , &
          SetMatrixSolver    , &
          SetSORRelax        , &
          SetMaxIteration    , &
          SetSolverTolerance , &
          SetZeroReset	     , &
          GetMaxIteration    , &
          GetNumIterations   , &
	  GetZeroReset	


CONTAINS



! *************************************************************
! *************************************************************
! **** SETTERS
! *************************************************************
! *************************************************************

  ! -------------------------------------------------------------
  ! --- SET SOLVER METHOD
  ! -------------------------------------------------------------
  SUBROUTINE SetMatrixSolver(Method)
    INTEGER,INTENT(IN)::Method

    MSOLVE = Method
   
  END SUBROUTINE SetMatrixSolver


  ! -------------------------------------------------------------
  ! --- SET MAXIMUM NUMBER OF ITERATIONS
  ! -------------------------------------------------------------
  SUBROUTINE SetMaxIteration(IterMax)
    INTEGER,INTENT(IN)::IterMax

    MXITER = IterMax

  END SUBROUTINE SetMaxIteration


  ! -------------------------------------------------------------
  ! --- SET CONVERGENCE CRITERIA
  ! -------------------------------------------------------------
  SUBROUTINE SetSolverTolerance(Toler)
    REAL(8),INTENT(IN)::Toler

    STOPC_MATRIX = Toler

  END SUBROUTINE SetSolverTolerance


  ! -------------------------------------------------------------
  ! --- SET RELAXATION PARAMETER FOR SOR
  ! -------------------------------------------------------------
  SUBROUTINE SetSORRelax(R)
    REAL(8),INTENT(IN)::R

    RELAX = R

  END SUBROUTINE SetSORRelax

  ! -------------------------------------------------------------
  ! --- SET ZERORESET FLAG FOR SOLVERS
  ! -------------------------------------------------------------
  SUBROUTINE SetZeroReset(FLAG)
    LOGICAL,INTENT(IN)::FLAG

    ZERORESET = FLAG

  END SUBROUTINE SetZeroReset



! *************************************************************
! *************************************************************
! **** GETTERS
! *************************************************************
! *************************************************************

  ! -------------------------------------------------------------
  ! --- GET MAXIMUM ITERATION
  ! -------------------------------------------------------------
  FUNCTION GetMaxIteration() RESULT(IterMax)
    INTEGER::IterMax

    IterMax = MXITER

  END FUNCTION GetMaxIteration

  ! -------------------------------------------------------------
  ! --- GET Number of ITERATIONs
  ! -------------------------------------------------------------
  
  FUNCTION GetNumIterations() RESULT(Iter)
    INTEGER::Iter

    Iter = ITS

  END FUNCTION GetNumIterations

  ! -------------------------------------------------------------
  ! --- GET ZeroReset Flag
  ! -------------------------------------------------------------
  
  FUNCTION GetZeroReset() RESULT(FLAG)
    LOGICAL::FLAG

    FLAG = ZERORESET

  END FUNCTION GetZeroReset


! *************************************************************
! *************************************************************
! **** GATEWAY METHOD FOR MATRIX SOLVERS
! *************************************************************
! *************************************************************
  SUBROUTINE       SOLVE(NJD       ,JND         ,COEFF      ,NLAKE     , &
     &                   ND        ,NL          ,NR         ,RHS       , &
     &                   U                          ) 
    INTEGER,INTENT(IN)::NLAKE ,ND ,NL ,NR
    INTEGER,DIMENSION(:),INTENT(IN)::NJD     ,JND      
    REAL(8), DIMENSION(:),INTENT(IN)::COEFF,RHS
    REAL(8),DIMENSION(:),INTENT(INOUT)::U
      
    !Local variables
    CHARACTER(LEN=5),PARAMETER::ThisProcedure='SOLVE'
    REAL(8),ALLOCATABLE,DIMENSION(:)::P       ,Z        ,WORK
    INTEGER, ALLOCATABLE, DIMENSION(:)::NJD_CRS, JND_CRS    			! INTERMEDIATE CRS FORMAT STORAGE ARRAYS
    REAL(8), ALLOCATABLE, DIMENSION(:)::COEFF_CRS

    INTEGER::NROW,ErrorCode,INFO(10),IDIM,IDIM_WORK,ITER,IDO,INDX,INDX_S,INDX_L,INDX1,INDX_DIV
    REAL(8)::DG8RES,DG9RES,DIAG


! *************************************************************
! ***** INITIALIZE THE VARIABLES
! *************************************************************
    IF ( ZERORESET .EQV. .TRUE. )  THEN 
       U = 0.D0
    END IF

    NROW        =NR+NLAKE+ND*NL


! *************************************************************
! ***** CHECK IF RELEVANT VARIABLES ARE DEFINED 
! *************************************************************
   IF (MSOLVE .EQ. -1)         CALL LogMessage('Matrix solver is not yet defined!',3,ThisProcedure)
   IF (MXITER .EQ. -1)         CALL LogMessage('Maximum number of iterations for matrix solver is not yet defined!',3,ThisProcedure)
   IF (STOPC_MATRIX .EQ. -1.0) CALL LogMessage('Matrix solver tolerance is not yet defined!',3,ThisProcedure)


! *************************************************************
! ***** MATRIX SOLUTION WITH SOR METHOD 
! *************************************************************
    IF (MSOLVE.EQ.1) THEN
       IF (RELAX .EQ. -1.0) THEN
          CALL LogMessage('Relaxation parameter for SOR method has not been defined!',3,ThisProcedure)
       END IF
      CALL CONVERTTOCRSFORMAT(COEFF, NJD, JND, COEFF_CRS, NJD_CRS, JND_CRS)
      CALL SOR(NROW,NJD_CRS,JND_CRS,RHS,COEFF_CRS,U)
      DEALLOCATE(COEFF_CRS, NJD_CRS, JND_CRS)
    END IF
      

! *************************************************************
! ***** MATRIX SOLUTION WITH PRECONDITIONED GMRES METHOD 
! *************************************************************
    IF (MSOLVE.EQ.2) THEN
      CALL CONVERTTOCRSFORMAT(COEFF, NJD, JND, COEFF_CRS, NJD_CRS, JND_CRS)
      CALL PGMRES(NROW,NJD_CRS,JND_CRS,RHS,COEFF_CRS,U)
      DEALLOCATE(COEFF_CRS, NJD_CRS, JND_CRS)
    END IF

  END SUBROUTINE SOLVE
      

!------------------------------------------------------------------------------------

! *************************************************************
! *************************************************************
! **** SOLVE THE MATRIX USING SUCCESSIVE OVER-RELAXATION METHOD
! *************************************************************
! *************************************************************
  SUBROUTINE SOR(NROW,NJD,JND,RHS,COEFF,U)
    INTEGER,INTENT(IN)::NROW
    INTEGER,DIMENSION(:),INTENT(IN)::NJD      ,JND      
    REAL(8),DIMENSION(:),INTENT(IN)::COEFF    ,RHS     
    REAL(8),DIMENSION(:),INTENT(INOUT)::U
    !Local variables
    
    CHARACTER(LEN=3),PARAMETER::ThisProcedure='SOR'
    INTEGER::IROW,INDX,INDX_S,INDX_L,NODEMAX
    REAL(8)::DIFF_L2,ADIFFMAX,U_INT,ACCUM,DIFF, DNRM2
    CHARACTER(LEN=128)::MessageArray(5)
    EXTERNAL DNRM2
! *************************************************************
! ***** INITIALIZE THE VARIABLES
! *************************************************************
   ITS = 0

! *************************************************************
! ***** SOLVE MQTRIX EQUATION USING SOR ITERATIVE METHOD
! *************************************************************
    DO
      ITS=ITS+1
      DIFF_L2=0.0
      ADIFFMAX=0.0

      DO IROW=1,NROW
        INDX_S=NJD(IROW)+1
        INDX_L=NJD(IROW+1)-1
        U_INT=U(IROW)
        ACCUM=0.0
        DO INDX=INDX_S,INDX_L
          IF (JND(INDX).EQ.0) CYCLE
          ACCUM=ACCUM+COEFF(INDX)*U(JND(INDX))
        END DO
        U(IROW)=(RHS(IROW)-ACCUM)/COEFF(INDX_S-1)
        DIFF=U(IROW)-U_INT
        DIFF_L2=DIFF_L2+DIFF**2.0
        U(IROW)=U_INT+DIFF*RELAX 
        IF (ABS(DIFF).GT.ADIFFMAX) THEN
          ADIFFMAX=ABS(DIFF)
          NODEMAX=IROW
        END IF
      END DO

! ***** CHECK CONVERGENCE OF MATRIX SOLVER
	
     IF (ITS.GT.MXITER .OR. SQRT(DIFF_L2).LE.STOPC_MATRIX) THEN 
!      WRITE(*,*) 'SOR Iteration: ', ITS, ' : Residual: ', SQRT(DIFF_L2), ' : ||x|| : ', DNRM2(NROW, U, 1)	     
	EXIT
      END IF
      
    END DO

! *************************************************************
! ***** STOP IF SOLUTION HAS NOT CONVERGED
! *************************************************************
    IF (SQRT(DIFF_L2).GT.STOPC_MATRIX) THEN
     MessageArray(1)='Convergence problem in the solution of equation system using' 
     MessageArray(2)='Succesive Over-Relaxation method.'     
     WRITE (MessageArray(3), '(A,I8)') 'Iteration =', ITS
     WRITE (MessageArray(4), '(A,I8)') 'Variable  =', NODEMAX
     WRITE (MessageArray(5),'(A,E12.3)') 'Difference=',ADIFFMAX
     CALL LogMessage(MessageArray(1:5),3,ThisProcedure)
   END IF

  END SUBROUTINE SOR


! *************************************************************
! *************************************************************
! **** SOLVE THE MATRIX USING (P)RECONDITIONED GMRES
! *************************************************************
! *************************************************************

 SUBROUTINE PGMRES(N,NJD,JND,RHS,COEFF,U)
!                 *** ILUT - Preconditioned GMRES ***                  *
!                                                                      *
!----------------------------------------------------------------------*
! This is a simple version of the ILUT preconditioned GMRES algorithm. *
! The ILUT preconditioner uses a dual strategy for dropping elements   *
! instead  of the usual level of-fill-in approach. See details in ILUT *
! subroutine documentation. PGMRES uses the L and U matrices generated *
! from the subroutine ILUT to precondition the GMRES algorithm.        *
! The preconditioning can be applied to the left or right. The 	       *
! stopping criterion utilized is based simply on reducing the 	       *
! residual norm by epsilon.   					       *
! This preconditioning is more reliable than ilu0 but requires more    *
! storage. It seems to be much less prone to difficulties related to   *
! strong nonsymmetries in the matrix. We recommend using a nonzero tol *
! (tol=.005 or .001 usually give good results) in ILUT. Use a large    *
! lfil whenever possible (e.g. lfil = 5 to 10). The higher lfil the    *
! more reliable the code is. Efficiency may also be much improved.     *
! Note that lfil=n and tol=0.0 in ILUT  will yield the same factors as *
! Gaussian elimination without pivoting.                               *


    INTEGER,INTENT(IN)::N
    INTEGER,DIMENSION(:),INTENT(IN)::NJD      ,JND      
    REAL(8),DIMENSION(:),INTENT(IN)::COEFF    ,RHS     
    REAL(8),DIMENSION(:),INTENT(INOUT)::U
    INTEGER, ALLOCATABLE, DIMENSION(:)::JLU, JU, IW, IPAR
    REAL(8), ALLOCATABLE, DIMENSION(:)::COEFFLU, FPAR, W
    REAL(8) DNRM2, DROPTOL, RES
    INTEGER LFIL, IM, IWK, IERR
    CHARACTER(LEN=6),PARAMETER::ThisProcedure='PGMRES'
    CHARACTER(LEN=128)::MessageArray(5)
   
    EXTERNAL DNRM2

    							! ILUT PRECONDITIONER PARAMETERS	
    LFIL 	= 	5				! THE LEVEL OF FILL-IN, TYPICALLY BETWEEN 5 AND 10
    DROPTOL	=	0.01				! THE DROP TOLERANCE, TYPICALLY SET TO < 1.0	
    IWK		= 	SIZE(COEFF) + LFIL + 1		! WORKSPACE SIZE REQUIRED BY THE PRECONDITIONER

    IM		=	20				! THE NUMBER OF ITERATIONS PERFORMED BEFORE RESTART OF GMRES(IM) 
    RES 	= 	0.0D0				! THE RESIDUAL ERROR USED TO ASSESS CONVERGENCE

    ALLOCATE(IPAR(16), FPAR(16))
							! PARMETERS OF THE SOLVER
    IPAR(1)	=	0				! INITIALIZE THE SOLVER
    IPAR(2) 	= 	1				! CHOOSE LEFT PRECONDITIONING
    IPAR(3) 	= 	2				! SPECIFY STOPPING CRITERIA BASED ON THE RESIDUAL	
    IPAR(4) 	= 	(N+3)*(IM+2) + (IM+1)*IM/2	! WORKSPACE SIZE NEED BY THE SOLVER
    IPAR(5) 	= 	IM
    IPAR(6) 	= 	MXITER

    FPAR(1) 	= 	STOPC_MATRIX/DNRM2(N,RHS,1)     ! TOLERANCE CONDITION
    FPAR(2) 	= 	2.22D-16			! MACHINE PRECISION
    
    ALLOCATE(COEFFLU(IWK), IW(3*N), JLU(IWK), JU(N), W(IPAR(4)))   

!   CALL SPARSKIT IMPLEMENTATION OF ILUT IN THE FILE pgmres.f
    CALL ILUT(N, COEFF, JND, NJD, LFIL, DROPTOL, COEFFLU, JLU, JU, IWK, W, IW, IERR)

    ITS         = 	0				! INITIALIZE THE NUMBER OF ITERATIONS
   

10  CALL GMRES(N, RHS, U, IPAR, FPAR, W)		! CALL SPARSKIT IMPLEMENTATION OF GMRES(M) IN THE FILE iters.f


    IF (IPAR(7).NE.ITS) THEN
         ITS = IPAR(7)
    ENDIF
    RES = FPAR(5)
!
 
    IF (IPAR(1).EQ.1) THEN
         CALL AMUX(N, W(IPAR(8)), W(IPAR(9)), COEFF, JND, NJD)	!MATRIX-VECTOR MULTIPLICATION
         GOTO 10
    ELSE IF (IPAR(1).EQ.3 .OR. IPAR(1).EQ.5) THEN
         CALL LUSOL(N,W(IPAR(8)),W(IPAR(9)),COEFFLU,JLU,JU)	!LU SOLVE
         GOTO 10 
    ELSE IF (IPAR(1).LT.0) THEN

         IF (IPAR(1).EQ.-1) THEN
     	   MessageArray(1)='Convergence problem in the solution of equation system using' 
           MessageArray(2)='PGMRES(M).'     
     	   WRITE(MessageArray(3),'(A,I8)') 'Iteration =', ITS
     	   WRITE (MessageArray(4),'(A,E12.3)') 'Residual =', RES
           CALL LogMessage(MessageArray(1:4),3,ThisProcedure)

         ELSE IF (IPAR(1).EQ.-2) THEN
           MessageArray(1)='ITERATIVE SOLVER WAS NOT GIVEN ENOUGH WORK SPACE.'
            WRITE (MessageArray(2),'(A,I12,A)') 'THE WORK SPACE SHOULD AT LEAST HAVE ', IPAR(4),' ELEMENTS.'
            CALL LogMessage(MessageArray(1:2),3,ThisProcedure)
         ELSE IF (IPAR(1).EQ.-3) THEN
           CALL LogMessage('ITERATIVE SOLVER IS FACING A BREAK-DOWN.',3,ThisProcedure)
         ELSE
     	   WRITE(MessageArray(1),'(A,I8)') 'ITERATIVE SOLVER TERMINATED. CODE =', IPAR(1) 
           CALL LogMessage(MessageArray(1),3,ThisProcedure)
         ENDIF
    ELSE IF (IPAR(1).EQ.0) THEN
 !      WRITE(*,*) 'GMRES Iteration: ', ITS, ' : Residual: ', RES, ' : ||x|| : ', DNRM2(N, U, 1)

    ENDIF
    
    DEALLOCATE(IPAR, FPAR)
    DEALLOCATE(COEFFLU, IW, JLU, JU, W)

END SUBROUTINE PGMRES

! *************************************************************
! *************************************************************
! **** CONVERT IWFM STORAGE FORMAT TO        	     **********
! **** COMPRESSED ROW STORAGE (CRS) FORMAT    	     **********
! *************************************************************
! *************************************************************


SUBROUTINE CONVERTTOCRSFORMAT(COEFF, NJD, JND, COEFF_CRS, NJD_CRS, JND_CRS)
    IMPLICIT NONE
    INTEGER, DIMENSION(:), INTENT(IN)::NJD
    INTEGER, DIMENSION(:), INTENT(IN)::JND
    REAL(8), DIMENSION(:), INTENT(IN)::COEFF

    INTEGER, ALLOCATABLE, DIMENSION(:), INTENT(OUT)::NJD_CRS, JND_CRS    		! INTERMEDIATE STORAGE ARRAYS
    REAL(8), ALLOCATABLE, DIMENSION(:), INTENT(OUT)::COEFF_CRS

    INTEGER I, K, IM, IshftIDX , N, NJ

    N = SIZE(NJD) - 1
    NJ= SIZE(COEFF)

    IshftIDX = 0
      
    DO I = 1, N   
       DO K = NJD(I),(NJD(I+1)-1)
         IF (COEFF(K).EQ.0.D0 .OR. JND(K).EQ.0) IshftIDX = IshftIDX + 1       
     ENDDO    
    ENDDO
      
    NJ = NJ-IshftIDX
    ALLOCATE(COEFF_CRS(NJ), NJD_CRS(N+1), JND_CRS(NJ))
      
    IM = 1
    NJD_CRS(1) = 1
    IshftIDX = 0

    DI: DO I = 1, N   
     DK: DO K = NJD(I),(NJD(I+1)-1)
         IF (COEFF(K).EQ.0.D0 .OR. JND(K).EQ.0) THEN
           IshftIDX = IshftIDX + 1
         ELSE
           COEFF_CRS(IM) = COEFF(K)
           JND_CRS(IM) = JND(K)
           IM = IM + 1
         ENDIF
        ENDDO DK
       NJD_CRS(I+1) = NJD(I+1) - IshftIDX 
    ENDDO DI
 

END SUBROUTINE CONVERTTOCRSFORMAT

END MODULE 
