#include "AUTODIFF_OPTIONS.h"
#ifdef ALLOW_CTRL
# include "CTRL_OPTIONS.h"
#endif

C--   File cg2d_mad.F: Code for hand-written (manual) adjoint of cg2d
C--    Contents
C--   CG2D_MAD        - computes the adjoint of cg2d
C--   CG2D_STORE      - saves input field to or restores from a common block

CBOP
C     !ROUTINE: CG2D_MAD
C     !INTERFACE:
      SUBROUTINE CG2D_MAD(
     U                    cg2d_b_ad,
     U                    cg2d_x_ad,
     U                    numIters, nIterMin,
     I                    myThid )
C     !DESCRIPTION: \bv
C     *==========================================================*
C     | SUBROUTINE CG2D_MAD
C     | o Manual ADjoint version of CG2D, the two-dimensional
C     |   grid problem conjugate-gradient inverter.
C     *==========================================================*
C     | This routine is called from solve_for_pressure_ad that
C     | is generated by TAF. We want a self-adjoint cg2d
C     | to avoid too many complications, so we provide flow
C     | directives in cg2d.flow and add the missing AD
C     | contributions manually (hence the name mad=Manuad ADjoint).
C     |
C     | This routine calls cg2d with reversed order of argumgent
C     | cg2d_x_ad and cg2d_b_ad, and sets dependendcies to zero
C     | that a non-self-adjoint cg2d would have produced.
C     | In addition, when runtime parameter cg2dFullAdjoint = T,
C     | it computes sensitivities to matrix coefficients. This
C     | makes it the full adjoint of cg2d, only if cg2d converges
C     | completely. Otherwise it is an excellent approximation.
C     *==========================================================*
C     \ev

C     !USES:
      IMPLICIT NONE
C     === Global data ===
#include "SIZE.h"
#include "EEPARAMS.h"
#include "PARAMS.h"
#if ( defined NONLIN_FRSURF || defined ALLOW_DEPTH_CONTROL )
# include "CG2D.h"
#endif
#ifdef ALLOW_AUTODIFF_TAMC
# include "AUTODIFF_PARAMS.h"
#endif

C     !INPUT/OUTPUT PARAMETERS:
C     === Routine arguments ===
C     cg2d_b_ad :: Adjoint of source term (on input = 0)
C     cg2d_x_ad :: Adjoint of the solution
C     firstResidual :: the initial residual before any iterations
C     minResidualSq :: the lowest residual reached (squared)
C     lastResidual  :: the actual residual reached
C     numIters  :: Inp: the maximum number of iterations allowed
C                  Out: the actual number of iterations used
C     nIterMin  :: Inp: decide to store (if >=0) or not (if <0) lowest res. sol.
C                  Out: iteration number corresponding to lowest residual
C     myThid    :: Thread on which I am working.
      _RL  cg2d_b_ad(1-OLx:sNx+OLx,1-OLy:sNy+OLy,nSx,nSy)
      _RL  cg2d_x_ad(1-OLx:sNx+OLx,1-OLy:sNy+OLy,nSx,nSy)
c     _RL  firstResidual, minResidualSq, lastResidual
      INTEGER numIters
      INTEGER nIterMin
      INTEGER myThid

#ifdef ALLOW_AUTODIFF_TAMC
#if ( defined NONLIN_FRSURF || defined ALLOW_DEPTH_CONTROL )
C     directly imported from TAF-generated code, make sure that they are
C     consistent with what is found in S/R update_cg2d_ad
      _RS aW2d_ad(1-OLx:sNx+OLx,1-OLy:sNy+OLy,nSx,nSy)
      _RS aS2d_ad(1-OLx:sNx+OLx,1-OLy:sNy+OLy,nSx,nSy)
      _RS aC2d_ad(1-OLx:sNx+OLx,1-OLy:sNy+OLy,nSx,nSy)
      _RS pW_ad  (1-OLx:sNx+OLx,1-OLy:sNy+OLy,nSx,nSy)
      _RS pS_ad  (1-OLx:sNx+OLx,1-OLy:sNy+OLy,nSx,nSy)
      _RS pC_ad  (1-OLx:sNx+OLx,1-OLy:sNy+OLy,nSx,nSy)
      COMMON /ADCG2D_I_RS/
     &     aW2d_ad, aS2d_ad, aC2d_ad, pW_ad, pS_ad, pC_ad
