! 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 (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, bacoli95_vals

    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 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 two layers of the approximate solution merge.
        ! The left layer corresponds to a transition in the solution
        ! value from 1 to 0.5. We therefore define the location of this 
	! layer to be the point, xL, where the solution has the value of
        ! 0.75, half way through the transition. The right layer 
	! corresponds to a transition in the solution value from 0.5 to 
	! 0.1. We therefore define the location of this layer to be the 
	! point, xR, where the solution has the value of 0.3, half way 
	! through the transition. Each layer has a width of O(epsilon) 
	! and we will define the layers to have merged when xL and xR
	! are within a distance of 5*epsilon. Thus there is one root and 
	! the gstop vector is [|xL-xR| .le. 5*epsilon].
	
	! For a given call to rt, we first search for xL and then search
	! for xR, in each case using a bisection algorithm to determine
	! the points xL and xR according to the definitions for the locations
	! of the layers given above.

        ! 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) ! approximate solution value at xsol(1)
        real(dp)            :: upnt(1) ! Target value 
        integer, parameter  :: npts   = 1 ! Number of evaluation points
        integer, parameter  :: nderiv = 0 ! Number of spatial derivatives of
                                          ! solution that are required
        real(dp)            :: a, b, m, UUa, UUb, UUm, xL, xR
		                       ! a,b: location of bisection interval and midpoint
				       ! UUa, UUB, UUb: corresponding solution values at 
				       ! these points

        upnt(1) = 0.005d0 ! Tolerance to within we want xL and xR to agree.
                          ! This is 5*epsilon, where epsilon = 1d-3.
							   
        ! Search for left layer
        a = 0.0d0

        ! Evaluate the solution at a. 
	! The value of the approximate solution at xsol and the current time is 
        ! returned in usol(1).
        xsol(1)=a
        call values(sol%kcol, xsol, nint, x, sol%npde, npts, nderiv, usol, ub, work)
        UUa = usol(1)

	! Evaluate the solution at b. 
	! The value of the approximate solution at xsol and the current time is 
        ! returned in usol(1).
        b = 1.0d0
        xsol(1)=b
        call values(sol%kcol, xsol, nint, x, sol%npde, npts, nderiv, usol, ub, work)
        UUb = usol(1)

        ! Compute midpoint of [a,b].
        m = a + (b-a)/2.0d0

	! Evaluate the solution at m. 
        xsol(1)=m
        call values(sol%kcol, xsol, nint, x, sol%npde, npts, nderiv, usol, ub, work)
        UUm = usol(1)

        !Bisection search until UUm is within tol=1d-6 of 0.75, location of left layer
        do while ((abs(UUm - 0.75d0) .gt. 1d-6))

          if ( ((UUa-0.75d0)*((UUm-0.75d0)) .gt. 0.0d0)) then
              !a and m are on the same side of the root 
              a = m
              UUa = UUm
          else
              !b and m are on the same side of the root 
              b = m
              UUb = UUm
          end if

          ! Compute midpoint of [a,b].
          m = a + (b-a)/2.0d0

	      ! Evaluate the solution at m. 
          xsol(1)=m
          call values(sol%kcol, xsol, nint, x, sol%npde, npts, nderiv, usol, ub, work)
          UUm = usol(1)

        end do

        ! xL is midpoint of current accepted interval.
        xL = m

        ! Search for right layer
        a = 0.0d0

        ! Evaluate the solution at a. 
	! The value of the approximate solution at xsol and the current time is 
        ! returned in usol(1).
        xsol(1)=a
        call values(sol%kcol, xsol, nint, x, sol%npde, npts, nderiv, usol, ub, work)
        UUa = usol(1)

	! Evaluate the solution at b. 
	! The value of the approximate solution at xsol and the current time is 
        ! returned in usol(1).
        b = 1.0d0
        xsol(1)=b
        call values(sol%kcol, xsol, nint, x, sol%npde, npts, nderiv, usol, ub, work)
        UUb = usol(1)

        ! Compute midpoint of [a,b].
        m = a + (b-a)/2.0d0

	! Evaluate the solution at m. 
        xsol(1)=m
        call values(sol%kcol, xsol, nint, x, sol%npde, npts, nderiv, usol, ub, work)
        UUm = usol(1)

        !Bisection search until UUm is within tol=1d-6 of 0.3, location of right layer
        do while ((abs(UUm - 0.3d0) .gt. 1d-6))

          if ( ((UUa-0.3d0)*((UUm-0.3d0)) .gt. 0.0d0)) then
              !a and m are on the same side of the root 
              a = m
              UUa = UUm
          else
              !b and m are on the same side of the root 
              b = m
              UUb = UUm
          end if

          ! Compute midpoint of [a,b].
          m = a + (b-a)/2.0d0

	  ! Evaluate the solution at m. 
          xsol(1)=m
          call values(sol%kcol, xsol, nint, x, sol%npde, npts, nderiv, usol, ub, work)
          UUm = usol(1)

        end do

        ! xR is midpoint of current accepted interval.
        xR = m

        ! We next assign the gstop value to be the difference of the absolute value of 
        ! the difference between the locations of the left and right layers and 0.005.
	! Initially xL and xR will be relatively far apart but as the layers
	! approach each other, eventually we will have xL and xR within 0.005
	! of each other.
        rval(1) = abs(xL-xR) - upnt(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
    ! In this example, all references to BACOLI95 are meant to be
    ! understood as referring to BACOLIKR95. Likewise BACOLI
    ! references actually refer to BACOLIKR.

    ! 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, Number of output points, Max number
    ! of subintervals allowed in the spatial mesh, Number of
    ! spatial derivatives of solution that will be output
    integer,  parameter    :: npde = 1, nout = 11, nint_max = 500, nderiv = 0

    ! Terminal time value = 1
    real(dp), parameter :: tout = 1.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

    common /burger/ eps
    data eps /1d-3/

!   Subroutines that define the rhs of the PDE, the boundary conditions,
!   and the initial condition.
    external f, bndxa, bndxb, uinit
!-----------------------------------------------------------------------------

    ! Set the number of roots
    nrt = 1

    ! Set tolerance
    atol(1) = 1d-6
    rtol(1) = atol(1)
 
    ! Initialize the sol structured type: Allocate storage and set algorithm 
    ! parameters. 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.
    call bacoli95(sol, tout, f, bndxa, bndxb, uinit, rt=rt)
    ! Alternative call that uses derivative subroutines
    ! call bacoli95(sol, tout, f, bndxa, bndxb, uinit, derivf, difbxa, difbxb, 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 a 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'("Successful return with idid=",i5)',sol%idid
            write(*,*) 'Root found at t = ', sol%t0
            write(*,*) 'Output solution'

            ! Set output points and evaluate solution at these points
            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


    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)
end program

