! -----------------------------------------------------------------
! Programmer(s): Cody J. Balos @ LLNL
! -----------------------------------------------------------------
! Acknowledgements: Based on idaAkzoNob_ASAi_dns.c by Radu Serban
! and Cosmin Petra @ LLNL.
! -----------------------------------------------------------------
! SUNDIALS Copyright Start
! Copyright (c) 2002-2024, Lawrence Livermore National Security
! and Southern Methodist University.
! All rights reserved.
!
! See the top-level LICENSE and NOTICE files for details.
!
! SPDX-License-Identifier: BSD-3-Clause
! SUNDIALS Copyright End
! -----------------------------------------------------------------
! Adjoint sensitivity example problem
!
! This IVP is a stiff system of 6 non-linear DAEs of index 1. The
! problem originates from Akzo Nobel Central research in Arnhern,
! The Netherlands, and describes a chemical process in which 2
! species are mixed, while carbon dioxide is continuously added.
! See http://pitagora.dm.uniba.it/~testset/report/chemakzo.pdf
!
! FIDAS also computes the sensitivities with respect to initial
! conditions of the following quantity:
!   G = int_t0^t1 y1 dt
! The sensitivity of G is the solution of the adjoint system at t0.
! -----------------------------------------------------------------

module dae_mod
  use, intrinsic :: iso_c_binding
  use fsundials_core_mod
  implicit none

  ! Since SUNDIALS can be compiled with 32-bit or 64-bit sunindextype
  ! we set the integer kind used for indices in this example based
  ! on the the index size SUNDIALS was compiled with so that it works
  ! in both configurations. This is not a requirement for user codes.
#if defined(SUNDIALS_INT32_T)
  integer, parameter :: myindextype = selected_int_kind(8)
#elif defined(SUNDIALS_INT64_T)
  integer, parameter :: myindextype = selected_int_kind(16)
#endif

  ! problem parameters
  integer(kind=myindextype), parameter :: NEQ = 6
  integer(c_long), parameter :: STEPS = 150
  real(c_double), parameter :: T0 = 0.0d0
  real(c_double), parameter :: TF = 180.d0
  real(c_double), parameter :: RTOL = 1e-08
  real(c_double), parameter :: ATOL = 1e-10
  real(c_double), parameter :: RTOLB = 1e-06
  real(c_double), parameter :: ATOLB = 1e-08
  real(c_double), parameter :: RTOLQ = 1e-10
  real(c_double), parameter :: ATOLQ = 1e-12

  ! problem constants
  real(c_double) :: ZERO = 0.0d0
  real(c_double) :: QUARTER = 0.25d0
  real(c_double) :: HALF = 0.5d0
  real(c_double) :: ONE = 1.0d0
  real(c_double) :: TWO = 2.0d0
  real(c_double) :: FOUR = 4.0d0
  real(c_double) :: EIGHT = 8.0d0

  ! problem data
  real(c_double) :: k1, k2, k3, k4
  real(c_double) :: K, klA, Ks, pCO2, H