#endif /* NONLIN_FRSURF or ALLOW_DEPTH_CONTROL */

C     !LOCAL VARIABLES:
C     === Local variables ===
C     bi,bj  :: tile indices
C     i,j    :: Loop counters
C     cg2d_x :: local copy of cg2d_x to be restored from "tape"
      INTEGER i,j,bi,bj
      _RL  firstResidual, minResidualSq, lastResidual
#if ( defined NONLIN_FRSURF || defined ALLOW_DEPTH_CONTROL )
      INTEGER numItersFwd, nIterMinFwd
      _RL  recip_cg2dNorm
      _RL  cg2d_x(1-OLx:sNx+OLx,1-OLy:sNy+OLy,nSx,nSy)
#endif
      CHARACTER*(MAX_LEN_MBUF) msgBuf
CEOP

C--   mark output of cg2d as being called from here
      IF ( debugLevel .GE. debLevZero ) THEN
       _BEGIN_MASTER( myThid )
       WRITE(standardMessageUnit,'(A)')
     &      ' Calling cg2d from S/R CG2D_MAD'
       _END_MASTER( myThid )
      ENDIF

#if ( defined NONLIN_FRSURF || defined ALLOW_DEPTH_CONTROL )
C     make local copies because we need them later
      numItersFwd = numIters
      nIterMinFwd = nIterMin
#endif

#ifdef ALLOW_SRCG
      IF ( useSRCGSolver ) THEN
C--   Call the single reduce CG solver
       CALL CG2D_SR(
     U      cg2d_x_ad, cg2d_b_ad,
     O      firstResidual, minResidualSq, lastResidual,
     U      numIters, nIterMin,
     I      myThid )
      ELSE
#else
      IF (.TRUE.) THEN
#endif /* ALLOW_SRCG */
C--   Call the default CG solver
       CALL CG2D(
     U      cg2d_x_ad, cg2d_b_ad,
     O      firstResidual, minResidualSq, lastResidual,
     U      numIters, nIterMin,
     I      myThid )
      ENDIF

CML      CALL ADEXCH_XY_RL( cg2d_b_ad, myThid )
CML       _EXCH_XY_RL( cg2d_b_ad, myThid )

C     cg2d_x_ad is reset to zero, because we assume that the solution of
C     the self-adjoint solver does not depend in the intial conditions
C     so that no sensitivities should be passed back to the caller.
      DO bj=myByLo(myThid),myByHi(myThid)
       DO bi=myBxLo(myThid),myBxHi(myThid)
        DO j=1-OLy,sNy+OLy
         DO i=1-OLx,sNx+OLx
          cg2d_x_ad(i,j,bi,bj) = 0. _d 0
         ENDDO
        ENDDO
       ENDDO
      ENDDO

C--   dump CG2D output
C     Since there is no clean way in importing myTime into this routine,
C     we just call it every timestep but only for a higher debugLevel
      IF ( debugLevel .GE. debLevC .OR. printResidualFreq.GE.0 ) THEN
       _BEGIN_MASTER( myThid )
       WRITE(msgBuf,'(A30,1PE23.14)')
     &       'CG2D_MAD: cg2d_init_res =', firstResidual
       CALL PRINT_MESSAGE(msgBuf,standardMessageUnit,SQUEEZE_RIGHT,1)
       WRITE(msgBuf,'(A37,2I8)')
     &       'CG2D_MAD: cg2d_iters(min,last) =', nIterMin, numIters
       CALL PRINT_MESSAGE(msgBuf,standardMessageUnit,SQUEEZE_RIGHT,1)
       IF ( minResidualSq.GE.0. ) THEN
        minResidualSq = SQRT(minResidualSq)
        WRITE(msgBuf,'(A30,1PE23.14)')
     &        'CG2D_MAD: cg2d_min_res  =', minResidualSq
        CALL PRINT_MESSAGE(msgBuf,standardMessageUnit,SQUEEZE_RIGHT,1)
       ENDIF
       WRITE(msgBuf,'(A30,1PE23.14)')
     &       'CG2D_MAD: cg2d_last_res =',lastResidual
       CALL PRINT_MESSAGE(msgBuf,standardMessageUnit,SQUEEZE_RIGHT,1)
       _END_MASTER( myThid )
      ENDIF

