module sw_utils
    implicit none
    private
    logical, save :: sw_allocated = .false.
    ! - double precision
    integer, parameter, public :: idp = kind(1.0d0)

    real(idp), parameter, public :: PI = 4.0_idp*atan(1.0)

    ! -------------------------------------
    ! number of point in x and y directions
    ! -------------------------------------
    integer, parameter, public :: NX = 101
    integer, parameter, public :: NY = 101

    ! --------------------
    ! number of time steps
    ! --------------------
    integer, parameter, public :: NT = 10001

    ! ----------------------------------------------
    ! upper and lower bounds for x and y dimensions
    ! ----------------------------------------------
    real(idp), parameter, public :: XUPPER = 5.0_idp     ! xend
    real(idp), parameter, public :: XLOWER = -5.0_idp    ! xstart
    real(idp), parameter, public :: YUPPER = 6.0_idp     ! ystart
    real(idp), parameter, public :: YLOWER = -6.0_idp    ! yend

    ! -----------------------------
    ! start and end point for time
    ! -----------------------------
    real(idp), parameter, public :: TSTART = 0.0_idp
    real(idp), parameter, public :: TEND = 1000.0_idp

    ! ----------------------------
    ! length of x-axis and y-axis
    ! ----------------------------
    real(idp), parameter, public :: XLENGTH = (XUPPER - XLOWER)
    real(idp), parameter, public :: YLENGTH = (YUPPER - YLOWER)

    ! -----------------------
    ! total time or duration
    ! -----------------------
    real(idp), parameter, public :: TLENGTH = (TEND - TSTART)

    ! ---------------------------
    ! spatial grid step dx and dy
    ! ---------------------------
    real(idp), parameter, public :: DX = XLENGTH/dble(NX-1)
    real(idp), parameter, public :: DY = YLENGTH/dble(NY-1)

    ! ----------
    ! time step
    ! ----------
    real(idp), parameter, public :: DT = TLENGTH/dble(NT-1)
    

    ! ---------------------
    ! gravitation constant
    ! ---------------------
    real(idp), parameter, public :: G = 9.8065_idp

    ! ------------
    ! base height
    ! ------------
    real(idp), parameter, public :: H0 = 1.0_idp

    ! --------------------------
    ! rotation parameter, omega
    ! --------------------------
    real(idp), parameter, public :: OMEGA = 0.1_idp

    ! -------------
    ! viscous drag
    ! -------------
    real(idp), parameter, public :: vd = 0.001_idp

    ! -----------------------
    ! viscosity of the fluid
    ! -----------------------
    real(idp), parameter, public :: nu = 0.000015_idp
    
    ! -------------
    ! perturbation
    ! -------------
    real(idp), parameter, public :: PERTURB = 0.01_idp

    ! ----------------------
    ! output data file name
    ! ----------------------
    character(len=*), parameter, public :: swfname = "sw_data.nc"

    ! =====================================
    ! public variables' name and constants
    ! =====================================
    real(idp), dimension(1:NX), public :: xdata
    real(idp), dimension(1:NY), public :: ydata
    real(idp), allocatable, public :: tdata(:)
    real(idp), allocatable, public :: xxdata(:, :)
    real(idp), allocatable, public :: yydata(:, :)
    real(idp), allocatable, public :: fdata(:, :)
    real(idp), allocatable, public :: udata(:, :, :)
    real(idp), allocatable, public :: vdata(:, :, :)
    real(idp), allocatable, public :: hdata(:, :, :)

    ! ---------------------------------
    ! public subroutines and functions
    ! ---------------------------------
    public swMeshgrid
    public swAllocateData
    public swDeallocateData
    public swSaveData

