!=======================================================================
! Module: Constants and Types
! Description: Defines constants, derived types for problem settings
!              and solution data.
!=======================================================================
module constants_and_types
  implicit none
  private

  ! -- Precision definition
  integer, parameter :: dp = kind(1.0d0) ! Double precision

  ! -- Mathematical constants
  real(dp), parameter :: PI = 4.0_dp * atan(1.0_dp)

  ! -- Derived type for problem settings
  type :: problem_settings_t
    integer             :: n_dims        ! Number of spatial dimensions (N)
    integer             :: n_grid        ! Number of grid points per dimension (M)
    real(dp)            :: domain_size   ! Physical size of the domain (e.g., L for [0, L]^N)
    real(dp)            :: dt            ! Time step size
    real(dp)            :: dx            ! Spatial grid spacing (dx = domain_size / n_grid)
    real(dp)            :: kappa         ! Diffusion coefficient
    real(dp)            :: alpha         ! Logistic growth coefficient
    real(dp)            :: beta          ! Sigmoid term coefficient
    real(dp), allocatable :: weights(:)  ! Weight vector w (size n_dims)
    real(dp)            :: bias          ! Bias b
    character(len=20)   :: boundary_cond ! Boundary condition type ('periodic', 'fixed', etc.)
    ! Add other parameters as needed (e.g., total simulation time)
  end type problem_settings_t

  ! -- Derived type for solution data
  type :: solution_data_t
    ! N-dimensional array for the solution u at current time step
    ! The size will be (n_grid, n_grid, ..., n_grid) - N times
    real(dp), allocatable, dimension(:, :, :, :, :, :, :) :: u ! Max 7 dimensions Fortran standard
    ! For N > 7, need different data structure or Fortran 2008+ feature
    ! For simplicity, let's assume N <= 7 or use a flattened 1D array with index calculation
    ! Example using a flattened array (more general for N):
    ! real(dp), allocatable :: u_flat(:)
    ! integer*8             :: total_grid_points
  end type solution_data_t

  public :: dp, PI, problem_settings_t, solution_data_t

end module constants_and_types

!=======================================================================
! Module: PDE Solver Utilities
! Description: Contains utility functions and subroutines for the PDE solver.
!=======================================================================
module pde_solver_utils
  use constants_and_types, only: dp
  implicit none
  private

  public :: sigmoid, initialize_solution, apply_boundary_conditions, compute_laplacian

