! Copyright (c) 2024, Paul Muir, Jack Pew, Tom Murtha, Connor Tannahill
! Paul Muir, Mathematics and Computing Science, Saint Mary's University.
! All rights reserved.
!
! Redistribution and use in source and binary forms, with or without
! modification, are permitted provided that the following conditions
! are met:
! * Redistributions of source code must retain the above copyright
!   notice, this list of conditions and the following disclaimer.
! * Redistributions in binary form must reproduce the above copyright
!   notice, this list of conditions and the following disclaimer in the
!   documentation and/or other materials provided with the distribution.
!
! THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
! "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
! LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
! A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
! HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
! SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
! LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
! DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
! THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
! (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
! OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
! This file contains the BACOLI source code (including the new code
! that implements the interpolation based spatial error estimates.)

module rootfinding
! This module enables the main program (below) to interact with
! the BACOLIKR software in order to implement root finding
! while computing a numerical solution to a PDE. The root finding
! routine that implements the gstop capability employs the 
! values routine from BACOLIKR.

    use bacoli95_mod, only: bacoli95_sol

    implicit none

!   Global variables and data structures
    integer, parameter          :: dp = kind(0d0)

    ! Structured type containing numerical solution info, accessed 
    ! in this module for convenience 
    type(bacoli95_sol), pointer :: sol

    ! Work array for use in root subroutine, rt, so that calls to
    ! the 'values' subroutine can be employed
    real(dp), pointer           :: work(:)

    contains

    ! 1. 'rt' subroutine. This routine is called in
    ! order to evaluate the gstop vector. This routine is called
    ! by DASKR at the end of every successful time step to 
    ! check if any components of the gstop vector have changed
    ! sign. If a sign change is detected, the 'rt' subroutine is
    ! then called within a search algorithm to locate the specfic 
    ! time at which some component of the gstop vector changes sign.

    subroutine rt(t, x, nint, ub, utb, neq, rval, nrt)
    ! t is the current time, x is the spatial mesh, nint is
    ! the number of spatial subintervals, ub is the vector
    ! of B-spline coefficients, ubt is the vector of time
    ! derivatives of the B-spline coefficients, neq is the
    ! number of B-spline coefficieints, rval is the vector
    ! of gstop values at time t, nrt is the number of components
    ! of the gstop vector.

    ! In this example we are interested in determing the times
    ! at which the u_3(0.2) = 0.24 and u_4(0.2,t) = 0.24. Thus there are two 
    ! roots and the gstop vector is [u_3(0.2,t) - 0.24, u_4(0.2,t) - 0.24]

    ! This routine accesses the sol structured type and the work array globally

        ! Input arguments
        integer ::  nint, neq, nrt 
        real(dp) :: t, x(nint+1), ub(neq), utb(neq)

        ! Output arguments
        real(dp) :: rval(nrt)

        ! Local variables
        real(dp)            :: xsol(1) ! x value at which solution is evaluated
        real(dp)            :: usol(4) ! approximate solution values at xsol(1)
        real(dp)            :: utarget(2) ! Target value of the solution
        integer, parameter  :: npts   = 1 ! Number of evaluation points
        integer, parameter  :: nderiv = 0 ! Number of spatial derivatives of
                                          ! solution that are required

        ! The point in space where we are interested in the solution value
        xsol(1) = 0.2d0

        ! The solution values we are interested in matching, i.e., the target values
        utarget(1) = 0.24d0
        utarget(2) = 0.24d0

        ! Evaluate the solution at x=0.2 to look for the target values, 0.24. Here we
        ! call the 'values' routine with input vector 'ub' which evaluates the 
        ! approximate solution at xsol and at the current time. (Replacing 'ub' with 'ubt' 
        ! would evaluate the time derivative of the approximate solution at xsol and the 
        ! current time.) kcol, the number of collocation points per subinterval and npde, 
        ! the number of PDEs, are provided through fields of the solution data structure,
        ! 'sol'.The 'values' routine is an F77 routine that can be accessed globally.
        call values(sol%kcol, xsol, nint, x, sol%npde, npts, nderiv, usol, ub, work)

        ! The values of the approximate solution at xsol and the current time are 
        ! returned in the usol vector.

        ! We next assign the gstop values to be the difference between each
        ! approximate solution value and the target solution value.
        rval(1) = usol(3) - utarget(1)
        rval(2) = usol(4) - utarget(2)
    end subroutine rt

    ! 2. Subroutine to allocate work array for use in 'rt' subroutine when it calls
    ! the 'values' routine to evalute the approximate solution. The required 
    ! dimension of the work array depends upon some of the fields of the 'sol' 
    ! structured type so this routine takes 'sol' as an input argument
    subroutine setrt(in_sol)
        type(bacoli95_sol), target :: in_sol

        ! Local variables used in defining size of work array
        integer, parameter :: nconti = 2
        integer, parameter :: nderiv = 0

        ! Assign the pointer in the module scope to be equal to in_sol
        sol => in_sol
        
        ! Allocate work memory for use in 'rt'
        allocate(work((sol%kcol+nconti)*(nderiv+1)+(sol%kcol+nconti+1)   &
            *(sol%kcol+nconti+2)/2+2*sol%kcol*(sol%nint_max+1)+2*nconti))
     end subroutine setrt

end module rootfinding
program simple_example_driver
    ! The BACOLIKR solver is called bacoli95.

    ! See the comments within bacoli95.f95 for details regarding use.

    ! Import required components of the bacoli95 module
    use bacoli95_mod, only: bacoli95_init, bacoli95, bacoli95_vals
    use bacoli95_mod, only: bacoli95_sol, bacoli95_sol_teardown

    ! Import components of the rootfinding module
    use rootfinding, only: rt, setrt

    implicit none
    integer, parameter     :: dp = kind(0d0)

    ! Solution data structure 
    type(bacoli95_sol), target  :: sol

    ! Number of PDEs = 4, only solution value is required so nderiv=0
    integer,  parameter    :: npde = 4, nderiv = 0

    ! Set endpoints of spatial domain
    real(dp), parameter    :: xa = 0, xb = 1

    ! Declare number of output points, nout, and arrays, xout and uout
    ! for storage of output points and corresponding solution values
    integer                :: nout
    real(dp), allocatable  :: xout(:), uout(:,:)

    ! Declare final time, tout, and arrays for absolute and relative
    ! tolerances, atol and rtol - these have one component per solution 
    ! value.
    real(dp)               :: tout, atol(npde), rtol(npde)

    ! Declare loop index, i, error flag for array allocation, ier;
    ! nrt is the number of gstop roots to be found
    integer                :: i, ier, nrt

    ! Max number of subintervals allowed in the spatial mesh = 500
    integer                :: nint_max = 500

    ! Declare number of collocation points, kcol, and the argument that
    ! specifies the type of error estimator that bacoli95 will use, estimator
    integer kcol, estimator

    ! Subroutines that define the rhs of the PDE, the boundary conditions,
    ! and the initial condition. We will ask bacoli95 to approximate the
    ! derivatives for these values so we do not need to declare the derivative
    ! routines to be external.
    external f, bndxa, bndxb, uinit 
    !external derivf, difbxa, difbxb

    ! Declare problem specific variables. Communicate these values to the
    ! problem dependent routines using a common block
    real(dp) :: a1,a2,d1,d2,r,c,n,pe1,pe2
    common /rcdstuff/ a1,a2,d1,d2,r,c,n,pe1,pe2

    !-------------------------------------------------------------------
    ! Set problem dependent parameters

    a1 = 30._dp
    a2 = 30._dp
    d1 = 1.5_dp
    d2 = 1.2_dp
    r = 1000._dp
    c = .96_dp
    n = 1._dp
    pe1 = 1e4_dp
    pe2 = 1e4_dp

    ! Set output time
    tout = 20.0d0

    ! Set tolerances
    atol = 1d-5
    rtol = atol

    ! Set number of roots to be found
    nrt = 2

    ! Initialize the sol structured type
    ! Initial mesh is uniform partition of [xa,ab], error estimator is the LOI,
    ! number of collocation points per subinterval, kcol, is 4
    call bacoli95_init(sol, npde, (/xa,xb/), atol=atol, rtol=rtol, &
               estimator=0, kcol=4, nint_max=nint_max, nrt = nrt)

    ! Allocate memory for the work array to be used during the root finding
    call setrt(sol)

    ! Set ouput at 101 uniformly spaced points across the spatial domain
    nout = 101

    ! Allocate memory for the xout and uout arrays.
    allocate(xout(nout), uout(npde,nout), stat=ier)
    if (ier /= 0 .or. sol%idid == -1000) goto 700

    ! Begin computation
    ! call bacoli95 with derivative routines derivf, difbxa, difbxb
    !call bacoli95(sol, tout, f, bndxa, bndxb, uinit, derivf, &
                                             !difbxa, difbxb, rt=rt)
    ! call bacoli95 without derivative routines
    call bacoli95(sol, tout, f, bndxa, bndxb, uinit, rt=rt)

    ! Check stop due to error condition; sol%idid > 0 implies successful 
    ! return from bacoli95
    if (sol%idid .le. 0) then
            write(*,*) 'Error? Idid = ', sol%idid
            stop 
    end if

    ! Successful return from bacoli95. If bacoli95 has stopped 
    ! upon finding the gstop root, this will be indicated by a return with 
    ! sol%idid == 5. If so, print out time and solution. Otherwise stop.
    if (sol%idid == 5) then
            print'("done. idid=",i5)',sol%idid
            write(*,*) 'Root found at t = ', sol%t0
            write(*,*) 'Output solution'
            ! Set output points and evaluate solution at these points
            xout = (/xa,(xa+i*(xb-xa)/(nout-1), i=1,nout-1)/)
            call bacoli95_vals(sol, xout, uout, nderiv=nderiv)

            print '(/"THE OUTPUT IS")'
            print '(a13,a27)', "XOUT", "UOUT"
            do i = 1, nout
                print*, xout(i), uout(:,i)
            end do
    else
            print'("done. idid=",i5)',sol%idid
            write(*,*) 'Root not found before reaching tout'
            stop
    end if

    ! Reaching this point means that bacoli95 returned after finding a gstop
    ! root. We have printed out the time and solution and now wish to
    ! continue the integration. We must set sol%idid=1 to continue the
    ! integration. If we had changed the PDE we would have to invoke a cold 
    ! restart by setting sol%mflag(1) = 2; since we have not changed the PDE, we 
    ! can continue with a warm restart, indicated by setting sol%mflag(1) = 1.

    sol%idid = 1
    sol%mflag(1) = 1

    ! Solve the PDE from the current time to the next gstop point.

    write(*,*)
    write(*,*) 'Continuing integration from t = ', sol%t0
    write(*,*)

    !call bacoli95(sol, tout, f, bndxa, bndxb, uinit, derivf, &
                                          !difbxa, difbxb, rt=rt)
    call bacoli95(sol, tout, f, bndxa, bndxb, uinit, rt=rt)

    ! Successful return from bacoli95. If bacoli95 has stopped 
    ! upon finding the next gstop root, this will be indicated by a return with 
    ! sol%idid == 5. If so, print out time and solution. Otherwise stop.
    if (sol%idid == 5) then
            print'("done. idid=",i5)',sol%idid
            write(*,*) 'Root found at t = ', sol%t0
            write(*,*) 'Output solution'
            ! Set output points and evaluate solution at these points
            xout = (/xa,(xa+i*(xb-xa)/(nout-1), i=1,nout-1)/)
            call bacoli95_vals(sol, xout, uout, nderiv=nderiv)

            print '(/"THE OUTPUT IS")'
            print '(a13,a27)', "XOUT", "UOUT"
            do i = 1, nout
                print*, xout(i), uout(:,i)
            end do
    else
            print'("done. idid=",i5)',sol%idid
            write(*,*) 'Root not found before reaching tout'
            stop
    end if

    ! Computation terminates after second root is found.

    call bacoli95_sol_teardown(sol) ; stop
600 print '("Error: Improperly formatted input")' ; stop
700 print '("Error: Could not allocate storage")' ; stop

    !-------------------------------------------------------------------
900 format("kcol = ",i2,", nint0 = ",i4,", npde = ",i3,", tout = ",es7.1)
901 format("atol = ",es7.1,", rtol = ",es7.1,",",17x,a3)
902 format("Number of subintervals in the current mesh:",i8)

end program
