!==============================================================================
! ThermoFluidCalc — Calculator #21 : Stability Regions
!==============================================================================
! Physics : Maps the |G(z)| = 1 boundary in the complex z = λΔt plane
!           for 6 classical ODE/PDE time-integration schemes.
!
! Schemes:
!   1 = Forward Euler      G = 1 + z
!   2 = Backward Euler     G = 1/(1 - z)
!   3 = Crank-Nicolson     G = (2 + z)/(2 - z)
!   4 = Leapfrog           G = z ± sqrt(z² + 1)  (two roots; max |G|)
!   5 = RK2 (Heun)         G = 1 + z + z²/2
!   6 = RK4 (Classical)    G = 1 + z + z²/2 + z³/6 + z⁴/24
!
! Reference : Gupta §2.8; Hairer & Wanner, "Solving ODEs II"
!
! Build:
!   gfortran -O2 -o stability_regions stability_regions.f90
!
! Input  (stdin):  npts  z_test_re  z_test_im
! Output (stdout): structured data (see below)
!==============================================================================
program stability_regions
  implicit none

  integer, parameter :: dp = selected_real_kind(15, 307)
  real(dp), parameter :: PI = 3.141592653589793238_dp
  integer, parameter :: NSCHEMES = 6
  integer, parameter :: MAX_BISECT = 80

  integer  :: npts, i, s
  real(dp) :: z_test_re, z_test_im
  complex(dp) :: z_test, gz
  real(dp) :: angle, dangle, r_lo, r_hi, r_mid, gm
  real(dp) :: cx, cy                ! center for sweep
  real(dp) :: x, y
  real(dp) :: probe_mag(NSCHEMES)
  logical  :: probe_stable(NSCHEMES)

  character(len=40) :: sname(NSCHEMES)
  character(len=80) :: sformula(NSCHEMES)
  character(len=80) :: sregion(NSCHEMES)

  ! Scheme metadata
  sname(1) = 'Forward Euler (Explicit)'
  sname(2) = 'Backward Euler (Implicit)'
  sname(3) = 'Crank-Nicolson'
  sname(4) = 'Leapfrog (Midpoint)'
  sname(5) = 'RK2 (Heun)'
  sname(6) = 'RK4 (Classical)'

  sformula(1) = 'G = 1 + z'
  sformula(2) = 'G = 1 / (1 - z)'
  sformula(3) = 'G = (2 + z) / (2 - z)'
  sformula(4) = 'G = z +/- sqrt(z^2 + 1)'
  sformula(5) = 'G = 1 + z + z^2/2'
  sformula(6) = 'G = 1 + z + z^2/2 + z^3/6 + z^4/24'

  sregion(1) = 'Disk center(-1,0) radius 1'
  sregion(2) = 'Exterior of disk center(1,0) radius 1'
  sregion(3) = 'Left half-plane Re(z) <= 0'
  sregion(4) = 'Imaginary segment [-i, i]'
  sregion(5) = 'Closed region (computed numerically)'
  sregion(6) = 'Closed region (computed numerically)'

  ! ── Read input ────────────────────────────────────────────────────────────
  read(*,*) npts, z_test_re, z_test_im
  if (npts < 40)   npts = 40
  if (npts > 4000)  npts = 4000
  z_test = cmplx(z_test_re, z_test_im, dp)

  dangle = 2.0_dp * PI / real(npts, dp)

  ! ── Loop over each scheme ─────────────────────────────────────────────────
  do s = 1, NSCHEMES

    ! --- Probe point evaluation ---
    gz = eval_G(s, z_test)
    probe_mag(s) = abs(gz)
    probe_stable(s) = (probe_mag(s) <= 1.0_dp + 1.0e-10_dp)

    ! --- Header for this scheme ---
    write(*,'(A,I1,A,A)')  'SCHEME_', s, '=', trim(sname(s))
    write(*,'(A,I1,A,A)')  'FORMULA_', s, '=', trim(sformula(s))
    write(*,'(A,I1,A,A)')  'REGION_', s, '=', trim(sregion(s))
    write(*,'(A,I1,A,ES15.8)') 'PROBE_', s, '=', probe_mag(s)
    if (probe_stable(s)) then
      write(*,'(A,I1,A)') 'STABLE_', s, '=YES'
    else
      write(*,'(A,I1,A)') 'STABLE_', s, '=NO'
    end if

    ! --- Contour data ---
    write(*,'(A,I1,A)') 'CONTOUR_', s, '_START'

    select case (s)

    case (1) ! Forward Euler: unit circle centered at (-1, 0)
      do i = 0, npts
        angle = real(i, dp) * dangle
        x = -1.0_dp + cos(angle)
        y = sin(angle)
        write(*,'(F12.8,A,F12.8)') x, ',', y
      end do

    case (2) ! Backward Euler: unit circle centered at (1, 0) — boundary
      do i = 0, npts
        angle = real(i, dp) * dangle
        x = 1.0_dp + cos(angle)
        y = sin(angle)
        write(*,'(F12.8,A,F12.8)') x, ',', y
      end do

    case (3) ! Crank-Nicolson: imaginary axis  x = 0
      do i = 0, npts
        y = -4.0_dp + 8.0_dp * real(i, dp) / real(npts, dp)
        write(*,'(F12.8,A,F12.8)') 0.0_dp, ',', y
      end do

    case (4) ! Leapfrog: segment on imaginary axis from -i to i
      do i = 0, npts
        y = -1.0_dp + 2.0_dp * real(i, dp) / real(npts, dp)
        write(*,'(F12.8,A,F12.8)') 0.0_dp, ',', y
      end do

    case (5, 6) ! RK2 or RK4: numerical boundary via bisection sweep
      ! Sweep angle from center, find |G|=1 boundary by bisection
      cx = -1.0_dp
      cy = 0.0_dp
      do i = 0, npts
        angle = real(i, dp) * dangle
        ! Bisection: find r such that |G(center + r*e^{i*angle})| = 1
        r_lo = 0.0_dp
        r_hi = 5.0_dp  ! generous upper bound
        ! First check if boundary exists along this ray
        x = cx + r_hi * cos(angle)
        y = cy + r_hi * sin(angle)
        gm = abs(eval_G(s, cmplx(x, y, dp)))
        if (gm <= 1.0_dp) then
          ! Entire ray is stable up to r_hi, write the outer point
          write(*,'(F12.8,A,F12.8)') x, ',', y
          cycle
        end if
        ! Bisection
        call bisect_boundary(s, cx, cy, angle, r_lo, r_hi, r_mid)
        x = cx + r_mid * cos(angle)
        y = cy + r_mid * sin(angle)
        write(*,'(F12.8,A,F12.8)') x, ',', y
      end do

    end select

    write(*,'(A,I1,A)') 'CONTOUR_', s, '_END'
  end do