#if ( defined NONLIN_FRSURF || defined ALLOW_DEPTH_CONTROL )
      recip_cg2dNorm = 1.0 _d 0
      IF ( cg2dNorm .NE. 0. _d 0 ) THEN
       recip_cg2dNorm = 1. _d 0 / cg2dNorm
      ENDIF

      IF ( cg2dFullAdjoint ) THEN
C     restore cg2d_x
       CALL CG2D_STORE( cg2d_x, .FALSE., myThid )

C     Compute derivative of coefficient matrix.
       DO bj=myByLo(myThid),myByHi(myThid)
        DO bi=myBxLo(myThid),myBxHi(myThid)
         DO j=1,sNy
          DO i=1,sNx
           aC2d_ad(i,j,bi,bj)   = aC2d_ad(i,j,bi,bj)
     &          - cg2d_b_ad(i,j,bi,bj) * cg2d_x(i,j,bi,bj)
     &          * recip_cg2dNorm
           aW2d_ad(i,  j,bi,bj) = aW2d_ad(i,j,bi,bj)
     &          - cg2d_b_ad(i,j,bi,bj) * cg2d_x(i-1,j,bi,bj)
     &          * recip_cg2dNorm
           aW2d_ad(i+1,j,bi,bj) = aW2d_ad(i+1,j,bi,bj)
     &          - cg2d_b_ad(i,j,bi,bj) * cg2d_x(i+1,j,bi,bj)
     &          * recip_cg2dNorm
           aS2d_ad(i,j,  bi,bj) = aS2d_ad(i,j,bi,bj)
     &          - cg2d_b_ad(i,j,bi,bj) * cg2d_x(i,j-1,bi,bj)
     &          * recip_cg2dNorm
           aS2d_ad(i,j+1,bi,bj) = aS2d_ad(i,j+1,bi,bj)
     &          - cg2d_b_ad(i,j,bi,bj) * cg2d_x(i,j+1,bi,bj)
     &          * recip_cg2dNorm
C     The sensitivities of preconditioner are something like
C          pC_ad = pC_ad + cg2d_z_ad * cg2d_r
C     and similar (cg2d_z = preconditioner * cg2d_r). That is, when the
C     solver has converged, cg2d_r -> 0 and pC_ad -> 0, unless cg2d_z_ad
C     is large, which would mean an instability. Therefore we set the
C     preconditioner adjoint to zero. This also reflects the fact, that
C     in an ideal world, the preconditioner does not have any effect on
C     the solution, but only on the rate of convergence.
           pW_ad(i,j,bi,bj) = 0. _d 0
           pS_ad(i,j,bi,bj) = 0. _d 0
           pC_ad(i,j,bi,bj) = 0. _d 0
          ENDDO
         ENDDO
        ENDDO
       ENDDO
      ELSE
C     Initialise AD fields to zero for an adjoint neglecting the
C     contributions from the matrix coefficients (default for now).
       DO bj=myByLo(myThid),myByHi(myThid)
        DO bi=myBxLo(myThid),myBxHi(myThid)
         DO j=1-OLy,sNy+OLy
          DO i=1-OLx,sNx+OLx
           aW2d_ad(i,j,bi,bj) = 0. _d 0
           aS2d_ad(i,j,bi,bj) = 0. _d 0
           aC2d_ad(i,j,bi,bj) = 0. _d 0
           pW_ad  (i,j,bi,bj) = 0. _d 0
           pS_ad  (i,j,bi,bj) = 0. _d 0
           pC_ad  (i,j,bi,bj) = 0. _d 0
          ENDDO
         ENDDO
        ENDDO
       ENDDO
      ENDIF