contains

  !---------------------------------------------------------------------
  ! Function: sigmoid
  ! Description: Computes the standard sigmoid function.
  ! Arguments:
  !   z (input, real(dp)): Input value
  ! Returns:
  !   real(dp): Sigmoid function value 1 / (1 + exp(-z))
  !---------------------------------------------------------------------
  function sigmoid(z) result(s)
    real(dp), intent(in) :: z
    real(dp) :: s
    s = 1.0_dp / (1.0_dp + exp(-z))
  end function sigmoid

  !---------------------------------------------------------------------
  ! Subroutine: initialize_solution
  ! Description: Initializes the solution field u.
  ! Arguments:
  !   settings (input): Problem settings
  !   solution (inout): Solution data structure to initialize
  ! Notes:
  !   Currently initializes with small random values. Modify for specific initial conditions.
  !   Handles allocation of the solution array based on settings.
  !---------------------------------------------------------------------
  subroutine initialize_solution(settings, solution)
    type(problem_settings_t), intent(in) :: settings
    type(solution_data_t), intent(inout) :: solution
    integer :: i, j, k, l, m, n, p, stat
    integer, dimension(settings%n_dims) :: shape
    real(dp) :: rand_val

    shape = settings%n_grid ! Array shape: (M, M, ..., M) N times

    ! Allocate the solution array u
    ! Fortran standard up to 7 dimensions directly.
    ! For N > 7, consider using a 1D array or Fortran 2008 coarrays/allocatable components.
    if (settings%n_dims > 7) then
      print *, "Error: N_DIMS > 7 not directly supported by this array structure."
      stop 1
    end if

    ! Create a dynamic allocation based on n_dims
    select case (settings%n_dims)
    case (1)
      allocate(solution%u(shape(1)), stat=stat)
    case (2)
      allocate(solution%u(shape(1), shape(2)), stat=stat)
    case (3)
      allocate(solution%u(shape(1), shape(2), shape(3)), stat=stat)
    case (4)
      allocate(solution%u(shape(1), shape(2), shape(3), shape(4)), stat=stat)
    case (5)
      allocate(solution%u(shape(1), shape(2), shape(3), shape(4), shape(5)), stat=stat)
    case (6)
      allocate(solution%u(shape(1), shape(2), shape(3), shape(4), shape(5), shape(6)), stat=stat)
    case (7)
      allocate(solution%u(shape(1), shape(2), shape(3), shape(4), shape(5), shape(6), shape(7)), stat=stat)
    case default
      print *, "Error: Invalid N_DIMS for allocation."
      stop 1
    end select

    if (stat /= 0) then
      print *, "Error: Could not allocate solution array u."
      stop 1
    end if

    ! Initialize with small random numbers (example)
    ! TODO: Implement specific initial condition logic here
    call random_seed()
    call random_number(solution%u)
    solution%u = (solution%u - 0.5_dp) * 0.1_dp + 0.5_dp ! Small variation around 0.5

    ! Allocate and initialize weights (example: random weights)
    if (.not. allocated(settings%weights)) then
       allocate(settings%weights(settings%n_dims), stat=stat)
       if (stat /= 0) then
          print *, "Error: Could not allocate weights array."
          stop 1
       end if
    end if
    call random_number(settings%weights)
    settings%weights = settings%weights * 2.0_dp - 1.0_dp ! Weights in [-1, 1]

  end subroutine initialize_solution

  !---------------------------------------------------------------------
  ! Subroutine: apply_boundary_conditions
  ! Description: Applies boundary conditions to the solution field u.
  ! Arguments:
  !   settings (input): Problem settings (contains boundary type)
  !   u        (inout): Solution array (N-dimensional)
  ! Notes:
  !   Implements periodic boundary conditions as an example.
  !   Requires careful handling of N-dimensional array indices.
  !---------------------------------------------------------------------
  subroutine apply_boundary_conditions(settings, u)
    type(problem_settings_t), intent(in) :: settings
    real(dp), dimension(:, :, :, :, :, :, :), intent(inout) :: u ! Assumes N <= 7
    integer :: dim, i, M
    integer, dimension(settings%n_dims) :: idx, low_idx, high_idx, shape

    M = settings%n_grid
    shape = shape(u) ! Get the actual shape

    if (trim(settings%boundary_cond) == 'periodic') then
      ! Loop through each dimension to apply periodic BC
      do dim = 1, settings%n_dims
        ! Create slices for the boundaries in the current dimension
        ! This is complex for general N. Fortran's array slicing syntax helps.
        ! Example for N=2:
        ! u(1,:) = u(M+1,:) ! Ghost cell left = right boundary (assuming M+2 points)
        ! u(M+2,:) = u(2,:) ! Ghost cell right = left boundary
        ! u(:,1) = u(:,M+1)
        ! u(:,M+2) = u(:,2)
        ! For N dimensions and M grid points (indices 1 to M), periodic means:
        ! index 0 corresponds to index M
        ! index M+1 corresponds to index 1
        ! We need ghost cells or careful indexing at boundaries (1 and M)
        ! Here, we implement wrap-around directly within the Laplacian calculation later.
        ! This subroutine might be better used for fixed value (Dirichlet) BCs.
        ! For now, leave this empty and handle periodicity in Laplacian.

      end do
    else if (trim(settings%boundary_cond) == 'fixed') then
      ! TODO: Implement fixed boundary conditions (Dirichlet)
      ! Set boundary cells to a specific value.
      print *, "Warning: Fixed boundary conditions not implemented yet."
    else
      print *, "Error: Unknown boundary condition type:", trim(settings%boundary_cond)
      stop 1
    end if

  end subroutine apply_boundary_conditions


  !---------------------------------------------------------------------
  ! Function: compute_laplacian
  ! Description: Computes the Laplacian term using N-dimensional central differences.
  !              Handles periodic boundary conditions internally.
  ! Arguments:
  !   settings (input): Problem settings (dx, n_grid, n_dims)
  !   u        (input): Solution array at current time step
  !   idx      (input): Integer array of size n_dims, representing the current grid index (i1, i2, ..., iN)
  ! Returns:
  !   real(dp): Approximate value of Laplacian at grid point idx
  !---------------------------------------------------------------------
   function compute_laplacian(settings, u, idx) result(lap_val)
     type(problem_settings_t), intent(in) :: settings
     real(dp), dimension(:, :, :, :, :, :, :), intent(in) :: u ! Assuming N <= 7
     integer, dimension(settings%n_dims), intent(in) :: idx
     real(dp) :: lap_val
     integer :: k, M
     integer, dimension(settings%n_dims) :: idx_plus, idx_minus
     real(dp) :: u_center, u_plus, u_minus
     real(dp) :: dx2_inv

     M = settings%n_grid
     dx2_inv = 1.0_dp / (settings%dx * settings%dx)
     lap_val = 0.0_dp

     ! Get the value at the center point (idx)
     ! This requires a way to access u dynamically based on the array 'idx'.
     ! Fortran doesn't directly support u(idx(1), idx(2), ...) easily.
     ! A helper function or careful select case might be needed.
     ! Let's assume a helper function `get_u_val(u, idx)` exists for now.
     u_center = get_u_val(u, idx)

     ! Sum contributions from each dimension
     do k = 1, settings%n_dims
       idx_plus = idx
       idx_minus = idx

       ! Indices for finite difference stencil, handling periodic BC
       idx_plus(k) = idx(k) + 1
       if (idx_plus(k) > M) idx_plus(k) = 1 ! Wrap around (Periodic BC)

       idx_minus(k) = idx(k) - 1
       if (idx_minus(k) < 1) idx_minus(k) = M ! Wrap around (Periodic BC)

       ! Get values at neighboring points
       u_plus = get_u_val(u, idx_plus)
       u_minus = get_u_val(u, idx_minus)

       ! Add contribution from dimension k
       lap_val = lap_val + (u_plus - 2.0_dp * u_center + u_minus)
     end do

     lap_val = lap_val * dx2_inv

   contains
     ! Helper function to get value from N-dim array using index vector
     ! Needs to be adapted based on how u is stored (N-dim or 1D flat)
     ! This is a placeholder - needs actual implementation.
     function get_u_val(arr, index_vec) result(val)
       real(dp), dimension(:, :, :, :, :, :, :), intent(in) :: arr
       integer, dimension(:), intent(in) :: index_vec
       real(dp) :: val
       integer :: n
       n = size(index_vec)
       ! This requires SELECT CASE based on n or direct indexing if N is fixed
       select case(n)
       case(1); val = arr(index_vec(1))
       case(2); val = arr(index_vec(1), index_vec(2))
       case(3); val = arr(index_vec(1), index_vec(2), index_vec(3))
       ! ... up to 7
       case default; val = 0.0_dp; print *, "Error in get_u_val" ! Error or stop
       end select
     end function get_u_val

   end function compute_laplacian