contains

  !------------------------------------------------------------------------
  ! Evaluate G(z) for scheme s
  !------------------------------------------------------------------------
  function eval_G(s, z) result(gz)
    integer, intent(in)    :: s
    complex(dp), intent(in) :: z
    complex(dp) :: gz, g1, g2, disc

    select case (s)
    case (1)  ! Forward Euler
      gz = cmplx(1.0_dp, 0.0_dp, dp) + z

    case (2)  ! Backward Euler
      gz = cmplx(1.0_dp, 0.0_dp, dp) / (cmplx(1.0_dp, 0.0_dp, dp) - z)

    case (3)  ! Crank-Nicolson
      gz = (cmplx(2.0_dp, 0.0_dp, dp) + z) / (cmplx(2.0_dp, 0.0_dp, dp) - z)

    case (4)  ! Leapfrog: two roots, return max |G|
      disc = z*z + cmplx(1.0_dp, 0.0_dp, dp)
      disc = sqrt(disc)
      g1 = z + disc
      g2 = z - disc
      if (abs(g1) >= abs(g2)) then
        gz = g1
      else
        gz = g2
      end if

    case (5)  ! RK2
      gz = cmplx(1.0_dp,0.0_dp,dp) + z + z*z/cmplx(2.0_dp,0.0_dp,dp)

    case (6)  ! RK4
      gz = cmplx(1.0_dp,0.0_dp,dp) + z &
         + z*z / cmplx(2.0_dp,0.0_dp,dp) &
         + z*z*z / cmplx(6.0_dp,0.0_dp,dp) &
         + z*z*z*z / cmplx(24.0_dp,0.0_dp,dp)

    case default
      gz = cmplx(0.0_dp, 0.0_dp, dp)
    end select
  end function eval_G

  !------------------------------------------------------------------------
  ! Bisection: find radius r from (cx,cy) at angle where |G| crosses 1
  !------------------------------------------------------------------------
  subroutine bisect_boundary(s, cx, cy, angle, r_lo_in, r_hi_in, r_out)
    integer,  intent(in)  :: s
    real(dp), intent(in)  :: cx, cy, angle, r_lo_in, r_hi_in
    real(dp), intent(out) :: r_out
    real(dp) :: rlo, rhi, rmid, xm, ym, gm
    complex(dp) :: zm
    integer :: iter

    rlo = r_lo_in
    rhi = r_hi_in
    do iter = 1, MAX_BISECT
      rmid = 0.5_dp * (rlo + rhi)
      xm = cx + rmid * cos(angle)
      ym = cy + rmid * sin(angle)
      zm = cmplx(xm, ym, dp)
      gm = abs(eval_G(s, zm))
      if (gm <= 1.0_dp) then
        rlo = rmid
      else
        rhi = rmid
      end if
      if (abs(rhi - rlo) < 1.0e-12_dp) exit
    end do
    r_out = 0.5_dp * (rlo + rhi)
  end subroutine bisect_boundary

end program stability_regions