contains

  integer(c_int) function res(t, nv_yy, nv_yd, nv_resval, userdata) &
    result(retval) bind(C, name='res')
    use, intrinsic :: iso_c_binding

    implicit none

    ! function arguments
    real(c_double), value :: t
    type(N_Vector)        :: nv_yy, nv_yd, nv_resval
    type(c_ptr)           :: userdata

    ! local variables
    real(c_double)          :: y1, y2, y3, y4, y5, y6
    real(c_double)          :: yd1, yd2, yd3, yd4, yd5
    real(c_double)          :: r1, r2, r3, r4, r5, Fin
    real(c_double), pointer :: yy(:), yd(:), resval(:)

    yy => FN_VGetArrayPointer(nv_yy)
    yd => FN_VGetArrayPointer(nv_yd)
    resval => FN_VGetArrayPointer(nv_resval)

    y1 = yy(1)
    y2 = yy(2)
    y3 = yy(3)
    y4 = yy(4)
    y5 = yy(5)
    y6 = yy(6)

    yd1 = yd(1)
    yd2 = yd(2)
    yd3 = yd(3)
    yd4 = yd(4)
    yd5 = yd(5)

    r1 = k1*(y1**4)*sqrt(y2)
    r2 = k2*y3*y4
    r3 = k2/K*y1*y5
    r4 = k3*y1*y4*y4
    r5 = k4*y6*y6*sqrt(y2)
    Fin = klA*(pCO2/H - y2)

    resval(1) = yd1 + TWO*r1 - r2 + r3 + r4
    resval(2) = yd2 + HALF*r1 + r4 + HALF*r5 - Fin
    resval(3) = yd3 - r1 + r2 - r3
    resval(4) = yd4 + r2 - r3 + TWO*r4
    resval(5) = yd5 - r2 + r3 - r5
    resval(6) = Ks*y1*y4 - y6

    retval = 0
    return
  end function res

  integer(c_int) function rhsQ(t, nv_yy, nv_yp, nv_qdot, userdata) &
    result(retval) bind(C, name='rhsQ')
    use, intrinsic :: iso_c_binding

    implicit none

    ! function arguments
    real(c_double), value :: t
    type(N_Vector)        :: nv_yy, nv_yp, nv_qdot
    type(c_ptr)           :: userdata

    ! local variables
    real(c_double), pointer :: qdot(:), yy(:)

    qdot => FN_VGetArrayPointer(nv_qdot)
    yy => FN_VGetArrayPointer(nv_yy)
    qdot(1) = yy(1)

    retval = 0
    return
  end function rhsQ

  integer(c_int) function resB(tt, nv_yy, nv_yp, nv_yyB, nv_ypB, nv_rrB, userdata) &
    result(retval) bind(C, name='resB')
    use, intrinsic :: iso_c_binding
    implicit none

    ! function arguments
    real(c_double), value :: tt
    type(N_Vector)        :: nv_yy, nv_yp, nv_yyB, nv_ypB, nv_rrB
    type(c_ptr)           :: userdata

    ! local variables
    real(c_double) :: y1, y2, y3, y4, y5, y6
    real(c_double) :: yB1, yB2, yB3, yB4, yB5, yB6
    real(c_double) :: ypB1, ypB2, ypB3, ypB4, ypB5
    real(c_double) :: y2tohalf, y1to3, k2overK, tmp1, tmp2
    real(c_double), pointer :: yy(:), yyB(:), ypB(:), rrb(:)

    yy => FN_VGetArrayPointer(nv_yy)
    yyB => FN_VGetArrayPointer(nv_yyB)
    ypB => FN_VGetArrayPointer(nv_ypB)
    rrB => FN_VGetArrayPointer(nv_rrB)

    y1 = yy(1)
    y2 = yy(2)
    y3 = yy(3)
    y4 = yy(4)
    y5 = yy(5)
    y6 = yy(6)

    yB1 = yyB(1)
    yB2 = yyB(2)
    yB3 = yyB(3)
    yB4 = yyB(4)
    yB5 = yyB(5)
    yB6 = yyB(6)

    ypB1 = ypB(1)
    ypB2 = ypB(2)
    ypB3 = ypB(3)
    ypB4 = ypB(4)
    ypB5 = ypB(5)

    y2tohalf = sqrt(y2)
    y1to3 = y1*y1*y1
    k2overK = k2/K

    tmp1 = k1*y1to3*y2tohalf
    tmp2 = k3*y4*y4
    rrB(1) = 1 + ypB1 - (EIGHT*tmp1 + k2overK*y5 + tmp2)*yB1 &
             - (TWO*tmp1 + tmp2)*yB2 + (FOUR*tmp1 + k2overK*y5)*yB3 &
             + k2overK*y5*(yB4 - yB5) - TWO*tmp2*yB4 + Ks*y4*yB6

    tmp1 = k1*y1*y1to3*(y2tohalf/y2)
    tmp2 = k4*y6*y6*(y2tohalf/y2)
    rrB(2) = ypB2 - tmp1*yB1 - (QUARTER*tmp1 + QUARTER*tmp2 + klA)*yB2 &
             + HALF*tmp1*yB3 + HALF*tmp2*yB5

    rrB(3) = ypB3 + k2*y4*(yB1 - yB3 - yB4 + yB5)

    tmp1 = k3*y1*y4
    tmp2 = k2*y3
    rrB(4) = ypB4 + (tmp2 - TWO*tmp1)*yB1 - TWO*tmp1*yB2 - tmp2*yB3 &
             - (tmp2 + FOUR*tmp1)*yB4 + tmp2*yB5 + Ks*y1*yB6

    rrB(5) = ypB5 - k2overK*y1*(yB1 - yB3 - yB4 + yB5)

    rrB(6) = k4*y6*y2tohalf*(2*yB5 - yB2) - yB6

    retval = 0
    return
  end function resB

  subroutine PrintOutput(nv_yB, nv_ypB)
    use, intrinsic :: iso_c_binding
    use fsundials_core_mod
    implicit none

    ! function arguments
    type(N_Vector) :: nv_yB, nv_ypB
    real(c_double), pointer :: yB(:)

    yB => FN_VGetArrayPointer(nv_yB)

    write (*, '(1x,A,es12.4)') "dG/dy0:         ", yB(1)
    write (*, '(1x,A,es12.4)') "                ", yB(2)
    write (*, '(1x,A,es12.4)') "                ", yB(3)
    write (*, '(1x,A,es12.4)') "                ", yB(4)
    write (*, '(1x,A,es12.4)') "                ", yB(5)
    write (*, '(1x,A,es12.4)') "                ", yB(6)
    write (*, *) "--------------------------------------------------------"
    write (*, *) ""

  end subroutine

