!+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! File: ct29set.f90
! Auxiliary subroutines - set up the derivative structures and bounds for the
! control problem.
!
! 28 Dec 2008: Updated for v4 of snctrl.
! 09 Feb 2010: v5.
!+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

module ct29set
  use precision,    only : ip, rp
  use ctWorkModule, only : ctWork, ctGrid
  use ct15usr,      only : usrodeS, usrodeA, usrodeD, usralgS, usralgA, usralgD
  use ct28derv,     only : odeSetupTR, algSetupTR, odeSetupHS, algSetupHS
  implicit none

  private
  public  :: c2setup ! contains c2deriv, c2algbds, c2vbdsA

contains

  !+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

  subroutine c2setup ( iDisc, iIntf, nPhs, nY, nU, nP, nC, objL, &
                       ytype, ctype, nnY, nlY, nnC, nlC, neJ, neG, x, hs, &
                       grid, work, odecon, algcon, vbounds, infBnd, &
                       cu, lencu, iu, leniu, ru, lenru, iw, leniw )
    integer(ip), intent(in)    :: iDisc, iIntf, nPhs, nY, nU, nP, nC, &
                                  objL(nPhs), ytype(nY,nPhs), ctype(nC,nPhs), &
                                  nnY(nPhs), nlY(nPhs), nnC(nPhs), nlC(nPhs), &
                                  neJ(nPhs), neG(nPhs), lencu, leniu, lenru, leniw
    real(rp),    intent(in)    :: infBnd
    integer(ip), intent(inout) :: iu(leniu), iw(leniw)
    real(rp),    intent(inout) :: ru(lenru)
    character(8),intent(inout) :: cu(lencu)
    integer(ip), intent(inout), pointer :: hs(:)
    real(rp),    intent(inout), pointer :: x(:)
    external             :: odecon, algcon, vbounds
    type(ctGrid), target :: grid
    type(ctWork)         :: work

    !===========================================================================
    ! Grid has already been setup.
    ! Sets the workspace, bounds, derivative structures.
    !
    ! 09 Feb 2010: Current version.
    !===========================================================================
    integer(ip) :: nInt, nNodes, n, nDisCon, m, lenY, lenC, nnCon, nlCon, &
                   ldiff, ndiff, ndS, ndE, indS, indE, jt, p, s
    integer(ip), pointer :: ndPtr(:), intPtr(:), phs(:)
    real(rp),    pointer :: step(:), bl(:), bu(:), rc(:), pi(:), px(:)
    real(rp),    allocatable :: clbds(:,:), cubds(:,:)

    nInt   =  grid%nInt
    nNodes =  grid%nNodes
    ndPtr  => grid%ndPtr
    intPtr => grid%intPtr
    step   => grid%step

    allocate ( clbds(nC,nPhs), cubds(nC,nPhs) )

    !---------------------------------------------------------------------------
    ! Compute values for the workspace.
    !   nnCon is the number of nonlinear constraints.
    !   m    is the number of constraints.
    !   n    is the number of variables.
    !   lenY is the (max) size of the Jacobian for the state equations.
    !   lenC is the (max) size of the Jacobian for the algebraic
    !        constraints.

    if ( iDisc == 0 ) then
       ! Variables
       n = (nY+nU)*nNodes + nP

       ! Constraints
       nDisCon = nInt*nY + &  ! discretized state eqns
                 nNodes*nC    ! discretized algebraic constraints

       ! Constraints/Objective
       m = 1 + nDisCon + nY*(nPhs-1)

       ! Derivatives
       lenY = 2*(nY+nU) * nInt*nY + &   ! wrt states/controls
              nP * nInt*nY + &          ! wrt parameters
              2 * (nPhs-1)*nY + &       ! phase continuity
              nPhs                      ! objectives
       lenC = (nY+nU) * nNodes*nC + &   ! wrt states/controls
              nP * nNodes*nC            ! wrt parameters

    else
       ! Variables
       n = (nY+nU)*nNodes + nP

       ! Constraints
       nDisCon = nInt*nU + &    ! linear controls
                 2*nInt*nY + &  ! discretized state eqns
                 (nInt+1)*nC    ! discretized algebraic constraints

       ! Constraints/Objective
       m = 1 + nDisCon + nY*(nPhs-1)

       ! Derivatives
       lenY = 3*(nY+nU) * 2*nInt*nY + &   ! wrt states/controls
              nP * 2*nInt*nY + &          ! wrt parameters
              2 * (nPhs-1)*nY + &         ! phase continuity
              nPhs + &                    ! objectives
              3 * nInt*nU                 ! linear controls

       lenC = (nY+nU) * (nInt+1)*nC + &   ! wrt states/controls
              nP * (nInt+1)*nC            ! wrt parameters

    end if

    ! Nonlinear constraints
    !    Count the number of nonlinear constraints based on user's info on
    !    nonlinear/linear states.
    nnCon = 0
    nlCon = 0

    do p = 1, nPhs
       ldiff = nlY(p)
       ndiff = nnY(p)
       if ( iDisc == 1 ) then
          ldiff = 2*nlY(p) + nU
          ndiff = 2*nnY(p)
       end if

       jt = intPtr(p+1) - intPtr(p)
       nnCon = nnCon + jt*(nnC(p) + ndiff)
       nlCon = nlCon + jt*(nlC(p) + ldiff)
    end do

    nnCon = nnCon + nnC(nPhs)
    nlCon = nlCon + nlC(nPhs)
    nlCon = nlCon + (nPhs-1)*nY

    if (nnCon + nlCon + 1 /= m) then
       call snPRNT (13,'snctrl: INTERNAL setup error.', iw, leniw )
       return
    end if

    ! Set the workspace:
    s = 0
    if ( associated(x) ) then
       s   = size(x)
       px  => x
       phs => hs
       nullify ( x, hs )
    end if

    allocate ( x(n+m), hs(n+m) )
    allocate ( work%bl(n+m), work%bu(n+m), work%rc(n+m), work%pi(m) )
    bl => work%bl
    bu => work%bu
    rc => work%rc
    pi => work%pi

    x  = 0.0
    hs = 0
    bl = -infBnd
    bu =  infBnd
    pi =  0.0
    rc =  0.0

    ! Initialize bounds on discretized constraints.
    bl(1+n:m+n) = 0.0
    bu(1+n:m+n) = 0.0

    ! Set the objective (free) bounds.
    bl(m+n) = -infBnd
    bu(m+n) =  infBnd

    ! Initialize workspace.
    work%n     = n
    work%m     = m
    work%lenJ  = 0
    work%nnCon = nnCon
    work%lenY  = lenY
    work%lenC  = lenC


    !---------------------------------------------------------------------------
    ! Variable bounds:
    if ( iIntF == 0 .or. iIntF == 1 ) then   ! S/D versions
       indS = 1
       do p = 1, nPhs
          ndS    = ndPtr(p)
          ndE    = ndPtr(p+1)
          nNodes = ndE - ndS + 1

          indE = indS-1 + nNodes*(nY+nU)

          call vbounds ( p, nPhs, nY, nU, nP, nC, nNodes, &
                         bl(indS:indE), bu(indS:indE), x(indS:indE), &
                         bl(n-nP+1:n), bu(n-nP+1:n), &
                         x(n-nP+1:n), clbds(:,p), cubds(:,p) )
          indS = indE + 1
       end do

    else     ! A version
       call c2vbdsA ( nPhs, nY, nU, nP, nC, ndPtr, n, &
                       bl(1:n), bu(1:n), x(1:n), clbds, cubds, vbounds )

    end if

    ! Set up constraint bounds.
    if ( nC > 0 ) &
         call c2algbds ( iDisc, nPhs, nC, nU, ctype, nnCon, intPtr, &
                          clbds, cubds, m, bl(1+n:m+n), bu(1+n:m+n) )


    ! Set up constraint derivative structures (work%lenJ, indJ, locJ, Jcol).
    call c2deriv ( iDisc, iIntf, nPhs, nY, nU, nP, nC, objL, &
                    ytype, ctype, nnY, nlY, nnC, nlC, neJ, neG, &
                    nnCon, ndPtr, step, nInt, work, odecon, algcon, &
                    cu, lencu, iu, leniu, ru, lenru )

    ! Copy input x, hs.
    if ( s > 0 ) then
       x(1:s)  = px(1:s)
       hs(1:s) = phs(1:s)
       deallocate ( px, phs )
    end if

    return

  contains

    !+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

    subroutine c2vbdsA ( nPhs, nY, nU, nP, nC, ndPtr, n, bl, bu, x, &
                         clb, cub, vbounds )
      integer(ip), intent(in)  :: nPhs, nY, nU, nP, nC, n, ndPtr(nPhs+1)
      real(rp),    intent(out) :: bl(n), bu(n), x(n), clb(nC,nPhs), cub(nC,nPhs)
      external     :: vbounds

      !=========================================================================
      ! Set up bounds for states, controls, parameters for the A version.
      ! Return algebraic constraints.
      !
      ! 31 Jan 2009: First version of c2vbdsA.
      ! 06 Feb 2010: Current version.
      !=========================================================================
      integer(ip) :: j, p, indv, nVar, ndS, ndE, nNodes
      real(rp)    :: y0low(nY), y0upp(nY), yflow(nY), yfupp(nY), &
                     ylow(nY),  yupp(nY),  ulow(nU),  uupp(nU)

      nVar = nY+nU
      indv = 0

      do p = 1, nPhs
         ndS    = ndPtr(p)
         ndE    = ndPtr(p+1)
         nNodes = ndE - ndS + 1

         ! Get user-defined bounds
         call vbounds ( p, nPhs, nY, nU, nP, nC, y0low, y0upp, &
                        yflow, yfupp, ylow, yupp, ulow, uupp, &
                        bl(n-nP+1:n), bu(n-nP+1:n), x(n-nP+1:n), &
                        clb(:,p), cub(:,p) )

         j = 1
         bl(1+indv:nY+indv) = y0low
         bu(1+indv:nY+indv) = y0upp
         bl(1+indv+nY:nU+indv+nY) = ulow
         bu(1+indv+nY:nU+indv+nY) = uupp
         indv = indv + nVar

         do j = 2, nNodes-1
            bl(1+indv:nY+indv) = ylow
            bu(1+indv:nY+indv) = yupp

            bl(1+indv+nY:nU+indv+nY) = ulow
            bu(1+indv+nY:nU+indv+nY) = uupp
            indv = indv + nVar
         end do

         j = nNodes
         bl(1+indv:nY+indv) = yflow
         bu(1+indv:nY+indv) = yfupp
         bl(1+indv+nY:nU+indv+nY) = ulow
         bu(1+indv+nY:nU+indv+nY) = uupp
         indv = indv + nVar

      end do

    end subroutine c2vbdsA

    !+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

    subroutine c2deriv ( iDisc, iIntf, nPhs, nY, nU, nP, nC, objL, &
                         ytype, ctype, nnY, nlY, nnC, nlC, neJ, neG, &
                         nnCon, ndPtr, step, s, work, odecon, algcon, &
                         cu, lencu, iu, leniu, ru, lenru )
      integer(ip),  intent(in) :: iDisc, iIntf, nPhs, nY, nU, nP, nC, s, &
                                  objL(nPhs), ytype(nY,nPhs), ctype(nC,nPhs), &
                                  nnY(nPhs), nlY(nPhs), nnC(nPhs), nlC(nPhs), &
                                  neJ(nPhs), neG(nPhs), nnCon, ndPtr(nPhs+1), &
                                  lencu, leniu, lenru, iu(leniu)
      real(rp),     intent(in) :: step(s), ru(lenru)
      character(8), intent(in) :: cu(lencu)
      external     :: odecon, algcon
      type(ctWork) :: work

      !=========================================================================
      ! Structures are set up in odeSetup and algSetup.  This routine calls them
      ! and combines the results.
      !
      ! 29 Dec 2008: Updated conSetupTR for v4 of snctrl.
      ! 16 Jan 2009: Implemented sparse data structures for user Jacobian.
      ! 19 Jan 2009: Combined TR/HS versions.
      ! 06 Feb 2010: Current version of c2deriv.
      !=========================================================================
      integer(ip) :: m, n, lenY, lenC, count, neY, neC, nJ, k, j
      integer(ip), allocatable :: ncindJ(:), lcindJ(:), nclocJ(:), lclocJ(:), &
                                  nfindJ(:), lfindJ(:), nflocJ(:), lflocJ(:)
      real(rp),    allocatable :: lfJcol(:), lcJcol(:)
      integer(ip), pointer :: locJ(:), indJ(:)
      real(rp),    pointer :: Jcol(:)


      n    = work%n
      m    = work%m
      lenY = work%lenY

      !-------------------------------------------------------------------
      ! ODE constraints
      !-------------------------------------------------------------------
      allocate ( nfindJ(lenY), nflocJ(n+1) )
      allocate ( lfindJ(lenY), lflocJ(n+1), lfJcol(lenY) )

      nfindJ = 0
      nflocJ = 0

      lfindJ = 0
      lflocJ = 0
      lfJcol = 0

      if ( iIntf == 0 ) then   ! S version
         if ( iDisc == 0 ) then
            call odeSetupTR ( nPhs, nY, nU, nP, objL, ytype, nnY, nlY, nnC, nlC, &
                              neJ, nnCon, ndPtr, step, s, n, m, &
                              lenY, nfindJ, nflocJ, lfJcol, lfindJ, lflocJ, &
                              usrodeS, odecon, cu, lencu, iu, leniu, ru, lenru )

         else
            call odeSetupHS ( nPhs, nY, nU, nP, objL, ytype, nnY, nlY, nnC, nlC, &
                              neJ, nnCon, ndPtr, step, s, n, m, &
                              lenY, nfindJ, nflocJ, lfJcol, lfindJ, lflocJ, &
                              usrodeS, odecon, cu, lencu, iu, leniu, ru, lenru )

         end if
      else if ( iIntf == 1 ) then   ! D version
         if ( iDisc == 0 ) then
            call odeSetupTR ( nPhs, nY, nU, nP, objL, ytype, nnY, nlY, nnC, nlC, &
                              neJ, nnCon, ndPtr, step, s, n, m, &
                              lenY, nfindJ, nflocJ, lfJcol, lfindJ, lflocJ, &
                              usrodeD, odecon, cu, lencu, iu, leniu, ru, lenru )

         else
            call odeSetupHS ( nPhs, nY, nU, nP, objL, ytype, nnY, nlY, nnC, nlC, &
                              neJ, nnCon, ndPtr, step, s, n, m, &
                              lenY, nfindJ, nflocJ, lfJcol, lfindJ, lflocJ, &
                              usrodeD, odecon, cu, lencu, iu, leniu, ru, lenru )

         end if

      else if ( iIntf == 2 ) then   ! A version
         if ( iDisc == 0 ) then
            call odeSetupTR ( nPhs, nY, nU, nP, objL, ytype, nnY, nlY, nnC, nlC, &
                              neJ, nnCon, ndPtr, step, s, n, m, &
                              lenY, nfindJ, nflocJ, lfJcol, lfindJ, lflocJ, &
                              usrodeA, odecon, cu, lencu, iu, leniu, ru, lenru )

         else
            call odeSetupHS ( nPhs, nY, nU, nP, objL, ytype, nnY, nlY, nnC, nlC, &
                              neJ, nnCon, ndPtr, step, s, n, m, &
                              lenY, nfindJ, nflocJ, lfJcol, lfindJ, lflocJ, &
                              usrodeA, odecon, cu, lencu, iu, leniu, ru, lenru )

         end if
      end if

      if ( nC == 0 ) then
         neY = nflocJ(n+1) + lflocJ(n+1) - 2

         work%lenJ = neY
         allocate ( work%locJ(n+1), work%Jcol(neY), work%indJ(neY) )

         locJ => work%locJ
         Jcol => work%Jcol
         indJ => work%indJ

         count = 1
         do k = 1, n
            locJ(k) = count

            do j = nflocJ(k), nflocJ(k+1)-1
               indJ(count) = nfindJ(j)
               Jcol(count) = 0.0
               count       = count + 1
            end do

            do j = lflocJ(k), lflocJ(k+1)-1
               indJ(count) = lfindJ(j)
               Jcol(count) = lfJcol(j)
               count       = count + 1
            end do
         end do
         locJ(n+1) = count

         deallocate ( nfindJ, nflocJ, lfindJ, lflocJ, lfJcol )
         return
      end if


      !-------------------------------------------------------------------
      ! Algebraic constraints
      !    Same as above but with the algebraic constraints.
      !-------------------------------------------------------------------
      lenC = work%lenC
      allocate ( ncindJ(lenC), nclocJ(n+1) )
      allocate ( lcindJ(lenC), lclocJ(n+1), lcJcol(lenC) )

      ncindJ = 0
      nclocJ = 0

      lcindJ = 0
      lclocJ = 0
      lcJcol = 0

      if ( iIntf == 0 ) then   ! S version
         if ( iDisc == 0 ) then
            call algSetupTR ( nPhs, nY, nU, nP, nC, ctype, nnY, nlY, nnC, nlC, &
                              neG, nnCon, ndPtr, n, lenC, &
                              ncindJ, nclocJ, lcJcol, lcindJ, lclocJ, &
                              usralgS, algcon, cu, lencu, iu, leniu, ru, lenru )

         else
            call algSetupHS ( nPhs, nY, nU, nP, nC, ctype, nnY, nlY, nnC, nlC, &
                              neG, nnCon, ndPtr, n, lenC, &
                              ncindJ, nclocJ, lcJcol, lcindJ, lclocJ, &
                              usralgS, algcon, cu, lencu, iu, leniu, ru, lenru )

         end if
      else if ( iIntf == 1 ) then   ! D version
         if ( iDisc == 0 ) then
            call algSetupTR ( nPhs, nY, nU, nP, nC, ctype, nnY, nlY, nnC, nlC, &
                              neG, nnCon, ndPtr, n, lenC, &
                              ncindJ, nclocJ, lcJcol, lcindJ, lclocJ, &
                              usralgD, algcon, cu, lencu, iu, leniu, ru, lenru )

         else
            call algSetupHS ( nPhs, nY, nU, nP, nC, ctype, nnY, nlY, nnC, nlC, &
                              neG, nnCon, ndPtr, n, lenC, &
                              ncindJ, nclocJ, lcJcol, lcindJ, lclocJ, &
                              usralgD, algcon, cu, lencu, iu, leniu, ru, lenru )

         end if

      else if ( iIntf == 2 ) then   ! A version
         if ( iDisc == 0 ) then
            call algSetupTR ( nPhs, nY, nU, nP, nC, ctype, nnY, nlY, nnC, nlC, &
                              neG, nnCon, ndPtr, n, lenC, &
                              ncindJ, nclocJ, lcJcol, lcindJ, lclocJ, &
                              usralgA, algcon, cu, lencu, iu, leniu, ru, lenru )

         else
            call algSetupHS ( nPhs, nY, nU, nP, nC, ctype, nnY, nlY, nnC, nlC, &
                              neG, nnCon, ndPtr, n, lenC, &
                              ncindJ, nclocJ, lcJcol, lcindJ, lclocJ, &
                              usralgA, algcon, cu, lencu, iu, leniu, ru, lenru )

         end if
      end if


      !-------------------------------------------------------------------
      ! Combine the linear/nonlinear/ode/alg parts
      !-------------------------------------------------------------------
      neY = nflocJ(n+1) + lflocJ(n+1) - 2
      neC = nclocJ(n+1) + lclocJ(n+1) - 2
      nJ  = neY + neC

      work%lenJ = nJ
      allocate ( work%locJ(n+1), work%Jcol(nJ), work%indJ(nJ) )

      locJ => work%locJ
      Jcol => work%Jcol
      indJ => work%indJ

      count = 1
      do k = 1, n
         locJ(k) = count

         do j = nclocJ(k), nclocJ(k+1)-1
            indJ(count) = ncindJ(j)
            Jcol(count) = 0.0
            count       = count + 1
         end do

         do j = nflocJ(k), nflocJ(k+1)-1
            indJ(count) = nfindJ(j)
            Jcol(count) = 0.0
            count       = count + 1
         end do

         do j = lclocJ(k), lclocJ(k+1)-1
            indJ(count) = lcindJ(j)
            Jcol(count) = lcJcol(j)
            count       = count + 1
         end do

         do j = lflocJ(k), lflocJ(k+1)-1
            indJ(count) = lfindJ(j)
            Jcol(count) = lfJcol(j)
            count       = count + 1
         end do
      end do
      locJ(n+1) = count

      deallocate ( nfindJ, nflocJ, lfindJ, lflocJ, lfJcol )
      deallocate ( ncindJ, nclocJ, lcindJ, lclocJ, lcJcol )

    end subroutine c2deriv

    !+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

    subroutine c2algbds ( iDisc, nPhs, nC, nU, ctype, nnCon, intPtr, &
                           clbds, cubds, m, bl, bu )
      integer(ip), intent(in) :: iDisc, nPhs, nC, nU, ctype(nC,nPhs), nnCon, &
                                 intPtr(nPhs), m
      real(rp),    intent(in) :: clbds(nC, nPhs), cubds(nC, nPhs)
      real(rp),   intent(out) :: bl(m), bu(m)

      !=========================================================================
      ! Set up algebraic constraint bounds.
      !
      ! 04 Jan 2009: First version of conSetBds.
      ! 06 Feb 2010: Current version of c2algbds.
      !=========================================================================
      integer(ip) :: indn, indl, ldiff, ndiff, p, i, j


      indn = 0
      indl = nnCon

      do p = 1, nPhs
         ldiff = nlY(p)
         ndiff = nnY(p)
         if ( iDisc == 1 ) then
            ldiff = 2*nlY(p) + nU
            ndiff = 2*nnY(p)
         end if

         ! alg constraint bds
         do j = intPtr(p), intPtr(p+1)-1
            do i = 1, nC
               if ( ctype(i,p) == 0 ) then ! nonlinear
                  indn = indn + 1
                  bl(indn) = clbds(i,p)
                  bu(indn) = cubds(i,p)
               else ! linear
                  indl = indl + 1
                  bl(indl) = clbds(i,p)
                  bu(indl) = cubds(i,p)
               end if
            end do
            indn = indn + ndiff
            indl = indl + ldiff
         end do
      end do

      ! at tf
      do i = 1, nC
         if ( ctype(i,nPhs) == 0 ) then ! nonlinear
            indn = indn + 1
            bl(indn) = clbds(i,nPhs)
            bu(indn) = cubds(i,nPhs)
         else ! linear
            indl = indl + 1
            bl(indl) = clbds(i,nPhs)
            bu(indl) = cubds(i,nPhs)
         end if
      end do

    end subroutine c2algbds

    !+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

  end subroutine c2setup

end module ct29set