#endif /* NONLIN_FRSURF or ALLOW_DEPTH_CONTROL */
#endif /* ALLOW_AUTODIFF_TAMC */

      RETURN
      END

C--------1---------2---------3---------4---------5---------6---------7--

CBOP
C     !ROUTINE: CG2D_STORE
C     !INTERFACE:
      SUBROUTINE CG2D_STORE(
     U                      cg2d_x,
     I                      doStore,
     I                      myThid )
C     !DESCRIPTION: \bv
C     *==========================================================*
C     | SUBROUTINE CG2D_STORE
C     | o Hand written tape routine
C     *==========================================================*
C     | This routine is called from solve_for_pressure to
C     | to force TAF to store cg2d_x
C     *==========================================================*
C     \ev

C     !USES:
      IMPLICIT NONE
C     === Global data ===
#include "SIZE.h"
#include "EEPARAMS.h"
#include "PARAMS.h"
#include "AUTODIFF_PARAMS.h"
#ifdef ALLOW_AUTODIFF_TAMC
# include "tamc.h"
#endif

C     !INPUT/OUTPUT PARAMETERS:
C     === Routine arguments ===
C     cg2d_x    :: The solution to be stored and restored
C     doStore   :: Store or restore based on the logical
C     myThid    :: Thread on which I am working.
      _RL  cg2d_x(1-OLx:sNx+OLx,1-OLy:sNy+OLy,nSx,nSy)
      LOGICAL doStore
      INTEGER myThid

#ifdef ALLOW_AUTODIFF_TAMC
#if ( defined NONLIN_FRSURF || defined ALLOW_DEPTH_CONTROL )
      _RL cg2d_tape(1-OLx:sNx+OLx,1-OLy:sNy+OLy,nSx,nSy,nchklev_1)
      COMMON /CG2D_TAPE_RL/ cg2d_tape

C     !LOCAL VARIABLES:
C     === Local variables ===
C     bi,bj  :: tile indices
C     i,j    :: Loop counters
      INTEGER i,j,bi,bj
CEOP

C--   mark output of cg2d as being called from here
      IF ( debugLevel .GE. debLevZero ) THEN
       _BEGIN_MASTER( myThid )
       WRITE(standardMessageUnit,'(A,L1)')
     &      ' Calling CG2D_STORE with doStore = ', doStore
       WRITE(standardMessageUnit,'(A,I6)')
     &      ' Calling CG2D_STORE with ikey_dynamics=', ikey_dynamics
       _END_MASTER( myThid )
      ENDIF

      IF ( doStore ) THEN
C     Save input cg2d_x to local common block
       DO bj = myByLo(myThid), myByHi(myThid)
        DO bi = myBxLo(myThid), myBxHi(myThid)
         DO j=1-OLy,sNy+OLy
          DO i=1-OLx,sNx+OLx
           cg2d_tape(i,j,bi,bj,ikey_dynamics) = cg2d_x(i,j,bi,bj)
          ENDDO
         ENDDO
        ENDDO
       ENDDO
      ELSE
C     Restore output cg2d_x from local common block
       DO bj = myByLo(myThid), myByHi(myThid)
        DO bi = myBxLo(myThid), myBxHi(myThid)
         DO j=1-OLy,sNy+OLy
          DO i=1-OLx,sNx+OLx
           cg2d_x(i,j,bi,bj) = cg2d_tape(i,j,bi,bj,ikey_dynamics)
          ENDDO
         ENDDO
        ENDDO
       ENDDO
      ENDIF
#endif /* NONLIN_FRSURF or ALLOW_DEPTH_CONTROL */
#endif /* ALLOW_AUTODIFF_TAMC */

      RETURN
      END