end module dae_mod

! Main program
program main
  use, intrinsic :: iso_c_binding
  use fsundials_core_mod
  use fidas_mod                  ! Fortran interface to IDA
  use fnvector_serial_mod        ! Fortran interface to serial N_Vector
  use fsunmatrix_dense_mod       ! Fortran interface to dense SUNMatrix
  use fsunlinsol_dense_mod       ! Fortran interface to dense SUNLinearSolver

  use dae_mod                    ! DAE problem module
  implicit none

  ! Local variables
  type(c_ptr)                    :: sunctx
  type(c_ptr)                    :: mem, memB
  integer(c_int)                 :: ncheck(1), retval
  real(c_double)                 :: time(1)
  integer(c_long)                :: nst(1), nstB(1)
  integer(c_int)                 :: indexB(1)
  real(c_double), pointer :: yy(:), q(:)
  type(N_Vector), pointer :: nv_yy, nv_yp, nv_rr, nv_q
  real(c_double), pointer :: ypB(:)
  type(N_Vector), pointer :: nv_yB, nv_ypB
  type(SUNMatrix), pointer :: A, AB
  type(SUNLinearSolver), pointer :: LS, LSB

  ! Consistent IC for  y, y'.
  real(c_double) :: y01 = 0.444d0
  real(c_double) :: y02 = 0.00123d0
  real(c_double) :: y03 = 0.0d0
  real(c_double) :: y04 = 0.007d0
  real(c_double) :: y05 = 0.0d0

  write (*, *) ""
  write (*, *) "Adjoint Sensitivity Example for Akzo-Nobel Chemical Kinetics"
  write (*, *) "-------------------------------------------------------------"
  write (*, *) "Sensitivity of G = int_t0^tf (y1) dt with respect to IC."
  write (*, *) "-------------------------------------------------------------"
  write (*, *) ""

  ! Fill problem data with the appropriate values for coefficients.
  k1 = 18.7d0
  k2 = 0.58d0
  k3 = 0.09d0
  k4 = 0.42d0
  K = 34.4d0
  klA = 3.3d0
  Ks = 115.83d0
  pCO2 = 0.9d0
  H = 737.0d0

  ! Create the SUNDIALS simulation context
  retval = FSUNContext_Create(SUN_COMM_NULL, sunctx)
  call check_retval(retval, "FSUNContext_Create")

  ! Allocate N-vectors.
  nv_yy => FN_VNew_Serial(NEQ, sunctx)
  nv_yp => FN_VNew_Serial(NEQ, sunctx)

  ! Set IC
  yy => FN_VGetArrayPointer(nv_yy)
  yy(1) = y01
  yy(2) = y02
  yy(3) = y03
  yy(4) = y04
  yy(5) = y05
  yy(6) = Ks*y01*y04

  ! Get y' = - res(t0, y, 0)
  call FN_VConst(ZERO, nv_yp)

  nv_rr => FN_VNew_Serial(NEQ, sunctx)
  retval = res(T0, nv_yy, nv_yp, nv_rr, c_null_ptr)
  call FN_VScale(-ONE, nv_rr, nv_yp)
  call FN_VDestroy(nv_rr)

  ! Create and initialize q0 for quadratures.
  nv_q => FN_VNew_Serial(1_myindextype, sunctx)
  if (.not. associated(nv_q)) then
    write (*, *) 'ERROR: FN_VNew_Serial returned NULL'
    stop 1
  end if

  q => FN_VGetArrayPointer(nv_q)
  if (.not. associated(q)) then
    write (*, *) 'ERROR: FN_VGetArrayPointer returned NULL'
    stop 1
  end if
  q(1) = ZERO

  ! Call FIDACreate and FIDAInit to initialize FIDA memory
  mem = FIDACreate(sunctx)
  if (.not. c_associated(mem)) then
    write (*, *) 'ERROR: FIDACreate returned NULL'
    stop 1
  end if

  retval = FIDAInit(mem, c_funloc(res), T0, nv_yy, nv_yp)
  call check_retval(retval, "FIDAInit")

  ! Set tolerances.
  retval = FIDASStolerances(mem, RTOL, ATOL)
  call check_retval(retval, "FIDASStolerances")

  ! Create dense SUNMatrix for use in linear solves
  A => FSUNDenseMatrix(NEQ, NEQ, sunctx)
  if (.not. associated(A)) then
    write (*, *) 'ERROR: FSUNDenseMatrix returned NULL'
    stop 1
  end if

  ! Create dense SUNLinearSolver object
  LS => FSUNLinSol_Dense(nv_yy, A, sunctx)
  if (.not. associated(LS)) then
    write (*, *) 'ERROR: FSUNLinSol_Dense returned NULL'
    stop 1
  end if

  ! Attach the matrix and linear solver
  retval = FIDASetLinearSolver(mem, LS, A)
  call check_retval(retval, "FIDASetLinearSolver")

  ! Initialize QUADRATURE(S).
  retval = FIDAQuadInit(mem, c_funloc(rhsQ), nv_q)
  call check_retval(retval, "FIDAQuadInit")

  ! Set tolerances and error control for quadratures.
  retval = FIDAQuadSStolerances(mem, RTOLQ, ATOLQ)
  call check_retval(retval, "FIDAQuadSStolerances")

  retval = FIDASetQuadErrCon(mem, 1)
  call check_retval(retval, "FIDASetQuadErrCon")

  ! Prepare ADJOINT.
  retval = FIDAAdjInit(mem, STEPS, IDA_HERMITE)
  call check_retval(retval, "FIDAAdjInit")

  ! FORWARD run.
  write (*, '(1x,A)', advance='no') "Forward integration ... "
  retval = FIDASolveF(mem, TF, time, nv_yy, nv_yp, IDA_NORMAL, ncheck)
  call check_retval(retval, "FIDASolveF")
  retval = FIDAGetNumSteps(mem, nst)
  write (*, '(A,i6,A)') "done ( nst = ", nst, " )"
  retval = FIDAGetQuad(mem, time, nv_q)

  write (*, '(1x,A,F24.16)') "G:          ", q(1)
  write (*, *) "--------------------------------------------------------"
  write (*, *) ""

  ! BACKWARD run

  ! Initialize yB
  nv_yB => FN_VNew_Serial(NEQ, sunctx)
  if (.not. associated(nv_yB)) then
    write (*, *) 'ERROR: FN_VNew_Serial returned NULL'
    stop 1
  end if
  call FN_VConst(ZERO, nv_yB)

  nv_ypB => FN_VNew_Serial(NEQ, sunctx)
  if (.not. associated(nv_ypB)) then
    write (*, *) 'ERROR: FN_VNew_Serial returned NULL'
    stop 1
  end if
  call FN_VConst(ZERO, nv_ypB)
  ypB => FN_VGetArrayPointer(nv_ypB)
  ypB(1) = -ONE

  retval = FIDACreateB(mem, indexB)
  call check_retval(retval, "FIDACreateB")

  retval = FIDAInitB(mem, indexB(1), c_funloc(resB), TF, nv_yB, nv_ypB)
  call check_retval(retval, "FIDAInitB")

  retval = FIDASStolerancesB(mem, indexB(1), RTOLB, ATOLB)
  call check_retval(retval, "FIDASStolerancesB")

  retval = FIDASetMaxNumStepsB(mem, indexB(1), 1000_c_long)
  call check_retval(retval, "FIDASetMaxNumStepsB")

  ! Create dense SUNMatrix for use in linear solves
  AB => FSUNDenseMatrix(NEQ, NEQ, sunctx)
  if (.not. associated(AB)) then
    write (*, *) 'ERROR: FSUNDenseMatrix returned NULL'
    stop 1
  end if

  ! Create dense SUNLinearSolver object
  LSB => FSUNLinSol_Dense(nv_yB, AB, sunctx)
  if (.not. associated(LSB)) then
    write (*, *) 'ERROR: FSUNLinSol_Dense returned NULL'
    stop 1
  end if

  ! Attach the matrix and linear solver
  retval = FIDASetLinearSolverB(mem, indexB(1), LSB, AB)
  call check_retval(retval, "FIDASetLinearSolverB")

  ! Do the backward integration
  write (*, '(1x,A)', advance='no') "Backward integration ... "
  retval = FIDASolveB(mem, T0, IDA_NORMAL)
  call check_retval(retval, "FIDASolveB")
  memB = FIDAGetAdjIDABmem(mem, indexB(1))
  retval = FIDAGetNumSteps(memB, nstB)
  write (*, '(A,i6,A)') "done ( nst = ", nstB, " )"
  retval = FIDAGetB(mem, indexB(1), time, nv_yB, nv_ypB)

  ! Print the solution
  call PrintOutput(nv_yB, nv_ypB)

  ! Free memory
  call FIDAFree(mem)
  retval = FSUNLinSolFree(LS)
  retval = FSUNLinSolFree(LSB)
  call FSUNMatDestroy(A)
  call FSUNMatDestroy(AB)
  call FN_VDestroy(nv_yy)
  call FN_VDestroy(nv_yp)
  call FN_VDestroy(nv_yB)
  call FN_VDestroy(nv_ypB)
  call FN_VDestroy(nv_q)
  retval = FSUNContext_Free(sunctx)

end program

subroutine check_retval(retval, name)
  use, intrinsic :: iso_c_binding

  character(len=*) :: name
  integer(c_int)   :: retval

  if (retval < 0) then
    write (*, '(A,A,A,I4)') 'ERROR: ', name, ' returned ', retval
    stop 1
  end if
end subroutine