end module pde_solver_utils


!=======================================================================
! Module: PDE Solver Core
! Description: Implements the main time-stepping loop for the PDE.
!=======================================================================
module pde_solver_core
  use constants_and_types, only: dp, problem_settings_t, solution_data_t
  use pde_solver_utils, only: sigmoid, compute_laplacian, apply_boundary_conditions
  implicit none
  private

  public :: time_step

contains

  !---------------------------------------------------------------------
  ! Subroutine: time_step
  ! Description: Performs one time step using Forward Euler method.
  ! Arguments:
  !   settings (input)    : Problem settings
  !   solution (inout)    : Solution data (contains u)
  !   u_new    (out, allocatable): Array to store the updated solution u^{n+1}
  !---------------------------------------------------------------------
  subroutine time_step(settings, solution, u_new)
    type(problem_settings_t), intent(in) :: settings
    type(solution_data_t), intent(inout) :: solution ! Contains u^n
    real(dp), dimension(:, :, :, :, :, :, :), allocatable, intent(out) :: u_new ! For u^{n+1}
    real(dp) :: lap_term, logistic_term, sigmoid_term, dot_product
    integer :: i, k, stat
    integer, dimension(settings%n_dims) :: current_idx
    integer, dimension(settings%n_dims) :: shape
    real(dp) :: current_u, current_x_k

    ! Allocate u_new with the same shape as solution%u
    shape = shape(solution%u)
    allocate(u_new(shape(1):shape(1)+size(solution%u,1)-1, &
                   shape(2):shape(2)+size(solution%u,2)-1, &
                   shape(3):shape(3)+size(solution%u,3)-1, &
                   shape(4):shape(4)+size(solution%u,4)-1, &
                   shape(5):shape(5)+size(solution%u,5)-1, &
                   shape(6):shape(6)+size(solution%u,6)-1, &
                   shape(7):shape(7)+size(solution%u,7)-1 ), stat=stat)

    ! Dynamic allocation based on n_dims (similar to initialization)
    select case (settings%n_dims)
      case (1); allocate(u_new(settings%n_grid), stat=stat)
      case (2); allocate(u_new(settings%n_grid, settings%n_grid), stat=stat)
      case (3); allocate(u_new(settings%n_grid, settings%n_grid, settings%n_grid), stat=stat)
      ! ... up to 7
      case default; print *, "Error: Invalid N_DIMS for u_new allocation."; stop 1
    end select
    if (stat /= 0) then
      print *, "Error: Could not allocate u_new array."
      stop 1
    end if

    ! Apply boundary conditions to u^n before computing derivatives
    ! (Especially important for non-periodic BCs, might involve ghost cells)
    ! call apply_boundary_conditions(settings, solution%u) ! May modify solution%u if ghost cells are used

    ! Iterate over all grid points (i1, i2, ..., iN)
    ! This requires N nested loops, which is inefficient and hard to write generally.
    ! A recursive subroutine or an iterative approach with an index counter is better.

    ! Iterative approach using a single counter and index calculation:
    ! integer*8 :: grid_idx_flat, total_points
    ! total_points = int(settings%n_grid, kind=8) ** settings%n_dims
    ! do grid_idx_flat = 1, total_points
    !    call get_ndim_index(grid_idx_flat, settings%n_grid, settings%n_dims, current_idx)
    !    ... calculate rhs ...
    ! end do

    ! Let's try nested loops for N=3 for illustration (replace with general method)
    if (settings%n_dims /= 3) then
       print *, "Error: This explicit loop structure is for N=3 only."
       print *, "       Implement a general N-dimensional loop."
       stop 1
    end if

    do k = 1, settings%n_grid ! Dimension 3
      do j = 1, settings%n_grid ! Dimension 2
        do i = 1, settings%n_grid ! Dimension 1
          current_idx = [i, j, k]

          ! Get value u^n at current point
          current_u = solution%u(i, j, k) ! Assumes N=3

          ! 1. Compute Laplacian term at (i,j,k) using u^n
          lap_term = compute_laplacian(settings, solution%u, current_idx)

          ! 2. Compute logistic term
          logistic_term = settings%alpha * current_u * (1.0_dp - current_u)

          ! 3. Compute sigmoid term
          ! Calculate w . x + b
          dot_product = settings%bias
          do dim_idx = 1, settings%n_dims
             current_x_k = dble(current_idx(dim_idx)) * settings%dx ! Physical coordinate x_k
             dot_product = dot_product + settings%weights(dim_idx) * current_x_k
          end do
          sigmoid_term = settings%beta * sigmoid(dot_product)

          ! 4. Compute u^{n+1} using Forward Euler
          u_new(i, j, k) = current_u + settings%dt * ( &
                             settings%kappa * lap_term + &
                             logistic_term - &
                             sigmoid_term )

        end do ! dim 1 (i)
      end do ! dim 2 (j)
    end do ! dim 3 (k)


    ! Note: For higher N, the nested loops become impractical. Use a recursive
    ! subroutine or an iterative approach mapping a single loop index to N-D indices.

    ! After computing all points in u_new, update the main solution array
    ! solution%u = u_new ! This copies the entire array

    ! Deallocate u_new if it's not needed outside this subroutine scope
    ! deallocate(u_new) ! Or pass it out and deallocate later

  end subroutine time_step

