! Copyright (c) 2024, Paul Muir, Jack Pew, Connor Tannahill
! 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 (see 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 in order to evaluate the approximate solution.

    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 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 repeatedly 
    ! within a search algorithm to locate the specific 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 current spatial mesh, nint is
        ! the current number of spatial subintervals, ub is the vector
        ! of B-spline coefficients at time t, ubt is the vector of time
        ! derivatives of the B-spline coefficients at time t, 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 time
        ! at which the approximate solution, evaluated at x = 0.4,
        ! takes on a value of 0.5. Since the solution transitions from
        ! approximately 1 to approximately 0, the middle of the transition
        ! layer can be defined as the point at which the solution takes on
        ! the value of 0.5. Thus the event can be viewed as asking for the
        ! time at which the layer is located at x=0.4. 

        ! Thus there is one root and the gstop vector is [u(0.45,t) - 0.5].

        ! 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(1) ! Solution value at x,t.
        real(dp)            :: utarget(1) ! 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 checking the solution value
        xsol(1) = 0.4d0

        ! The solution value we are interested in checking, i.e., the target value
        utarget(1) = 0.5d0

        ! Evaluate the solution at x=0.4 to determine if it equals the target value,
        ! 0.5. 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 value of the approximate solution at xsol and the current time is 
        ! returned in usol(1).

        ! We next assign the gstop value to be the difference between the
        ! approximate solution value at x,t, and the target solution value.
        rval(1) = usol(1) - utarget(1)

    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 setsol(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 setsol

end module rootfinding

program simple_example_driver
    ! A minimal example driver for the BACOLIKR95 wrapper of BACOLIKR.
    ! 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, &
                            bacoli95_sol, bacoli95_sol_teardown

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

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

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

    ! Number of PDEs = 1, Number of output points = 11, Max number
    ! of subintervals allowed in the spatial mesh = 500
    integer,  parameter :: npde = 1, nout = 11, nint_max = 500

    ! Terminal time value 
    real(dp), parameter :: tout = 2.0d0

    ! xout(:) contains output points, uout(:) contains solution
    ! approximations at output points, eps is a problem parameter
    real(dp) :: xout(nout), uout(nout), eps

    ! atol(1) is absolute tolerance, rtol(1) is relative tolerance
    real(dp) :: atol(1), rtol(1)

    ! loop index, i; nrt is the number of gstop roots to be found
    integer :: i, nrt


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


    ! Communicate eps value to problem dependent routines using common 
    ! block; set eps value
    common /burger/ eps
    data eps /1d-3/

    ! 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

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

    ! One event to be found
    nrt = 1

    ! Set tolerances
    atol(1) = 1d-6
    rtol(1) = atol(1)
 
    ! Initialize the sol structured type
    ! Initial mesh is uniform partition of [0,1], error estimator is the LOI,
    ! number of collocation points per subinterval, kcol, is 4
    call bacoli95_init(sol, npde, (/0._dp, 1._dp/), nint_max=nint_max, &
        nrt=nrt, estimator=0, kcol=4, atol=atol, rtol=rtol)

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

    ! Solve the PDE until the gstop root is found. 
    ! bacoli95 can be provided with the user routines that provide
    ! exact derivatives for f, bndxa, bndxb, namely, derivf, difbxa, and difbxb,
    ! see call below - but we choose to exclude these optional arguments, in which
    ! case bacoli95 will approximate the required derivative values.
    !call bacoli95(sol, tout, f, bndxa, bndxb, uinit, derivf, &
                                             !difbxa, difbxb, rt=rt)
    call bacoli95(sol, tout, f, bndxa, bndxb, uinit, rt=rt)

    ! Check stop due to an error condition; sol%idid > 0 implies a 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'
            xout = (/(i*0.1_dp, i=0,nout-1)/)
            call bacoli95_vals(sol, xout, uout)
            print'(5x,"x",5x5x,"u(t,x)")'
            print'(2f11.7)', transpose(reshape((/xout,uout/),(/nout,2/)))
    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 to tout. We must set sol%idid=1 to continue the
    ! integration. Since we have not changed the PDE, we can continue with a 
    ! warm restart, indicated by setting sol%mflag(1) = 1. (If we had changed the 
    ! PDE we would have to invoke a cold restart by setting sol%mflag(1) = 2).

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

    ! Solve the PDE from the current time to tout. 
    write(*,*)
    write(*,*) 'Continuing integration from t = ', sol%t0
    write(*,*)
    call bacoli95(sol, tout, f, bndxa, bndxb, uinit, rt=rt)

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

    ! Otherwise successful return from bacoli95
    if (sol%idid > 0) then 
        print'("done. idid=",i5)',sol%idid
        write(*,*) 'Return at t = ',  sol%t0
        xout = (/(i*0.1_dp, i=0,nout-1)/)
        call bacoli95_vals(sol, xout, uout)
        write(*,*) 'Output solution'
         print'(5x,"x",5x5x,"u(t,x)")'
        print'(2f11.7)', transpose(reshape((/xout,uout/),(/nout,2/)))
    end if

    write(*,*) 'stats'
    write(*,*) 'remeshings: ', sol%num_remeshings
    write(*,*) 'num_ini_remeshings: ', sol%num_ini_remeshings
    write(*,*) 'num_cold_starts: ', sol%num_cold_restarts
    write(*,*) 'num_accepted_time_steps: ', sol%num_accepted_time_steps

    call bacoli95_sol_teardown(sol)
end program

