! 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 
! bacoli95_vals routine from BACOLIKR.

    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 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 mesh 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 time
    ! at which exp(t) * the integral of the approximate solution
    ! takes on a value of 10.0. Thus there is one root and the
    ! gstop vector is integral - 10.0.

    ! This routine accesses sol structured type and 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)            :: approx   ! Integral approximation
        real(dp)            :: int      ! Target value of the integral
        real(dp)            :: u(6500)   ! Solution values at x values
                                        ! Breaks if nint_max > 500
        real(dp)            :: h(500)   ! h values; assumes nint_max <= 500
        real(dp)            :: rho(13)  ! Gauss points per subinterval
        real(dp)            :: wts(13)  ! Gauss weights per subinterval
        real(dp)            :: excol(6500), ewts(6500) 
		                    ! Total set of Gauss points and weights over 
				    ! entire domain; assumes nint_max <= 500.
                
        integer             :: npts     ! Number of evaluation points
        integer, parameter  :: nderiv = 0 
                               ! Number of spatial derivatives of
                               ! solution that are required
        integer             :: i, ii, j        ! loop indexes

	!-----------------------------------------------------------
	! We are looking for the time when the integral of the tumor
        ! concentration is 10. We look for exp(t)*int-5to5 u(x,t)dx= 10.
        ! So the target value is 10.
        int = 10.0d0
	
	!------------------------------------------------------------
		
	! For efficiency, the computation of the Gauss pts and wts should
	! be done once beforehand and then the array can be available
	! globally. For now we leave these computations here.
        ! For the kcol+3 Gaussian Quadrature rule approximation to the
        ! integral we will need to evaluate the approximate solution
        ! at nini*(kcol+3) points, i.e., the kcol+3 Gauss points on 
	! each subinterval. Therefore
        npts = nint*(sol%kcol+3)
		
	! Next we compute the entire set of Gauss points and weights 
	! for the entire spatial domain. The code below is taken from bacoli.f.
		
	! Calculate the subinterval size sequence.
        do i = 1, nint
            h(i) = x(i+1)-x(i)
        end do

        ! Compute the kcol+3 Gaussian points and Gaussian weights on [0,1].
	! (In the call below to the F77 routine, gauleg, the work array
	! needs to be of dimension at least (kcol+3)**2, which has a max
	! of 169 (since the max kcol value is 10). The size of work is
	! set in the next routine below, called setsol. As long as nint_max
	! is large enough, work will be large enough for the call to gauleg.)
	! The points and weights are returned in rho and wts resepctively.
         call gauleg(sol%kcol+3, (sol%kcol+3)*(sol%kcol+3), rho, wts, work, 4)

        ! Define the Gaussian quadrature point and weight sequence over the 
	! whole domain
        do i = 1, nint
            ii = (i - 1) * (sol%kcol + 3)
            do j = 1, sol%kcol+3
                excol(ii + j) = x(i) + h(i) * rho(j)
                ewts(ii + j) = h(i) * wts(j)
            end do
        end do
		
        ! Evaluate the approximate solution to obtain values for the integral.
        ! It turns out that values can compute all of these values
        ! in one call.
 
        ! Here we call 'values' with input vector 'ub' which evaluates
        ! the approximate solution at excol and at the current time.
        ! Replacing 'ub' with 'ubt' would evaluate the time derivative
        ! of the approximate solution at excol and the current 
        ! time. The 'values' routine is an F77 routine that can be
        ! accessed globally.
        call values(sol%kcol, excol, nint, x, sol%npde, npts, nderiv,  &
                                                         u, ub, work)

        ! The values of the approximate solution at x and the current
        ! time are returned in u

        ! Apply the Gauss rule to obtain the integral approximation
        approx = 0.0d0
        do i = 1, nint
            ii = (i - 1) * (sol%kcol + 3)
            do j = 1, sol%kcol+3
                approx = approx + ewts(ii+j)*u(ii+j)
            end do
        end do
        
	!------------------------------------------------------------
        ! Scale result by exp(t) to get integral of tumor cocnetraton
        approx = approx * exp(t)

        ! We next assign the gstop value to be difference between the
        ! approximate integral value and the target integral value. 
        rval(1) = approx - int

    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 BACCOLI
    ! references actually refer to BACOLIKR.

    ! A minimal example driver for the BACOLI95 wrapper of BACOLI.
    ! See the comments within bacoli95.f95 for details regarding use.
    use bacoli95_mod, only: bacoli95_init, bacoli95, bacoli95_vals, &
                            bacoli95_sol, bacoli95_sol_teardown

    use rootfinding, only: rt, setsol

    implicit none

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

    type(bacoli95_sol), target  :: sol

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

    ! Choose xa = -5 and xb = 5 for the brain tumor example
    real(dp), parameter  :: xa = -5, xb = 5

    ! Declare initial mesh for tumor problem to have 16 subintervals
    real(dp)            :: tumorMesh(17)

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

    ! xout(:) contains output points, uout(:) contains solution
    ! approximations at output points
    real(dp) :: xout(nout), uout(nout)

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

    ! nrt is the number of gstop roots to be found
    integer :: nrt

    ! Loop index
    integer :: i

    ! The following parameters are used in the brain tumor example.
    real(dp) :: a, b, w1, w2, gamma, xi, eta
    common /tumor/ a, b, w1, w2, gamma, xi, eta

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

!   Number of gstop roots
    nrt = 1

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

    ! Set problem dependent parameters
    a = xa
    b = xb
    w1 = -0.5_dp
    w2 = 0.5_dp
    gamma = 0.2_dp
    xi = -2.0_dp
    eta = 0.2_dp
    
    ! Initialization of spatial mesh
    tumorMesh = (/-5.0D0,-4.0D0,-3.0D0,-2.0D0,-1.0D0, &
        -0.6D0,-0.5D0,-0.4D0,0.0D0,0.4D0,0.5D0,0.6D0,1.0D0,2.0D0,3.0D0,4.0D0, &
        5.0D0/)
 
    ! Initialize the sol structured type. error estimator is the LOI,
    ! number of collocation points per subinterval, kcol, is 4
    call bacoli95_init(sol, npde, tumorMesh, 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 to tout or stop early if a gstop root is found.

    call bacoli95(sol, tout, f, bndxa, bndxb, uinit, rt=rt)

    ! Stop due to an error condition
    if (sol%idid .le. 0) then
            write(*,*) 'Error? Idid = ', sol%idid
            stop 123
    end if

    ! Successful return from bacoli95. We are stopping at a time
    ! that corresponds to finding a gstop root or at tout if no root 
    ! was found.
    if (sol%idid == 5) then
            print'("done. idid=",i5)',sol%idid
            write(*,*) 'Root found at t = ', sol%t0
    else
            print'("done. idid=",i5)',sol%idid
            write(*,*) 'Root not found before reaching tout'
    end if
 
    ! sol%idid > 0 implies a successful return from bacoli95.
    if (sol%idid > 0) then

        write(*,*) 'Output solution'

        xout = (/xa,(xa+i*(xb-xa)/(nout-1), i=1,nout-1)/)
        call bacoli95_vals(sol, xout, uout)

        write(*,*)

        ! Scale uout by exp(sol%t0) to get brain tumor concentrations
        do i = 1, nout
            uout(i) = uout(i)*exp(sol%t0)
        end do

        print'("at t=",f4.2)', sol%t0
        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