end module pde_solver_core


!=======================================================================
! Main Program: PDE Simulation Driver
! Description: Sets up the problem, runs the simulation loop,
!              and handles output.
!=======================================================================
program main_pde_solver
  use constants_and_types, only: dp, problem_settings_t, solution_data_t
  use pde_solver_utils, only: initialize_solution
  use pde_solver_core, only: time_step
  implicit none

  type(problem_settings_t) :: settings
  type(solution_data_t) :: solution
  real(dp), allocatable, dimension(:, :, :, :, :, :, :) :: u_next_step ! For N <= 7
  integer :: n, n_steps, output_interval
  real(dp) :: t_current, t_final

  ! -- 1. Setup Problem Parameters --
  settings%n_dims        = 2       ! <<<<< Number of dimensions N (e.g., 2 or 3)
  settings%n_grid        = 50      ! <<<<< Grid points per dimension M
  settings%domain_size   = 10.0_dp ! <<<<< Physical domain size L ([0,L]^N)
  settings%dt            = 0.0001_dp ! <<<<< Time step (Adjust for stability!)
  settings%kappa         = 1.0_dp  ! <<<<< Diffusion coefficient
  settings%alpha         = 5.0_dp  ! <<<<< Logistic term coefficient
  settings%beta          = 1.0_dp  ! <<<<< Sigmoid term coefficient
  settings%bias          = 0.0_dp  ! <<<<< Sigmoid bias b
  settings%boundary_cond = 'periodic' ! Boundary condition type

  t_final           = 1.0_dp  ! <<<<< Total simulation time
  output_interval   = 100     ! <<<<< Output every N steps

  ! Calculate derived parameters
  settings%dx = settings%domain_size / dble(settings%n_grid)
  n_steps = int(t_final / settings%dt)

  ! Print simulation parameters
  print *, "==========================================="
  print *, " PDE Simulation Settings"
  print *, "==========================================="
  print *, " Dimensions (N): ", settings%n_dims
  print *, " Grid points/dim (M): ", settings%n_grid
  print *, " Domain size (L): ", settings%domain_size
  print *, " Grid spacing (dx): ", settings%dx
  print *, " Time step (dt): ", settings%dt
  print *, " Final time: ", t_final
  print *, " Total steps: ", n_steps
  print *, " Kappa: ", settings%kappa
  print *, " Alpha: ", settings%alpha
  print *, " Beta: ", settings%beta
  print *, " Bias: ", settings%bias
  print *, " Boundary Conditions: ", trim(settings%boundary_cond)
  print *, "==========================================="

  ! Stability check hint (very rough estimate for diffusion)
  if (settings%kappa > 0.0_dp) then
     if (settings%dt > (settings%dx**2) / (2.0_dp * settings%kappa * settings%n_dims)) then
       print *, "Warning: Time step dt may be too large for stability (CFL condition)."
     end if
  end if

  ! -- 2. Initialization --
  ! Allocate and initialize weights vector w (size N) within initialization
  call initialize_solution(settings, solution)
  print *, "Initialization complete. Weights (w):", settings%weights(1:settings%n_dims)
  print *, "Solution array 'u' allocated with shape: ", lbound(solution%u), ubound(solution%u)

  ! -- 3. Time Stepping Loop --
  t_current = 0.0_dp
  do n = 1, n_steps
    t_current = t_current + settings%dt

    ! Perform one time step (u^{n+1} = EulerStep(u^n))
    call time_step(settings, solution, u_next_step)

    ! Update solution: u^n <- u^{n+1}
    solution%u = u_next_step

    ! Deallocate the temporary array for the next step result
    deallocate(u_next_step)

    ! -- 4. Output / Analysis (Optional) --
    if (mod(n, output_interval) == 0) then
      print '(A, I0, A, F8.4)', "Step: ", n, " / ", n_steps, " Time: ", t_current
      ! Add code here to write solution to file, compute statistics, etc.
      ! Example: Print min/max values
      print *, "  Min/Max u: ", minval(solution%u), maxval(solution%u)
    end if

  end do

  print *, "==========================================="
  print *, " Simulation Finished"
  print *, " Final Time: ", t_current
  print *, "==========================================="

  ! -- 5. Cleanup --
  if (allocated(solution%u)) deallocate(solution%u)
  if (allocated(settings%weights)) deallocate(settings%weights)

end program main_pde_solver