!==============================================================================
! ThermoFluidCalc — Calculator #19 : Von Neumann Stability Analysis
!==============================================================================
! Physics : Amplification factor G(k*dx) for classical FD schemes.
!           A scheme is stable iff |G(k*dx)| <= 1  for all wave numbers k.
! Reference : Gupta, "Numerical Methods for Engineers", §2.8
!
! Schemes supported (scheme_id):
!   1 = FTCS  Advection           G = 1 - i*sigma*sin(theta)
!   2 = FTCS  Diffusion           G = 1 - 4*r*sin^2(theta/2)
!   3 = Upwind (1st order)        G = 1 - sigma*(1-cos(theta)) - i*sigma*sin(theta)
!   4 = Lax-Friedrichs            G = cos(theta) - i*sigma*sin(theta)
!   5 = Lax-Wendroff              G = 1 - sigma^2*(1-cos(theta)) - i*sigma*sin(theta)
!   6 = Crank-Nicolson Diffusion  G = [1 - r*(1-cos(theta))] / [1 + r*(1-cos(theta))]
!   7 = Implicit Euler Diffusion  G = 1 / [1 + 2*r*(1-cos(theta))]
!   8 = Leapfrog Advection        G = -i*sigma*sin(theta) +/- sqrt(1 - sigma^2*sin^2(theta))
!
! Build:
!   gfortran -O2 -o von_neumann von_neumann.f90
!
! Input (stdin, one line):   scheme_id  sigma  r_diff  npts
! Output (stdout):           structured text  (see below)
!==============================================================================
program von_neumann
  implicit none

  ! --- Double precision kind -------------------------------------------------
  integer, parameter :: dp = selected_real_kind(15, 307)

  ! --- Constants -------------------------------------------------------------
  real(dp), parameter :: PI      = 3.141592653589793238_dp
  real(dp), parameter :: TOL     = 1.0e-10_dp

  ! --- Variables -------------------------------------------------------------
  integer  :: scheme_id, npts, i
  real(dp) :: sigma, r_diff, dtheta, theta
  real(dp) :: g_mag, g_max, theta_gmax
  complex(dp) :: g_complex, g1, g2
  real(dp) :: re_g, im_g, denom
  character(len=60) :: scheme_name, param_label, criterion
  real(dp) :: param_value
  logical  :: stable

  ! --- Read input ------------------------------------------------------------
  read(*,*) scheme_id, sigma, r_diff, npts

  ! Clamp npts
  if (npts < 10)   npts = 10
  if (npts > 10000) npts = 10000

  ! --- Scheme metadata -------------------------------------------------------
  select case (scheme_id)
    case (1)
      scheme_name = 'FTCS Advection'
      param_label = 'sigma (Courant)'
      param_value = sigma
      criterion   = 'Always unstable (|G|>1 for all sigma>0)'
    case (2)
      scheme_name = 'FTCS Diffusion'
      param_label = 'r (diffusion number)'
      param_value = r_diff
      criterion   = 'Stable if r <= 0.5'
    case (3)
      scheme_name = 'Upwind 1st Order'
      param_label = 'sigma (Courant)'
      param_value = sigma
      criterion   = 'Stable if 0 <= sigma <= 1'
    case (4)
      scheme_name = 'Lax-Friedrichs'
      param_label = 'sigma (Courant)'
      param_value = sigma
      criterion   = 'Stable if |sigma| <= 1'
    case (5)
      scheme_name = 'Lax-Wendroff'
      param_label = 'sigma (Courant)'
      param_value = sigma
      criterion   = 'Stable if |sigma| <= 1'
    case (6)
      scheme_name = 'Crank-Nicolson Diffusion'
      param_label = 'r (diffusion number)'
      param_value = r_diff
      criterion   = 'Unconditionally stable'
    case (7)
      scheme_name = 'Implicit Euler Diffusion'
      param_label = 'r (diffusion number)'
      param_value = r_diff
      criterion   = 'Unconditionally stable'
    case (8)
      scheme_name = 'Leapfrog Advection'
      param_label = 'sigma (Courant)'
      param_value = sigma
      criterion   = 'Stable if |sigma| <= 1'
    case default
      write(*,'(A)') 'ERROR=Invalid scheme_id (must be 1-8)'
      stop
  end select

  ! --- Print header ----------------------------------------------------------
  write(*,'(A,A)')        'SCHEME=',      trim(scheme_name)
  write(*,'(A,A)')        'PARAM_LABEL=', trim(param_label)
  write(*,'(A,ES15.8)')   'PARAM_VALUE=', param_value
  write(*,'(A,A)')        'CRITERION=',   trim(criterion)

  ! --- Sweep theta from 0 to PI ---------------------------------------------
  dtheta    = PI / real(npts - 1, dp)
  g_max     = 0.0_dp
  theta_gmax = 0.0_dp

  write(*,'(A)') 'DATA_START'

  do i = 0, npts - 1
    theta = real(i, dp) * dtheta

    select case (scheme_id)

      case (1)  ! FTCS Advection
        re_g = 1.0_dp
        im_g = -sigma * sin(theta)
        g_mag = sqrt(re_g**2 + im_g**2)

      case (2)  ! FTCS Diffusion  (G is real)
        re_g = 1.0_dp - 4.0_dp * r_diff * sin(theta / 2.0_dp)**2
        g_mag = abs(re_g)

      case (3)  ! Upwind 1st order
        re_g = 1.0_dp - sigma * (1.0_dp - cos(theta))
        im_g = -sigma * sin(theta)
        g_mag = sqrt(re_g**2 + im_g**2)

      case (4)  ! Lax-Friedrichs
        re_g = cos(theta)
        im_g = -sigma * sin(theta)
        g_mag = sqrt(re_g**2 + im_g**2)

      case (5)  ! Lax-Wendroff
        re_g = 1.0_dp - sigma**2 * (1.0_dp - cos(theta))
        im_g = -sigma * sin(theta)
        g_mag = sqrt(re_g**2 + im_g**2)

      case (6)  ! Crank-Nicolson Diffusion (G is real)
        denom = 1.0_dp + r_diff * (1.0_dp - cos(theta))
        re_g  = (1.0_dp - r_diff * (1.0_dp - cos(theta))) / denom
        g_mag = abs(re_g)

      case (7)  ! Implicit Euler Diffusion (G is real)
        denom = 1.0_dp + 2.0_dp * r_diff * (1.0_dp - cos(theta))
        re_g  = 1.0_dp / denom
        g_mag = abs(re_g)

      case (8)  ! Leapfrog Advection — two roots, take max |G|
        im_g = -sigma * sin(theta)
        g_complex = cmplx(0.0_dp, im_g, dp)
        ! discriminant under sqrt : 1 - sigma^2 * sin^2(theta)
        re_g = 1.0_dp - sigma**2 * sin(theta)**2
        if (re_g >= 0.0_dp) then
          g1 = cmplx(sqrt(re_g), im_g, dp)
          g2 = cmplx(-sqrt(re_g), im_g, dp)
        else
          ! Complex square root
          g1 = cmplx(0.0_dp, im_g, dp) + sqrt(cmplx(re_g, 0.0_dp, dp))
          g2 = cmplx(0.0_dp, im_g, dp) - sqrt(cmplx(re_g, 0.0_dp, dp))
        end if
        g_mag = max(abs(g1), abs(g2))

    end select

    ! Track maximum
    if (g_mag > g_max) then
      g_max      = g_mag
      theta_gmax = theta
    end if

    write(*,'(F12.8,A,F15.10)') theta, ',', g_mag
  end do

  write(*,'(A)') 'DATA_END'

  ! --- Stability verdict -----------------------------------------------------
  stable = (g_max <= 1.0_dp + TOL)

  write(*,'(A,F15.10)')  'G_MAX=',      g_max
  write(*,'(A,F12.8)')   'THETA_GMAX=', theta_gmax
  if (stable) then
    write(*,'(A)') 'STABLE=YES'
  else
    write(*,'(A)') 'STABLE=NO'
  end if

end program von_neumann
