! 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 (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. Subroutine 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 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 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 times
        ! at which the solution at the left and right boundaries reaches
        ! a given critical value, ucrit. Thus there are two roots and the 
        ! gstop vector is [u(0,t) - ucrit, u(1,t) - ucrit]. (The spatial 
        ! domain is [0,1].)

        ! 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(2) ! x values at which solution is evaluated
        real(dp)            :: usol(4) ! approximate solution values at xsol
        real(dp)            :: upnt(2) ! Target values of the solution at the
                                       ! two boundaries.
        integer, parameter  :: npts   = 2 ! Number of evaluation points
        integer, parameter  :: nderiv = 1 ! Number of spatial derivatives of
                                          ! solution that are required
        real(dp)            :: bvala1, bvalb1, bvala2, bvalb2

        !PROBLEM CONSTANTS
        real(dp)            ::  kappa, alpha, gamma, ucrit
        common /HEATPARMS/      kappa, alpha, gamma, ucrit
        !---------------------------------------------------------------------------

        ! The points in space where we are interested in the solution value; these
        ! are the boundary points for this problem.
        xsol(1) = 0.0d0
        xsol(2) = 1.0d0

        ! The solution values we are interested in matching, i.e., the target value
        upnt(1) = ucrit
        upnt(2) = ucrit
        
        ! Evaluate the solution to look for the target value. Here we call
        ! 'values' 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') 
        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 usol. Inside 'values', usol is of dimension (npde, npts, nderiv+1).
        ! Since npde=1 and nderiv=0, and usol is declared to be a 1D array in this routine,
        ! usol(1) is the solution at xsol(1)=0, and usol(2) is the solution at xsol(2)=1.

        ! We next assign the gstop value to be the difference between the
        ! approximate solution values and the target solution values at the
        ! two boundary points.
        rval(1) = usol(1) - upnt(1) ! at x = 0.
        rval(2) = usol(2) - upnt(2) ! at x = 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 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

    ! An example driver for the BACOLIKR95 wrapper of BACOLIKR.
    ! The BACOLIKR solver is actually 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

    ! Set number of PDEs and left and right boundary points.
    integer,  parameter    :: npde = 1
    real(dp), parameter    :: xa = 0, xb = 1

    ! Set number of spatial derivatives of the solution to be printed out 
    ! in addition to the solution itself; set the number of output points;
    ! set number of events, i.e., number of gstop roots; set max number of
    ! spatial subintervals.
    integer,  parameter    :: nderiv = 1, nout = 101, nrt = 2, nint_max = 1000

    ! Declare arrays for output points and output solution values 
    real(dp) :: xout(nout), uout(npde*nout*(nderiv+1))

    ! Declare output time and absolute and relative tolerances
    real(dp)               :: tout, atol(npde), rtol(npde)

    ! Loop indicies
    integer                :: i, ij, k

    ! There are two gstop roots, one at x = 0, one at x = 1.
    ! jroot indicates which root has been found when bacolikr
    ! returns with an indication that it has found a root.
    ! jroot is a field of the sol structured type.
    integer                :: jroot(nrt)

    ! f defines the rhs of the PDE; bndxa1, bndxb1 are the left and
    ! right original boundary conditions; bndxa2, bndxb2 are the new
    ! left and right boundary conditions after an event. derivf, 
    ! difbxa1, difbxb1, difbxa2, difbxb2, provide the erivatives of
    ! f and the boundary conditions
    external f, bndxa1, bndxb1, uinit, derivf, difbxa1, difbxb1
    external    bndxa2, bndxb2,                difbxa2, difbxb2

    ! Problem specific variables
    real(dp)            ::  kappa, alpha, gamma, ucrit
    common /HEATPARMS/      kappa, alpha, gamma, ucrit

    !-------------------------------------------------------------------
    ! kappa = 0.1 for an easy version of the heat equation where the 
    ! diffusion is not too strong which means that the solution  
    ! changes not too quickly. 
    
    kappa = 0.1d0 

    ! The critical value for the solution at either boundary is ucrit. 

    ! A trivial form of this problem that gives a smooth transition from 
    ! the pre-event boundaries, u_x(a,t) = alpha and u_x(b,t) = alpha, to 
    ! the post-event boundaries, u_x(a,t) = -gamma*u(a,t) and u_x(b,t) = 
    ! -gamma*u(b,t), is obtained by choosing alpha = gamma = 0. A more 
    ! interesting version of this problem, that still gives a smooth 
    ! transition from the pre-event boundaries to the post-event boundaries,
    ! is obtained by requiring that alpha = -gamma * ucrit.)

    ! (Discontinuous boundary condition transitions are obtained by not
    ! imposing the condition that alpha = -gamma * ucrit. An example is to
    ! choose alpha = 0, so the spatial derivative of the solution at x=0 and 
    ! x=1 is zero for the pre-event part of the integration, and then gamma = 
    ! 0.5 and ucrit = 1 so that the post-event boundary conditions require 
    ! the spatial derivative of the solution at x=0 and x=1 to be -0.5 times  
    ! the corresponding boundary solution value.)

    ! Choose one of the following two options:

    ! Continuous change in boundary conditions
    alpha = 1.0d0
    ucrit = 1.1d0
    gamma = -alpha/ucrit

    ! Discontinuous change in boundary conditions
    !alpha = 0.0d0
    !ucrit = 1.0d0
    !gamma = 0.5d0

    ! Set output time
    tout = 5.0D0

    ! Write out value of npde, tout, and the problem parameters.
    
    write(6,*) 'Number of PDEs and tout', npde, tout
    write(6,*) 'kappa, alpha, gamma, ucrit = ', kappa, alpha, gamma, ucrit

    ! Set tolerance
    atol = 1d-6
    rtol = atol

    write(6,*) 'Absolute and relative tolerances', atol, rtol
    write(6,*)

    ! Initialization
    call bacoli95_init(sol, npde, (/xa,xb/), atol=atol, rtol=rtol, &
                                              nint_max=nint_max, nrt = nrt)
    call setrt(sol)

    ! Begin computation; start with the first versions of the boundary
    ! routines and the corresponding derivative routines.
    call bacoli95(sol, tout, f, bndxa1, bndxb1, uinit, rt=rt)
    write(*,*) 'First return from BACOLI'
    write(7,*)

    ! 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
            write(*,*) 'Return with t = ', sol%t0
            stop
    end if

    ! Successful return from bacoli95. If bacoli95 has stopped 
    ! upon finding a gstop root, this will be indicated by a return with 
    ! sol%idid == 5. If so, first determine which event occurred and then
    ! print out time, solution, and first spatial derivative. Otherwise, 
    ! print out time, solution, and first spatial derivative and stop.

    if (sol%idid == 5) then

            print'("Return from bacoli95. idid=",i5)',sol%idid
            write(*,*) 'Root found at t = ', sol%t0

            ! Assign jroot field from sol to jroot
            jroot = sol%jroot

            if (jroot(1) .ne. 0) then ! jroot(1) not equal to 0 implies solution 
                                      ! at left boundary reached ucrit.
                write(6,*) 'Solution at left boundary has reached critical value'
            else                      ! jroot(1) equal to 0, implies jroot(2) not
                                      ! equal to zero, which implies solution 
                                      ! at right boundary reached ucrit.
                write(6,*) 'Solution at right boundary has reached critical value'
            end if

            write(*,*) 'Output solution'
            ! Set output points and evaluate solution and first spatial derivative 
            ! 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
              ij = (i-1)*npde
              write(*,*) xout(i), (uout(ij+k), k = 1, npde), (uout(nout+ij+k), k = 1, npde)
            end do
            write(*,*)

    else    ! bacoli95 has terminated at tout without finding an event.

            print'("Return from bacoli95. idid=",i5)',sol%idid
            write(*,*) 'Root not found before reaching tout'
            write(*,*) 'Integration completed at t = ', sol%t0

            ! 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
              ij = (i-1)*npde
              write(*,*) xout(i), (uout(ij+k), k = 1, npde), (uout(nout+ij+k), k = 1, npde)
            end do
            write(*,*)
            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. 
    sol%idid = 1

    ! Continue with a warm restart by setting sol%mflag(1) = 1 or a cold restart by
    ! setting sol%mflag(1) = 2
    sol%mflag(1) = 2

    ! We need to change one of the two boundary conditions, depending on which
    ! gstop root was found, either the one at the left boundary or the one at the
    ! right boundary. This can be detected by accessing the jroot vector again.

    if (jroot(1) .ne. 0) then
        ! The gstop root just detected is associated with the left boundary condition. 
        ! Call bacoli95 to solve the modified problem from the current time to the next 
        ! gstop point. Use the second set of routines assoicated with the left
        ! boundary, i.e., bndxa2 and difbxa2, but continue to use the first set of
        ! routines associated with the right boundary, i.e., bndxb1 and difbxb1.

        write(*,*)
        write(*,*) 'Continuing integration from t = ', sol%t0
        write(*,*)
        call bacoli95(sol, tout, f, bndxa2, bndxb1, uinit, rt=rt)

   else 
        ! The gstop root is associated with the right boundary condition. Call
        ! bacoli95 to solve the modified problem from the current time to the next 
        ! gstop point. Use the second set of routines assoicated with the right
        ! boundary, i.e., bndxb2 and difbxb2, but continue to use the first set of
        ! routines associated with the left boundary, i.e., bndxa1 and difbxa1.

        write(*,*)
        write(*,*) 'Continuing integration from t = ', sol%t0
        write(*,*)
        call bacoli95(sol, tout, f, bndxa1, bndxb2, uinit, rt=rt)

    end if

    write(*,*) 'Second return from BACOLI'
    write(7,*)

    ! We expect, at this point, that bacoli95 has returned at the second gstop time.
    ! 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 a gstop root, this will be indicated by a return with 
    ! sol%idid == 5. If so, first determine which event occurred and then
    ! print out time, solution, and first spatial derivative. Otherwise, 
    ! print out time, solution, and first spatial derivative and stop.

    if (sol%idid == 5) then

            print'("Return from bacoli95. idid=",i5)',sol%idid
            write(*,*) 'Root found at t = ', sol%t0
            !tstar = sol%t0 ! Set tstar to event time for use in transitioning
                           ! the boundary condition

            ! Assign jroot field from sol to jroot
            jroot = sol%jroot

            if (jroot(1) .ne. 0) then ! jroot(1) not equal to 0 implies solution 
                                      ! at left boundary reached ucrit.
                write(6,*) 'Solution at left boundary has reached critical value'
            else                      ! jroot(1) equal to 0, implies jroot(2) not
                                      ! equal to zero, which implies solution 
                                      ! at right boundary reached ucrit.
                write(6,*) 'Solution at right boundary has reached critical value'
            end if

            write(*,*) 'Output solution'
            ! Set output points and evaluate solution and first spatial 
            ! derivative 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
              ij = (i-1)*npde
              write(*,*) xout(i), (uout(ij+k), k = 1, npde), (uout(nout+ij+k), k = 1, npde)
            end do
            write(*,*)           

    else    ! bacoli95 has terminated at tout without finding the second event.

            print'("Return from bacoli95. idid=",i5)',sol%idid
            write(*,*) 'Root not found before reaching tout'
            write(*,*) 'Integration completed at t = ', sol%t0
            ! Set output points and evaluate solution and first spatial derivative
            ! 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
              ij = (i-1)*npde
              write(*,*) xout(i), (uout(ij+k), k = 1, npde), (uout(nout+ij+k), k = 1, npde)
            end do
            
            stop

    end if

    ! At this point in time, both gstop roots have been found and we now wish to
    ! continue the integration. We must set sol%idid=1 to continue the integration. 
    sol%idid = 1

    ! Continue with a warm restart by setting sol%mflag(1) = 1 or a cold restart by
    ! setting sol%mflag(1) = 2
    sol%mflag(1) = 2

    ! Call bacoli95 to solve the modified problem from the current time to tout. 
    ! Use the second set of routines associated with the left and right
    ! boundary conditions for both boundaries. 

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

    call bacoli95(sol, tout, f, bndxa2, bndxb2, 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

    print'("Return from bacoli95. idid=",i5)',sol%idid
    write(*,*) 'Integration completed at t = ', sol%t0
    ! Set output points and evaluate solution and first spatial derivative
    ! 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
      ij = (i-1)*npde
      write(*,*) xout(i), (uout(ij+k), k = 1, npde), (uout(nout+ij+k), k = 1, npde)
    end do

    write(*,*)
    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) ; stop

end program