contains
    subroutine swAllocateData()
        implicit none
        if ( .not. sw_allocated ) then
            allocate(tdata(NT))
            allocate(xxdata(NX, NY))
            allocate(yydata(NX, NY))
            allocate(fdata(NX, NY))
            allocate(udata(NX, NY, NT))
            allocate(vdata(NX, NY, NT))
            allocate(hdata(NX, NY, NT))
            sw_allocated = .true.
        endif
    end subroutine swAllocateData

    subroutine swDeallocateData()
        implicit none
        if (sw_allocated) then
            deallocate(tdata)
            deallocate(xxdata)
            deallocate(yydata)
            deallocate(fdata)
            deallocate(udata)
            deallocate(vdata)
            deallocate(hdata)
            sw_allocated = .false.
        endif
    end subroutine swDeallocateData

    ! ---
    subroutine swMeshgrid(x, y, xx, yy, nx, ny)
        implicit none
        integer, intent(in) :: nx, ny
        real(kind=8), dimension(1:nx), intent(in) :: x
        real(kind=8), dimension(1:ny), intent(in) :: y
        real(kind=8), dimension(1:nx, 1:ny), intent(out) :: xx, yy
        xx = spread(x, 1, size(y))
        yy = spread(y, 2, size(x))
    end subroutine swMeshgrid

    subroutine swSaveData()
        use netcdf
        implicit none

        ! -----------------------------------------
        ! number of dimensions for large variables
        ! -----------------------------------------
        integer, parameter :: NDIMS3 = 3
        integer, parameter :: NDIMS2 = 2

        ! ------------------------
        ! id for netcdf variables
        ! ------------------------
        integer :: ncid

        ! ----------------
        ! dimension names
        ! ----------------
        character(len=*), parameter :: xname = "X Distance"
        character(len=*), parameter :: yname = "Y Distance"
        character(len=*), parameter :: tname = "T Time"

        ! --------------
        ! dimension ids
        ! --------------
        integer :: xdimID, ydimID, tdimID

        ! --------------------
        ! shape of dimensions
        ! --------------------
        integer :: dimID3(NDIMS3)
        integer :: dimID2(NDIMS2)

        ! ---------------
        ! variable names
        ! ---------------
        character(len=*), parameter :: uname = "U_Velocity"
        character(len=*), parameter :: vname = "V_Velocity"
        character(len=*), parameter :: hname = "H_Level"
        character(len=*), parameter :: fname = "F_Coriolis"

        ! ----------------------------------
        ! Units for the dimension variables
        ! ----------------------------------
        character(len=*), parameter :: units = "units"
        character(len=*), parameter :: u_units = "m/s"
        character(len=*), parameter :: v_units = "m/s"
        character(len=*), parameter :: h_units = "m"
        character(len=*), parameter :: f_units = "1/s"
        character(len=*), parameter :: x_units = "meters_east"
        character(len=*), parameter :: y_units = "meters_north"
        character(len=*), parameter :: t_units = "seconds_forward"

        ! ---------------------------
        ! variable IDs for variables
        ! ---------------------------
        integer :: uvarID, vvarID, hvarID, fvarID

        ! -----------------------------------------
        ! variable IDs for the dimension variables
        ! -----------------------------------------
        integer :: xvarID, yvarID, tvarID

        write(*, *) "Defined all data variables and their IDs."

        ! ----------------
        ! create the file
        ! ----------------
        call check(nf90_create(                                     &
            swfname, or(nf90_clobber, nf90_64bit_offset), ncid)   &
        )
        
        ! ----------------------
        ! define the dimensions
        ! ----------------------
        call check(nf90_def_dim(ncid, xname, NX, xdimID))
        call check(nf90_def_dim(ncid, yname, NY, ydimID))
        call check(nf90_def_dim(ncid, tname, NT, tdimID))

        ! ---------------------------
        ! create dimension variables
        ! ---------------------------
        call check(nf90_def_var(ncid, xname, NF90_DOUBLE, xdimID, xvarID))
        call check(nf90_def_var(ncid, yname, NF90_DOUBLE, ydimID, yvarID))
        call check(nf90_def_var(ncid, tname, NF90_DOUBLE, tdimID, tvarID))

        ! ----------------------------------
        ! set attributes for the dimensions
        ! ----------------------------------
        call check(nf90_put_att(ncid, xvarID, units, x_units))
        call check(nf90_put_att(ncid, yvarID, units, y_units))
        call check(nf90_put_att(ncid, tvarID, units, t_units))

        write(*, *) "Dimensions' unit attribute added"

        ! ----------------------
        ! create variable slots
        ! ----------------------
        dimID3 = (/ xdimID, ydimID, tdimID /)
        dimID2 = (/ xdimID, ydimID /)

        call check(nf90_def_var(ncid, uname, NF90_DOUBLE, dimID3, uvarID))
        call check(nf90_def_var(ncid, vname, NF90_DOUBLE, dimID3, vvarID))
        call check(nf90_def_var(ncid, hname, NF90_DOUBLE, dimID3, hvarID))
        call check(nf90_def_var(ncid, fname, NF90_DOUBLE, dimID2, fvarID))

        write(*, *) "Data variables created"

        ! ----------------------------------
        ! set attribute for data variabales
        ! ----------------------------------
        call check(nf90_put_att(ncid, uvarID, units, u_units))
        call check(nf90_put_att(ncid, vvarID, units, v_units))
        call check(nf90_put_att(ncid, hvarID, units, h_units))
        call check(nf90_put_att(ncid, fvarID, units, f_units))

        write(*, *) "Data variable unit attribute added"

        ! --------------------------
        ! clode variable definition
        ! --------------------------
        call check(nf90_enddef(ncid))

        ! ------------------------------
        ! write the dimension variables
        ! ------------------------------
        call check(nf90_put_var(ncid, xvarID, xdata))
        write(*, *) "Wrote x dimension variable"
        call check(nf90_put_var(ncid, yvarID, ydata))
        write(*, *) "Wrote y dimension variable"
        call check(nf90_put_var(ncid, tvarID, tdata))
        write(*, *) "Wrote t dimension variable"

        ! -------------------------
        ! write the data variables
        ! -------------------------
        call check(nf90_put_var(ncid, uvarID, udata))
        write(*, *) "Wrote u data variable"
        call check(nf90_put_var(ncid, vvarID, vdata))
        write(*, *) "Wrote v data variable"
        call check(nf90_put_var(ncid, hvarID, hdata))
        write(*, *) "Wrote h (i.e height) data variable"
        call check(nf90_put_var(ncid, fvarID, fdata))
        write(*, *) "Wrote f (i.e coriolis) data variable"

        ! --------------
        ! close the file
        ! ---------------
        call check(nf90_close(ncid))
        
    contains
        subroutine check(status)
            implicit none
            integer, intent(in) :: status
            if(status /= nf90_noerr) then
                write(*, *) trim(nf90_strerror(status))
                stop "Stopped"
            endif
        end subroutine check
    end subroutine swSaveData
end module sw_utils
