! TSU-WAVE NSWE Solver (Fortran 90)
! Nonlinear Shallow-Water Equations for Tsunami Propagation

module nswe_solver
  implicit none
  real, parameter :: g = 9.81  ! gravity
  real, parameter :: rho = 1025.0  ! seawater density
  
contains

  subroutine solve_nswe(eta, u, v, H, nx, ny, dt, dx, dy, nt)
    ! Solve 2D Nonlinear Shallow-Water Equations
    implicit none
    integer, intent(in) :: nx, ny, nt
    real, intent(inout) :: eta(nx, ny), u(nx, ny), v(nx, ny)
    real, intent(in) :: H(nx, ny)
    real, intent(in) :: dt, dx, dy
    
    integer :: i, j, t
    real :: eta_new(nx, ny), u_new(nx, ny), v_new(nx, ny)
    real :: du_dx, du_dy, dv_dx, dv_dy, deta_dx, deta_dy
    
    do t = 1, nt
      ! Continuity equation
      do i = 2, nx-1
        do j = 2, ny-1
          du_dx = (u(i+1,j) - u(i-1,j)) / (2*dx)
          dv_dy = (v(i,j+1) - v(i,j-1)) / (2*dy)
          eta_new(i,j) = eta(i,j) - dt * ((H(i,j) + eta(i,j)) * (du_dx + dv_dy))
        end do
      end do
      
      ! Momentum equations (simplified)
      do i = 2, nx-1
        do j = 2, ny-1
          deta_dx = (eta_new(i+1,j) - eta_new(i-1,j)) / (2*dx)
          u_new(i,j) = u(i,j) - dt * (u(i,j) * (u(i+1,j)-u(i-1,j))/(2*dx) + &
                                       v(i,j) * (u(i,j+1)-u(i,j-1))/(2*dy) + &
                                       g * deta_dx)
          
          deta_dy = (eta_new(i,j+1) - eta_new(i,j-1)) / (2*dy)
          v_new(i,j) = v(i,j) - dt * (u(i,j) * (v(i+1,j)-v(i-1,j))/(2*dx) + &
                                       v(i,j) * (v(i,j+1)-v(i,j-1))/(2*dy) + &
                                       g * deta_dy)
        end do
      end do
      
      ! Update boundaries
      call apply_boundary_conditions(eta_new, u_new, v_new, nx, ny)
      
      ! Copy back
      eta = eta_new
      u = u_new
      v = v_new
    end do
    
  end subroutine solve_nswe
  
  subroutine apply_boundary_conditions(eta, u, v, nx, ny)
    implicit none
    integer, intent(in) :: nx, ny
    real, intent(inout) :: eta(nx, ny), u(nx, ny), v(nx, ny)
    integer :: i, j
    
    ! Zero-gradient at boundaries
    ! Left boundary
    eta(1,:) = eta(2,:)
    u(1,:) = u(2,:)
    v(1,:) = v(2,:)
    
    ! Right boundary
    eta(nx,:) = eta(nx-1,:)
    u(nx,:) = u(nx-1,:)
    v(nx,:) = v(nx-1,:)
    
    ! Bottom boundary
    eta(:,1) = eta(:,2)
    u(:,1) = u(:,2)
    v(:,1) = v(:,2)
    
    ! Top boundary
    eta(:,ny) = eta(:,ny-1)
    u(:,ny) = u(:,ny-1)
    v(:,ny) = v(:,ny-1)
    
  end subroutine apply_boundary_conditions

end module nswe_solver