MODULE Step

  USE bellhopMod
  USE sspMod
  IMPLICIT NONE
CONTAINS

  SUBROUTINE Step2D( ray0, ray2, Topx, Topn, Botx, Botn )

    ! Does a single step along the ray
    ! x denotes the ray coordinate, ( r, z )
    ! t denotes the scaled tangent to the ray (previously ( rho, zeta ) )
    ! c * t would be the unit tangent

    TYPE( ray2DPt ), INTENT( IN  ) :: ray0
    TYPE( ray2DPt )                :: ray1
    TYPE( ray2DPt ), INTENT( OUT ) :: ray2
    REAL ( KIND=8 ), INTENT( IN  ) :: Topx( 2 ), Topn( 2 ), Botx( 2 ), Botn( 2 )
    INTEGER         :: iSegz0, iSegr0
    REAL ( KIND=8 ) :: gradc0( 2 ), gradc1( 2 ), gradc2( 2 ), rho, &
         c0, cimag0, csq0, crr0, crz0, czz0, cnn0_csq0, &
         c1, cimag1, csq1, crr1, crz1, czz1, cnn1_csq1, &
         c2, cimag2,       crr2, crz2, czz2, &
         urayt0( 2 ), urayt1( 2 ), &
         h, halfh, hw0, hw1, w0, w1, RM, RN, &
         ray2n( 2 ), gradcjump( 2 ), cnjump, csjump

    ! write( *, * ) ray0%x, ray0%p, ray0%q

    ! The numerical integrator used here is a version of the polygon (a.k.a. midpoint, leapfrog, or Box method), and similar
    ! to the Heun (second order Runge-Kutta method).
    ! However, it's modified to allow for a dynamic step change, while preserving the second-order accuracy).

    ! *** Phase 1 (an Euler step)

    CALL EvaluateSSP( ray0%x, c0, cimag0, gradc0, crr0, crz0, czz0, rho, freq, 'TAB' )

    csq0      = c0 * c0
    cnn0_csq0 = crr0 * ray0%t( 2 )**2 - 2.0 * crz0 * ray0%t( 1 ) * ray0%t( 2 ) + czz0 * ray0%t( 1 )**2
    urayt0    = c0 * ray0%t   ! unit tangent to ray

    iSegz0    = iSegz         ! make note of current layer
    iSegr0    = iSegr

    h = Beam%deltas           ! initially set the step h, to the basic one, deltas

    CALL ReduceStep2D( ray0%x, urayt0, iSegz0, iSegr0, Topx, Topn, Botx, Botn, h ) ! reduce h to land on boundary
    halfh = 0.5 * h   ! first step of the modified polygon method is a half step

    ray1%x = ray0%x + halfh * urayt0
    ray1%t = ray0%t - halfh * gradc0 / csq0
    ray1%p = ray0%p - halfh * cnn0_csq0 * ray0%q
    ray1%q = ray0%q + halfh * c0        * ray0%p

    ! write( *, * ) ray1%x, ray1%p, ray1%q

    ! *** Phase 2

    CALL EvaluateSSP( ray1%x, c1, cimag1, gradc1, crr1, crz1, czz1, rho, freq, 'TAB' )
    csq1      = c1 * c1
    cnn1_csq1 = crr1 * ray1%t( 2 )**2 - 2.0 * crz1 * ray1%t( 1 ) * ray1%t( 2 ) + czz1 * ray1%t( 1 )**2
    urayt1    = c1 * ray1%t   ! unit tangent to ray

    ! The Munk test case with a horizontally launched ray caused problems.
    ! The ray vertexes on an interface and can ping-pong around that interface.
    ! Have to be careful in that case about big changes to the stepsize (that invalidate the leap-frog scheme) in phase II.
    ! A modified Heun or Box method could also work.

    CALL ReduceStep2D( ray0%x, urayt1, iSegz0, iSegr0, Topx, Topn, Botx, Botn, h ) ! reduce h to land on boundary

    ! use blend of f' based on proportion of a full step used.
    w1  = h / ( 2.0d0 * halfh )
    w0  = 1.0d0 - w1
    hw0 = h * w0
    hw1 = h * w1

    ray2%x   = ray0%x   + hw0 * urayt0                      + hw1 * urayt1
    ray2%t   = ray0%t   - hw0 * gradc0 / csq0               - hw1 * gradc1 / csq1
    ray2%tau = ray0%tau + hw0 / CMPLX( c0, cimag0, KIND=8 ) + hw1 / CMPLX( c1, cimag1, KIND=8 )
    ray2%p   = ray0%p   - hw0 * cnn0_csq0 * ray0%q          - hw1 * cnn1_csq1 * ray1%q
    ray2%q   = ray0%q   + hw0 * c0        * ray0%p          + hw1 * c1        * ray1%p

    ray2%Amp       = ray0%Amp
    ray2%Phase     = ray0%Phase
    ray2%NumTopBnc = ray0%NumTopBnc
    ray2%NumBotBnc = ray0%NumBotBnc

    ! If we crossed an interface, apply jump condition

    CALL EvaluateSSP( ray2%x, c2, cimag2, gradc2, crr2, crz2, czz2, rho, freq, 'TAB' )
    ray2%c = c2

    IF ( iSegz /= iSegz0 .OR. iSegr /= iSegr0 ) THEN
      !write( *, * ) 'new seg', iSegr0, iSegz0, iSegr, iSegz
      SELECT CASE ( SSP%Type )     ! is their a discontinuity in the first derivative?
       CASE ( 'N', 'C', 'Q', 'H' )  ! N2-linear or C-linear profile option
          gradcjump =  gradc2 - gradc0  ! this is precise only for c-linear layers
          ray2n     = [ -ray2%t( 2 ), ray2%t( 1 ) ]   ! ray normal

          cnjump    = DOT_PRODUCT( gradcjump, ray2n  )
          csjump    = DOT_PRODUCT( gradcjump, ray2%t )

          !!! what if we cross isegr or isegz at the same time?
          IF ( iSegz /= iSegz0 ) THEN         ! crossing in depth
             RM = +ray2%t( 1 ) / ray2%t( 2 )  ! this is tan( alpha ) where alpha is the angle of incidence
          ELSE                                ! crossing in range
             RM = -ray2%t( 2 ) / ray2%t( 1 )  ! this is tan( alpha ) where alpha is the angle of incidence
          END IF

          RN     = RM * ( 2 * cnjump - RM * csjump ) / c2
          ray2%p = ray2%p - ray2%q * RN
       CASE ( 'P', 'S' )            ! monotone PCHIP ACS or Cubic spline profile option
          gradcjump = 0.0
       CASE ( 'A' )                 ! Analytic profile option
          gradcjump = 0.0
       CASE DEFAULT
          WRITE( PRTFile, * ) 'Profile option: ', SSP%Type
          CALL ERROUT( 'BELLHOP: Step2D', 'Invalid profile option' )
       END SELECT

    END IF

    ! write( *, * ) ray2%x, ray2%p, ray2%q

  END SUBROUTINE Step2D

  ! **********************************************************************!

  SUBROUTINE ReduceStep2D( x0, urayt, iSegz0, iSegr0, Topx, Topn, Botx, Botn, h )

    ! calculate a reduced step size, h, that lands on any points where the environment changes

    USE BdryMod
    INTEGER,       INTENT( IN    ) :: iSegz0, iSegr0                             ! SSP layer the ray is in
    REAL (KIND=8), INTENT( IN    ) :: x0( 2 ), urayt( 2 )                        ! ray coordinate and tangent
    REAL (KIND=8), INTENT( IN    ) :: Topx( 2 ), Topn( 2 ), Botx( 2 ), Botn( 2 ) ! Top, bottom coordinate and normal
    REAL (KIND=8), INTENT( INOUT ) :: h                                          ! reduced step size 
    REAL (KIND=8)                  :: hInt, hTop, hBot, hSeg, hBoxr, hBoxz       ! step sizes
    REAL (KIND=8)                  :: x( 2 ), d( 2 ), d0( 2 ), rSeg( 2 )

    ! Detect interface or boundary crossing and reduce step, if necessary, to land on that crossing.
    ! Keep in mind possibility that user put source right on an interface
    ! and that multiple events can occur (crossing interface, top, and bottom in a single step).

    x = x0 + h * urayt ! make a trial step

    ! interface crossing in depth
    hInt = huge( hInt )
    IF ( ABS( urayt( 2 ) ) > EPSILON( hInt ) ) THEN
       IF        ( SSP%z( iSegz0     ) > x(  2 ) ) THEN
          hInt = ( SSP%z( iSegz0     ) - x0( 2 ) ) / urayt( 2 )
       ELSE IF   ( SSP%z( iSegz0 + 1 ) < x(  2 ) ) THEN
          hInt = ( SSP%z( iSegz0 + 1 ) - x0( 2 ) ) / urayt( 2 )
       END IF
    END IF

    ! top crossing
    hTop = huge( hTop )
    d    = x - Topx             ! vector from top to ray
    IF ( DOT_PRODUCT( Topn, d ) > EPSILON( hTop ) ) THEN
       d0   = x0 - Topx         ! vector from top    node to ray origin
       hTop = -DOT_PRODUCT( d0, Topn ) / DOT_PRODUCT( urayt, Topn )
    END IF

    ! bottom crossing
    hBot = huge( hBot )
    d    = x - Botx             ! vector from bottom to ray
    IF ( DOT_PRODUCT( Botn, d ) > EPSILON( hBot ) ) THEN
       d0   = x0 - Botx         ! vector from bottom node to ray origin
       hBot = -DOT_PRODUCT( d0, Botn ) / DOT_PRODUCT( urayt, Botn )
    END IF

    ! top or bottom segment crossing in range
    rSeg( 1 ) = MAX( rTopSeg( 1 ), rBotSeg( 1 ) )
    rSeg( 2 ) = MIN( rTopSeg( 2 ), rBotSeg( 2 ) )

    IF ( SSP%Type == 'Q' ) THEN
       rSeg( 1 ) = MAX( rSeg( 1 ), SSP%Seg%r( iSegr0     ) )
       rSeg( 2 ) = MIN( rSeg( 2 ), SSP%Seg%r( iSegr0 + 1 ) )
    END IF

    hSeg = huge( hSeg )
    IF ( ABS( urayt( 1 ) )  > EPSILON( hSeg ) ) THEN
       IF         ( x(  1 ) < rSeg( 1 ) ) THEN
          hSeg = -( x0( 1 ) - rSeg( 1 ) ) / urayt( 1 )
       ELSE IF    ( x(  1 ) > rSeg( 2 ) ) THEN
          hSeg = -( x0( 1 ) - rSeg( 2 ) ) / urayt( 1 )
       END IF
    END IF

    ! ray mask using a box centered at ( 0, 0 )
    hBoxr    = huge( hBoxr )
    hBoxz    = huge( hBoxz )
    
    IF ( ABS( x( 1 ) ) > Beam%Box%r ) hBoxr = ( Beam%Box%r - ABS( x0( 1 ) ) ) / ABS( urayt( 1 ) )
    IF ( ABS( x( 2 ) ) > Beam%Box%z ) hBoxz = ( Beam%Box%z - ABS( x0( 2 ) ) ) / ABS( urayt( 2 ) )
    
    h = MIN( h, hInt, hTop, hBot, hSeg, hBoxr, hBoxz )  ! take limit set by shortest distance to a crossing

    IF ( h < 1.0d-3 * Beam%deltas ) THEN   ! is it taking an infinitesimal step?
       h = 1.0d-3 * Beam%deltas            ! make sure we make some motion
       iSmallStepCtr = iSmallStepCtr + 1   ! keep a count of the number of sequential small steps
       ! IF ( iSmallStepCtr > 10 ) h = Beam%deltas   ! if stuck, take the default step
       ! write( *, * ) 'h used', h, iSmallStepCtr
    ELSE
       iSmallStepCtr = 0                   ! didn't do a small step so reset the counter
    END IF
    
  END SUBROUTINE ReduceStep2D

END MODULE Step
