!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!! miniWeather
!! Author: Matt Norman <normanmr@ornl.gov>  , Oak Ridge National Laboratory
!! This code simulates dry, stratified, compressible, non-hydrostatic fluid flows
!! For documentation, please see the attached documentation in the "documentation" folder
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!! Copyright (c) 2018, National Center for Computational Sciences, Oak Ridge National Laboratory.  All rights reserved.
!!
!! Portions Copyright (c) 2020, NVIDIA Corporation.  All rights reserved.
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

program miniweather
  use nvtx
  !implicit none

  !Declare the precision for the real constants (at least 15 digits of accuracy = double precision)
  integer , parameter :: rp = selected_real_kind(15)
  !Define some physical constants to use throughout the simulation
  real(rp), parameter :: pi        = 3.14159265358979323846264338327_rp   !Pi
  real(rp), parameter :: grav      = 9.8_rp                               !Gravitational acceleration (m / s^2)
  real(rp), parameter :: cp        = 1004._rp                             !Specific heat of dry air at constant pressure
  real(rp), parameter :: rd        = 287._rp                              !Dry air constant for equation of state (P=rho*rd*T)
  real(rp), parameter :: p0        = 1.e5_rp                              !Standard pressure at the surface in Pascals
  real(rp), parameter :: C0        = 27.5629410929725921310572974482_rp   !Constant to translate potential temperature into pressure (P=C0*(rho*theta)**gamma)
  real(rp), parameter :: gamma     = 1.40027894002789400278940027894_rp   !gamma=cp/Rd
  !Define domain and stability-related constants
  real(rp), parameter :: xlen      = 2.e4_rp    !Length of the domain in the x-direction (meters)
  real(rp), parameter :: zlen      = 1.e4_rp    !Length of the domain in the z-direction (meters)
  real(rp), parameter :: hv_beta   = 0.25_rp     !How strong to diffuse the solution: hv_beta \in [0:1]
  real(rp), parameter :: cfl       = 1.50_rp    !"Courant, Friedrichs, Lewy" number (for numerical stability)
  real(rp), parameter :: max_speed = 450        !Assumed maximum wave speed during the simulation (speed of sound + speed of wind) (meter / sec)
  integer , parameter :: hs        = 2          !"Halo" size: number of cells needed for a full "stencil" of information for reconstruction
  integer , parameter :: sten_size = 4          !Size of the stencil used for interpolation

  !Parameters for indexing and flags
  integer , parameter :: NUM_VARS = 4           !Number of fluid state variables
  integer , parameter :: ID_DENS  = 1           !index for density ("rho")
  integer , parameter :: ID_UMOM  = 2           !index for momentum in the x-direction ("rho * u")
  integer , parameter :: ID_WMOM  = 3           !index for momentum in the z-direction ("rho * w")
  integer , parameter :: ID_RHOT  = 4           !index for density * potential temperature ("rho * theta")
  integer , parameter :: DIR_X = 1              !Integer constant to express that this operation is in the x-direction
  integer , parameter :: DIR_Z = 2              !Integer constant to express that this operation is in the z-direction

  !Gauss-Legendre quadrature points and weights on the domain [0:1]
  integer , parameter :: nqpoints = 3
  real(rp) :: qpoints (nqpoints) = (/ 0.112701665379258311482073460022E0_rp , 0.500000000000000000000000000000E0_rp , 0.887298334620741688517926539980E0_rp /)
  real(rp) :: qweights(nqpoints) = (/ 0.277777777777777777777777777779E0_rp , 0.444444444444444444444444444444E0_rp , 0.277777777777777777777777777779E0_rp /)

  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  !! Variables that are initialized but remain static over the course of the simulation
  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  real(rp) :: sim_time                          !total simulation time in seconds
  real(rp) :: output_freq                       !frequency to perform output in seconds
  real(rp) :: dt                                !Model time step (seconds)
  integer  :: nx, nz                            !Number of local grid cells in the x- and z- dimensions
  real(rp) :: dx, dz                            !Grid space length in x- and z-dimension (meters)
  integer  :: nx_glob, nz_glob                  !Number of total grid cells in the x- and z- dimensions
  integer  :: i_beg, k_beg                      !beginning index in the x- and z-directions
  integer  :: nranks, myrank                    !my rank id
  integer  :: left_rank, right_rank
  real(rp), allocatable :: hy_dens_cell      (:)      !hydrostatic density (vert cell avgs).   Dimensions: (1-hs:nz+hs)
  real(rp), allocatable :: hy_dens_theta_cell(:)      !hydrostatic rho*t (vert cell avgs).     Dimensions: (1-hs:nz+hs)
  real(rp), allocatable :: hy_dens_int       (:)      !hydrostatic density (vert cell interf). Dimensions: (1:nz+1)
  real(rp), allocatable :: hy_dens_theta_int (:)      !hydrostatic rho*t (vert cell interf).   Dimensions: (1:nz+1)
  real(rp), allocatable :: hy_pressure_int   (:)      !hydrostatic press (vert cell interf).   Dimensions: (1:nz+1)

  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  !! Variables that are dynamics over the course of the simulation
  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  real(rp) :: etime                             !Elapsed model time
  real(rp) :: output_counter                    !Helps determine when it's time to do output
  !Runtime variable arrays
  real(rp), allocatable :: state             (:,:,:)  !Fluid state.                            Dimensions: (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
  real(rp), allocatable :: state_tmp         (:,:,:)  !Fluid state.                            Dimensions: (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
  real(rp), allocatable :: flux              (:,:,:)  !Cell interface fluxes.                  Dimensions: (nx+1,nz+1,NUM_VARS)
  real(rp), allocatable :: tend              (:,:,:)  !Fluid state tendencies.                 Dimensions: (nx,nz,NUM_VARS)
  integer(8) :: t1, t2, rate                          !For CPU Timings

  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  !! THE MAIN PROGRAM STARTS HERE
  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  !! BEGIN USER-CONFIGURABLE PARAMETERS
  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  !The x-direction length is twice as long as the z-direction length
  !So, you'll want to have nx_glob be twice as large as nz_glob
  nx_glob = 40        !Number of total cells in the x-dirction
  nz_glob = 20        !Number of total cells in the z-dirction
  sim_time = 1000 !How many seconds to run the simulation
  output_freq = 100 !How frequently to output data to file (in seconds)
  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  !! END USER-CONFIGURABLE PARAMETERS
  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !Allocate arrays, initialize the grid and the data
  call nvtxStartRange("Total")
  call init()

  !Output the initial state
  !call output(state,etime)

  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  !! MAIN TIME STEP LOOP
  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 call system_clock(t1)
  call nvtxStartRange("while")
  do while (etime < sim_time)
    !If the time step leads to exceeding the simulation time, shorten it for the last step
    if (etime + dt > sim_time) dt = sim_time - etime
    !Perform a single time step
    call nvtxStartRange("perform_timestep")
    call perform_timestep(state,state_tmp,flux,tend,dt)
    call nvtxEndRange
    !Inform the user
write(*,*) 'Elapsed Time: ', etime , ' / ' , sim_time
    !Update the elapsed time and output counter
    etime = etime + dt
    output_counter = output_counter + dt
    !If it's time for output, reset the counter, and do output
    if (output_counter >= output_freq) then
      output_counter = output_counter - output_freq
      !call output(state,etime)
    endif
  enddo
  call nvtxEndRange

    call system_clock(t2,rate)
    write(*,*) "CPU Time: ",dble(t2-t1)/dble(rate)

  !Deallocate
  call finalize()
  call nvtxEndRange

contains

  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  !! HELPER SUBROUTINES AND FUNCTIONS GO HERE
  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !Performs a single dimensionally split time step using a simple low-storate three-stage Runge-Kutta time integrator
  !The dimensional splitting is a second-order-accurate alternating Strang splitting in which the
  !order of directions is alternated each time step.
  !The Runge-Kutta method used here is defined as follows:
  ! q*     = q[n] + dt/3 * rhs(q[n])
  ! q**    = q[n] + dt/2 * rhs(q*  )
  ! q[n+1] = q[n] + dt/1 * rhs(q** )
  subroutine perform_timestep(state,state_tmp,flux,tend,dt)
    implicit none
    real(rp), intent(inout) :: state    (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
    real(rp), intent(  out) :: state_tmp(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
    real(rp), intent(  out) :: flux     (nx+1,nz+1,NUM_VARS)
    real(rp), intent(  out) :: tend     (nx,nz,NUM_VARS)
    real(rp), intent(in   ) :: dt
    logical, save :: direction_switch = .true.
    if (direction_switch) then
      !x-direction first
      call semi_discrete_step( state , state     , state_tmp , dt / 3 , DIR_X , flux , tend )
      call semi_discrete_step( state , state_tmp , state_tmp , dt / 2 , DIR_X , flux , tend )
      call semi_discrete_step( state , state_tmp , state     , dt / 1 , DIR_X , flux , tend )
      !z-direction second
      call semi_discrete_step( state , state     , state_tmp , dt / 3 , DIR_Z , flux , tend )
      call semi_discrete_step( state , state_tmp , state_tmp , dt / 2 , DIR_Z , flux , tend )
      call semi_discrete_step( state , state_tmp , state     , dt / 1 , DIR_Z , flux , tend )
    else
      !z-direction second
      call semi_discrete_step( state , state     , state_tmp , dt / 3 , DIR_Z , flux , tend )
      call semi_discrete_step( state , state_tmp , state_tmp , dt / 2 , DIR_Z , flux , tend )
      call semi_discrete_step( state , state_tmp , state     , dt / 1 , DIR_Z , flux , tend )
      !x-direction first
      call semi_discrete_step( state , state     , state_tmp , dt / 3 , DIR_X , flux , tend )
      call semi_discrete_step( state , state_tmp , state_tmp , dt / 2 , DIR_X , flux , tend )
      call semi_discrete_step( state , state_tmp , state     , dt / 1 , DIR_X , flux , tend )
    endif
    direction_switch = .not. direction_switch
  end subroutine perform_timestep

  !Perform a single semi-discretized step in time with the form:
  !state_out = state_init + dt * rhs(state_forcing)
  !Meaning the step starts from state_init, computes the rhs using state_forcing, and stores the result in state_out
  subroutine semi_discrete_step( state_init , state_forcing , state_out , dt , dir , flux , tend )
    implicit none
    real(rp), intent(in   ) :: state_init   (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
    real(rp), intent(inout) :: state_forcing(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
    real(rp), intent(  out) :: state_out    (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
    real(rp), intent(  out) :: flux         (nx+1,nz+1,NUM_VARS)
    real(rp), intent(  out) :: tend         (nx,nz,NUM_VARS)
    real(rp), intent(in   ) :: dt
    integer , intent(in   ) :: dir
    integer :: i,k,ll

    if     (dir == DIR_X) then
      !Set the halo values in the x-direction
      call set_halo_values_x(state_forcing)
      !Compute the time tendencies for the fluid state in the x-direction
      call compute_tendencies_x(state_forcing,flux,tend)
    elseif (dir == DIR_Z) then
      !Set the halo values in the z-direction
      call set_halo_values_z(state_forcing)
      !Compute the time tendencies for the fluid state in the z-direction
      call compute_tendencies_z(state_forcing,flux,tend)
    endif

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! TODO: THREAD ME
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !Apply the tendencies to the fluid state
    !$acc parallel loop collapse(3)
    do ll = 1 , NUM_VARS
      do k = 1 , nz
        do i = 1 , nx
          state_out(i,k,ll) = state_init(i,k,ll) + dt * tend(i,k,ll)
        enddo
      enddo
    enddo
  end subroutine semi_discrete_step

  !Compute the time tendencies of the fluid state using forcing in the x-direction

  !First, compute the flux vector at each cell interface in the x-direction (including hyperviscosity)
  !Then, compute the tendencies using those fluxes
  subroutine compute_tendencies_x(state,flux,tend)
    implicit none
    real(rp), intent(in   ) :: state(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
    real(rp), intent(  out) :: flux (nx+1,nz+1,NUM_VARS)
    real(rp), intent(  out) :: tend (nx,nz,NUM_VARS)
    integer :: i,k,ll,s
    real(rp) :: r,u,w,t,p, stencil(4), d3_vals(NUM_VARS), vals(NUM_VARS), hv_coef
    !Compute the hyperviscosity coeficient
    hv_coef = -hv_beta * dx / (16*dt)
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! TODO: THREAD ME
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !Compute fluxes in the x-direction for each cell
    !$acc parallel loop
    do k = 1 , nz
      !$acc loop private(stencil,vals,d3_vals)
      do i = 1 , nx+1
        !Use fourth-order interpolation from four cell averages to compute the value at the interface in question
        do ll = 1 , NUM_VARS
          do s = 1 , sten_size
            stencil(s) = state(i-hs-1+s,k,ll)
          enddo
          !Fourth-order-accurate interpolation of the state
          vals(ll) = -stencil(1)/12 + 7*stencil(2)/12 + 7*stencil(3)/12 - stencil(4)/12
          !First-order-accurate interpolation of the third spatial derivative of the state (for artificial viscosity)
          d3_vals(ll) = -stencil(1) + 3*stencil(2) - 3*stencil(3) + stencil(4)
        enddo

        !Compute density, u-wind, w-wind, potential temperature, and pressure (r,u,w,t,p respectively)
        r = vals(ID_DENS) + hy_dens_cell(k)
        u = vals(ID_UMOM) / r
        w = vals(ID_WMOM) / r
        t = ( vals(ID_RHOT) + hy_dens_theta_cell(k) ) / r
        p = C0*(r*t)**gamma

        !Compute the flux vector
        flux(i,k,ID_DENS) = r*u     - hv_coef*d3_vals(ID_DENS)
        flux(i,k,ID_UMOM) = r*u*u+p - hv_coef*d3_vals(ID_UMOM)
        flux(i,k,ID_WMOM) = r*u*w   - hv_coef*d3_vals(ID_WMOM)
        flux(i,k,ID_RHOT) = r*u*t   - hv_coef*d3_vals(ID_RHOT)
      enddo
    enddo

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! TODO: THREAD ME
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !Use the fluxes to compute tendencies for each cell
    !$acc parallel loop collapse(3)
    do ll = 1 , NUM_VARS
      do k = 1 , nz
        do i = 1 , nx
          tend(i,k,ll) = -( flux(i+1,k,ll) - flux(i,k,ll) ) / dx
        enddo
      enddo
    enddo
  end subroutine compute_tendencies_x

  !Compute the time tendencies of the fluid state using forcing in the z-direction

  !First, compute the flux vector at each cell interface in the z-direction (including hyperviscosity)
  !Then, compute the tendencies using those fluxes
  subroutine compute_tendencies_z(state,flux,tend)
    implicit none
    real(rp), intent(in   ) :: state(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
    real(rp), intent(  out) :: flux (nx+1,nz+1,NUM_VARS)
    real(rp), intent(  out) :: tend (nx,nz,NUM_VARS)
    integer :: i,k,ll,s
    real(rp) :: r,u,w,t,p, stencil(4), d3_vals(NUM_VARS), vals(NUM_VARS), hv_coef
    !Compute the hyperviscosity coeficient
    hv_coef = -hv_beta * dz / (16*dt)
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! TODO: THREAD ME
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !Compute fluxes in the x-direction for each cell
    !$acc parallel loop
    do k = 1 , nz+1
      !$acc loop private(stencil,vals,d3_vals)
      do i = 1 , nx
        !Use fourth-order interpolation from four cell averages to compute the value at the interface in question
        do ll = 1 , NUM_VARS
          do s = 1 , sten_size
            stencil(s) = state(i,k-hs-1+s,ll)
          enddo
          !Fourth-order-accurate interpolation of the state
          vals(ll) = -stencil(1)/12 + 7*stencil(2)/12 + 7*stencil(3)/12 - stencil(4)/12
          !First-order-accurate interpolation of the third spatial derivative of the state
          d3_vals(ll) = -stencil(1) + 3*stencil(2) - 3*stencil(3) + stencil(4)
        enddo

        !Compute density, u-wind, w-wind, potential temperature, and pressure (r,u,w,t,p respectively)
        r = vals(ID_DENS) + hy_dens_int(k)
        u = vals(ID_UMOM) / r
        w = vals(ID_WMOM) / r
        t = ( vals(ID_RHOT) + hy_dens_theta_int(k) ) / r
        p = C0*(r*t)**gamma - hy_pressure_int(k)
        !Enforce vertical boundary condition and exact mass conservation
        if (k == 1 .or. k == nz+1) then
          w                = 0
          d3_vals(ID_DENS) = 0
        endif

        !Compute the flux vector with hyperviscosity
        flux(i,k,ID_DENS) = r*w     - hv_coef*d3_vals(ID_DENS)
        flux(i,k,ID_UMOM) = r*w*u   - hv_coef*d3_vals(ID_UMOM)
        flux(i,k,ID_WMOM) = r*w*w+p - hv_coef*d3_vals(ID_WMOM)
        flux(i,k,ID_RHOT) = r*w*t   - hv_coef*d3_vals(ID_RHOT)
      enddo
    enddo

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! TODO: THREAD ME
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !Use the fluxes to compute tendencies for each cell
    !$acc parallel loop collapse(3)
    do ll = 1 , NUM_VARS
      do k = 1 , nz
        do i = 1 , nx
          tend(i,k,ll) = -( flux(i,k+1,ll) - flux(i,k,ll) ) / dz
          if (ll == ID_WMOM) then
            tend(i,k,ID_WMOM) = tend(i,k,ID_WMOM) - state(i,k,ID_DENS)*grav
          endif
        enddo
      enddo
    enddo
  end subroutine compute_tendencies_z

  !Set halo values in the x-direction.
  subroutine set_halo_values_x(state)
    implicit none
    real(rp), intent(inout) :: state(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
    integer :: k, ll
    real(rp) :: z

    !$acc parallel loop collapse(2)
    do ll = 1 , NUM_VARS
      do k = 1 , nz
        state(-1  ,k,ll) = state(nx-1,k,ll)
        state(0   ,k,ll) = state(nx  ,k,ll)
        state(nx+1,k,ll) = state(1   ,k,ll)
        state(nx+2,k,ll) = state(2   ,k,ll)
      enddo
    enddo

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

      if (myrank == 0) then
        !$acc parallel loop
        do k = 1 , nz
          z = (k_beg-1 + k-0.5_rp)*dz
          if (abs(z-3*zlen/4) <= zlen/16) then
            state(-1:0,k,ID_UMOM) = (state(-1:0,k,ID_DENS)+hy_dens_cell(k)) * 50._rp
            state(-1:0,k,ID_RHOT) = (state(-1:0,k,ID_DENS)+hy_dens_cell(k)) * 298._rp - hy_dens_theta_cell(k)
          endif
        enddo
      endif
  end subroutine set_halo_values_x

  !Set halo values in the z-direction.
  !decomposition in the vertical direction
  subroutine set_halo_values_z(state)
    implicit none
    real(rp), intent(inout) :: state(1-hs:nx+hs,1-hs:nz+hs,NUM_VARS)
    integer :: i, ll
    real(rp), parameter :: mnt_width = xlen/8
    real(rp) :: x, xloc, mnt_deriv
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! TODO: THREAD ME
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !$acc parallel loop collapse(2) private(xloc,mnt_deriv,x)
    do ll = 1 , NUM_VARS
      do i = 1-hs,nx+hs
        if (ll == ID_WMOM) then
          state(i,-1  ,ll) = 0
          state(i,0   ,ll) = 0
          state(i,nz+1,ll) = 0
          state(i,nz+2,ll) = 0
        else
          state(i,-1  ,ll) = state(i,1 ,ll)
          state(i,0   ,ll) = state(i,1 ,ll)
          state(i,nz+1,ll) = state(i,nz,ll)
          state(i,nz+2,ll) = state(i,nz,ll)
        endif
      enddo
    enddo
  end subroutine set_halo_values_z

  !Initialize some grid settings, allocate data, and initialize the full grid and data
  subroutine init()
    implicit none
    integer :: i, k, ii, kk, ll, ierr
    real(rp) :: x, z, r, u, w, t, hr, ht

    !Set the cell grid size
    dx = xlen / nx_glob
    dz = zlen / nz_glob

    nranks = 1
    myrank = 0
    i_beg = 1
    nx = nx_glob
    left_rank = 0
    right_rank = 0

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! YOU DON'T NEED TO ALTER ANYTHING BELOW THIS POINT IN THE CODE
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

    k_beg = 1
    nz = nz_glob

    !Allocate the model data
    allocate(state             (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS))
    allocate(state_tmp         (1-hs:nx+hs,1-hs:nz+hs,NUM_VARS))
    allocate(flux              (nx+1,nz+1,NUM_VARS))
    allocate(tend              (nx,nz,NUM_VARS))
    allocate(hy_dens_cell      (1-hs:nz+hs))
    allocate(hy_dens_theta_cell(1-hs:nz+hs))
    allocate(hy_dens_int       (nz+1))
    allocate(hy_dens_theta_int (nz+1))
    allocate(hy_pressure_int   (nz+1))

    !Define the maximum stable time step based on an assumed maximum wind speed
    dt = min(dx,dz) / max_speed * cfl
    !Set initial elapsed model time and output_counter to zero
    etime = 0
    output_counter = 0

    !Display some grid information

      write(*,*) 'nx_glob, nz_glob:', nx_glob, nz_glob
      write(*,*) 'dx,dz: ',dx,dz
      write(*,*) 'dt: ',dt

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! Initialize the cell-averaged fluid state via Gauss-Legendre quadrature
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    state = 0
    do k = 1-hs , nz+hs
      do i = 1-hs , nx+hs
        !Initialize the state to zero
        !Use Gauss-Legendre quadrature to initialize a hydrostatic balance + temperature perturbation
        do kk = 1 , nqpoints
          do ii = 1 , nqpoints
            !Compute the x,z location within the global domain based on cell and quadrature index
            x = (i_beg-1 + i-0.5_rp) * dx + (qpoints(ii)-0.5_rp)*dx
            z = (k_beg-1 + k-0.5_rp) * dz + (qpoints(kk)-0.5_rp)*dz

            !Set the fluid state based on the user's specification
            call injection      (x,z,r,u,w,t,hr,ht)

            !Store into the fluid state array
            state(i,k,ID_DENS) = state(i,k,ID_DENS) + r                         * qweights(ii)*qweights(kk)
            state(i,k,ID_UMOM) = state(i,k,ID_UMOM) + (r+hr)*u                  * qweights(ii)*qweights(kk)
            state(i,k,ID_WMOM) = state(i,k,ID_WMOM) + (r+hr)*w                  * qweights(ii)*qweights(kk)
            state(i,k,ID_RHOT) = state(i,k,ID_RHOT) + ( (r+hr)*(t+ht) - hr*ht ) * qweights(ii)*qweights(kk)
          enddo
        enddo
        do ll = 1 , NUM_VARS
          state_tmp(i,k,ll) = state(i,k,ll)
        enddo
      enddo
    enddo
    !Compute the hydrostatic background state over vertical cell averages
    hy_dens_cell = 0
    hy_dens_theta_cell = 0
    do k = 1-hs , nz+hs
      do kk = 1 , nqpoints
        z = (k_beg-1 + k-0.5_rp) * dz + (qpoints(kk)-0.5_rp)*dz
        !Set the fluid state based on the user's specification
        call injection      (0._rp,z,r,u,w,t,hr,ht)
        hy_dens_cell(k)       = hy_dens_cell(k)       + hr    * qweights(kk)
        hy_dens_theta_cell(k) = hy_dens_theta_cell(k) + hr*ht * qweights(kk)
      enddo
    enddo
    !Compute the hydrostatic background state at vertical cell interfaces
    do k = 1 , nz+1
      z = (k_beg-1 + k-1) * dz
      call injection      (0._rp,z,r,u,w,t,hr,ht)
      hy_dens_int      (k) = hr
      hy_dens_theta_int(k) = hr*ht
      hy_pressure_int  (k) = C0*(hr*ht)**gamma
    enddo
  end subroutine init

  !This test case is initially balanced but injects fast, cold air from the left boundary near the model top
  subroutine injection(x,z,r,u,w,t,hr,ht)
    implicit none
    real(rp), intent(in   ) :: x, z        !x- and z- location of the point being sampled
    real(rp), intent(  out) :: r, u, w, t  !Density, uwind, wwind, and potential temperature
    real(rp), intent(  out) :: hr, ht      !Hydrostatic density and potential temperature
    call hydro_const_theta(z,hr,ht)
    r = 0
    t = 0
    u = 0
    w = 0
  end subroutine injection

  !Establish hydrstatic balance using constant potential temperature (thermally neutral atmosphere)
  subroutine hydro_const_theta(z,r,t)
    implicit none
    real(rp), intent(in   ) :: z  !x- and z- location of the point being sampled
    real(rp), intent(  out) :: r, t  !Density and potential temperature at this point
    real(rp), parameter :: theta0 = 300._rp  !Background potential temperature
    real(rp), parameter :: exner0 = 1._rp    !Surface-level Exner pressure
    real(rp) :: p,exner,rt
    !Establish hydrostatic balance first using Exner pressure
    t = theta0                                  !Potential Temperature at z
    exner = exner0 - grav * z / (cp * theta0)   !Exner pressure at z
    p = p0 * exner**(cp/rd)                     !Pressure at z
    rt = (p / c0)**(1._rp / gamma)              !rho*theta at z
    r = rt / t                                  !Density at z
  end subroutine hydro_const_theta

  !Deallocate and call Finalize
  subroutine finalize()
    implicit none
    integer :: ierr
    deallocate(state             )
    deallocate(state_tmp         )
    deallocate(flux              )
    deallocate(tend              )
    deallocate(hy_dens_cell      )
    deallocate(hy_dens_theta_cell)
    deallocate(hy_dens_int       )
    deallocate(hy_dens_theta_int )
    deallocate(hy_pressure_int   )
  end subroutine finalize

end program miniweather


      double precision function Rgen(pjet,i,p,j)
c--- Calculate the angular separation between pjet(i) and p(j)
c--- This routine is a generalization of R.f: Rgen(p,i,p,j) == R(p,i,j)
      implicit none
      include 'constants.f'
      double precision pjet(mxpart,4),p(mxpart,4),
     & r1,r2,dely,delphi,ei,ej
      integer i,j

      ei=dsqrt(pjet(i,1)**2+pjet(i,2)**2+pjet(i,3)**2)
      ej=dsqrt(p(j,1)**2+p(j,2)**2+p(j,3)**2)

      r1= (ei+pjet(i,3))*(ej-p(j,3))/
     .     ((ej+p(j,3))*(ei-pjet(i,3)))
      dely=0.5d0*dlog(r1)

      r2= (pjet(i,1)*p(j,1)+pjet(i,2)*p(j,2))
     .     /dsqrt((pjet(i,1)**2+pjet(i,2)**2)*(p(j,1)**2+p(j,2)**2))
      if (r2 .gt. +0.9999999D0) r2=+1D0
      if (r2 .lt. -0.9999999D0) r2=-1D0
      delphi=dacos(r2)

      Rgen=dsqrt(dely**2+delphi**2)

      return
      end


c               Copyright 1993 Colorado State University
c                       All Rights Reserved

      double precision function line(x, x1, y1, x2, y2)

      implicit none

c ... Argument declarations
      double precision x, x1, y1, x2, y2

c ... This function is the generic equation of a line from
c ... two points.
c ... Given 2 known points and a new X, calculate Y.
c ... slope = (y2 - y1) / (x2 - x1)
c ... y = slope * (x - x2) + y2
c ...
c ... CALLED FROM:  co2eff.f
c ...               froota.f
c ...               grem.f
c ...               prelim.f
c ...               ramp.f

      line = (y2 - y1) / (x2 - x1) * (x - x2) + y2

      return
      end


!  *********************************************************************
!  *                                                                   *
!  *                        subroutine thin1d                          *
!  *                                                                   *
!  *********************************************************************
!  Mixed Precision Version 1.11
!  Written by Gordon A. Fenton, TUNS, Dec. 8, 1993
!  Latest Update: Jun 9, 1999
!
!  PURPOSE  creates the parameter matrices required for the interior cell
!           subdivisions of LAS2G for the case k1 or k2 = 1 and for LAS3G
!           for k1=k2=1, k1=k3=1, or k2=k3=1.
!
!  Requires:
!   1) from libGAFsim:	DSIFA, DSISL, DCHOL2, DAXPY, DSWAP, DDOT, IDAMAX
!
!  i1, i2, and i3 are the indexes into the global covariance matrix R to be
!  extracted to form the local covariance matrix.
!  n is the number of cells in the subdivision less 1. This is equal to
!  3 for LAS2G and to 7 for LAS3G.
!
!  REVISION HISTORY:
!  1.11	replaced dummy dimensions with a (*) for GNU's compiler (Jun 9/99)
!-----------------------------------------------------------------------------
      subroutine thin1d(R,ir,B,ib,S,is,AS,ias,AI,iai,CS,CI,i1,i2,i3,n,iout,tol)
      real AS(ias,n,*), AI(iai,*), CS(*), CI(*)
      real*8 R(ir,*), B(ib,*), S(is,*)
      real*8 RT(2,2), RI(3,3), DA(3), B1(7,7), B2(7,7)
      integer indT(2), indI(3)

   1  format(a,a,a)
   2  format(a,e13.4)
!							extract R
      RT(1,1) = R(i2,i2)
      RT(1,2) = R(i2,i3)
      RT(2,2) = R(i3,i3)

      RI(1,1) = R(i1,i1)
      RI(1,2) = R(i1,i2)
      RI(1,3) = R(i1,i3)
      RI(2,2) = R(i2,i2)
      RI(2,3) = R(i2,i3)
      RI(3,3) = R(i3,i3)
!							factorize RT (edges)
      call dsifa(RT,2,2,indT,ierr)
      if( ierr .ne. 0 ) then
         write(iout,1)'Error: unable to factorize covariance matrix RT in THIN1D.'
         stop
      endif
!							factorize RI (interior)
      call dsifa(RI,3,3,indI,ierr)
      if( ierr .ne. 0 ) then
         write(iout,1)'Error: unable to factorize covariance matrix RI in THIN1D.'
         stop
      endif

      do 30 j = 1, n
         DA(1) = S(i2,j)
         DA(2) = S(i3,j)
         call dsisl( RT, 2, 2, indT, DA )
         AS(1,j,1) = DA(1)
         AS(2,j,1) = DA(2)
         do 10 i = 1, j
            B1(i,j) = B(i,j) - S(i2,i)*DA(1) - S(i3,i)*DA(2)
  10     continue

         DA(1) = S(i1,j)
         DA(2) = S(i2,j)
         DA(3) = S(i3,j)
         call dsisl( RI, 3, 3, indI, DA )
         AI(1,j) = DA(1)
         AI(2,j) = DA(2)
         AI(3,j) = DA(3)
         do 20 i = 1, j
            B2(i,j) = B(i,j) - S(i1,i)*DA(1) - S(i2,i)*DA(2) - S(i3,i)*DA(3)
  20     continue
  30  continue
!							Cholesky Decomposition
      call dchol2( B1, 7, n, rerr )
      if( rerr .gt. tol ) then
         write(iout,1)'Warning: Cholesky decomposition of THIN1D covariance matrix B1'
         write(iout,2)'         has maximum relative error of ',rerr
      endif

      call dchol2( B2, 7, n, rerr )
      if( rerr .gt. tol ) then
         write(iout,1)'Warning: Cholesky decomposition of THIN1D covariance matrix B2'
         write(iout,2)'         has maximum relative error of ',rerr
      endif
!							store in real*4
      ii = 0
      do 40 j = 1, n
      do 40 i = 1, j
         ii = ii + 1
         CS(ii) = B1(i,j)
         CI(ii) = B2(i,j)
  40  continue

      return
      end


!*==dlasr.f90  processed by SPAG 7.51RB at 20:08 on  3 Mar 2022
!> \brief \b DLASR applies a sequence of plane rotations to a general rectangular matrix.
!
!  =========== DOCUMENTATION ===========
!
! Online html documentation available at
!            http://www.netlib.org/lapack/explore-html/
!
!> \htmlonly
!> Download DLASR + dependencies
!>
!> [TGZ]
!>
!> [ZIP]
!>
!> [TXT]
!> \endhtmlonly
!
!  Definition:
!  ===========
!
!       SUBROUTINE DLASR( SIDE, PIVOT, DIRECT, M, N, C, S, A, LDA )
!
!       .. Scalar Arguments ..
!       CHARACTER          DIRECT, PIVOT, SIDE
!       INTEGER            LDA, M, N
!       ..
!       .. Array Arguments ..
!       DOUBLE PRECISION   A( LDA, * ), C( * ), S( * )
!       ..
!
!
!> \par Purpose:
!  =============
!>
!> \verbatim
!>
!> DLASR applies a sequence of plane rotations to a real matrix A,
!> from either the left or the right.
!>
!> When SIDE = 'L', the transformation takes the form
!>
!>    A := P*A
!>
!> and when SIDE = 'R', the transformation takes the form
!>
!>    A := A*P**T
!>
!> where P is an orthogonal matrix consisting of a sequence of z plane
!> rotations, with z = M when SIDE = 'L' and z = N when SIDE = 'R',
!> and P**T is the transpose of P.
!>
!> When DIRECT = 'F' (Forward sequence), then
!>
!>    P = P(z-1) * ... * P(2) * P(1)
!>
!> and when DIRECT = 'B' (Backward sequence), then
!>
!>    P = P(1) * P(2) * ... * P(z-1)
!>
!> where P(k) is a plane rotation matrix defined by the 2-by-2 rotation
!>
!>    R(k) = (  c(k)  s(k) )
!>         = ( -s(k)  c(k) ).
!>
!> When PIVOT = 'V' (Variable pivot), the rotation is performed
!> for the plane (k,k+1), i.e., P(k) has the form
!>
!>    P(k) = (  1                                            )
!>           (       ...                                     )
!>           (              1                                )
!>           (                   c(k)  s(k)                  )
!>           (                  -s(k)  c(k)                  )
!>           (                                1              )
!>           (                                     ...       )
!>           (                                            1  )
!>
!> where R(k) appears as a rank-2 modification to the identity matrix in
!> rows and columns k and k+1.
!>
!> When PIVOT = 'T' (Top pivot), the rotation is performed for the
!> plane (1,k+1), so P(k) has the form
!>
!>    P(k) = (  c(k)                    s(k)                 )
!>           (         1                                     )
!>           (              ...                              )
!>           (                     1                         )
!>           ( -s(k)                    c(k)                 )
!>           (                                 1             )
!>           (                                      ...      )
!>           (                                             1 )
!>
!> where R(k) appears in rows and columns 1 and k+1.
!>
!> Similarly, when PIVOT = 'B' (Bottom pivot), the rotation is
!> performed for the plane (k,z), giving P(k) the form
!>
!>    P(k) = ( 1                                             )
!>           (      ...                                      )
!>           (             1                                 )
!>           (                  c(k)                    s(k) )
!>           (                         1                     )
!>           (                              ...              )
!>           (                                     1         )
!>           (                 -s(k)                    c(k) )
!>
!> where R(k) appears in rows and columns k and z.  The rotations are
!> performed without ever forming P(k) explicitly.
!> \endverbatim
!
!  Arguments:
!  ==========
!
!> \param[in] SIDE
!> \verbatim
!>          SIDE is CHARACTER*1
!>          Specifies whether the plane rotation matrix P is applied to
!>          A on the left or the right.
!>          = 'L':  Left, compute A := P*A
!>          = 'R':  Right, compute A:= A*P**T
!> \endverbatim
!>
!> \param[in] PIVOT
!> \verbatim
!>          PIVOT is CHARACTER*1
!>          Specifies the plane for which P(k) is a plane rotation
!>          matrix.
!>          = 'V':  Variable pivot, the plane (k,k+1)
!>          = 'T':  Top pivot, the plane (1,k+1)
!>          = 'B':  Bottom pivot, the plane (k,z)
!> \endverbatim
!>
!> \param[in] DIRECT
!> \verbatim
!>          DIRECT is CHARACTER*1
!>          Specifies whether P is a forward or backward sequence of
!>          plane rotations.
!>          = 'F':  Forward, P = P(z-1)*...*P(2)*P(1)
!>          = 'B':  Backward, P = P(1)*P(2)*...*P(z-1)
!> \endverbatim
!>
!> \param[in] M
!> \verbatim
!>          M is INTEGER
!>          The number of rows of the matrix A.  If m <= 1, an immediate
!>          return is effected.
!> \endverbatim
!>
!> \param[in] N
!> \verbatim
!>          N is INTEGER
!>          The number of columns of the matrix A.  If n <= 1, an
!>          immediate return is effected.
!> \endverbatim
!>
!> \param[in] C
!> \verbatim
!>          C is DOUBLE PRECISION array, dimension
!>                  (M-1) if SIDE = 'L'
!>                  (N-1) if SIDE = 'R'
!>          The cosines c(k) of the plane rotations.
!> \endverbatim
!>
!> \param[in] S
!> \verbatim
!>          S is DOUBLE PRECISION array, dimension
!>                  (M-1) if SIDE = 'L'
!>                  (N-1) if SIDE = 'R'
!>          The sines s(k) of the plane rotations.  The 2-by-2 plane
!>          rotation part of the matrix P(k), R(k), has the form
!>          R(k) = (  c(k)  s(k) )
!>                 ( -s(k)  c(k) ).
!> \endverbatim
!>
!> \param[in,out] A
!> \verbatim
!>          A is DOUBLE PRECISION array, dimension (LDA,N)
!>          The M-by-N matrix A.  On exit, A is overwritten by P*A if
!>          SIDE = 'L' or by A*P**T if SIDE = 'R'.
!> \endverbatim
!>
!> \param[in] LDA
!> \verbatim
!>          LDA is INTEGER
!>          The leading dimension of the array A.  LDA >= max(1,M).
!> \endverbatim
!
!  Authors:
!  ========
!
!> \author Univ. of Tennessee
!> \author Univ. of California Berkeley
!> \author Univ. of Colorado Denver
!> \author NAG Ltd.
!
!> \date December 2016
!
!> \ingroup OTHERauxiliary
!
!  =====================================================================
      SUBROUTINE DLASR(Side,Pivot,Direct,M,N,C,S,A,Lda)
      IMPLICIT NONE
!*--DLASR203
!
!  -- LAPACK auxiliary routine (version 3.7.0) --
!  -- LAPACK is a software package provided by Univ. of Tennessee,    --
!  -- Univ. of California Berkeley, Univ. of Colorado Denver and NAG Ltd..--
!     December 2016
!
!     .. Scalar Arguments ..
      CHARACTER Direct , Pivot , Side
      INTEGER Lda , M , N
!     ..
!     .. Array Arguments ..
      DOUBLE PRECISION A(Lda,*) , C(*) , S(*)
!     ..
!
!  =====================================================================
!
!     .. Parameters ..
      DOUBLE PRECISION ONE , ZERO
      PARAMETER (ONE=1.0D+0,ZERO=0.0D+0)
!     ..
!     .. Local Scalars ..
      INTEGER i , info , j
      DOUBLE PRECISION ctemp , stemp , temp
!     ..
!     .. External Functions ..
      LOGICAL LSAME
      EXTERNAL LSAME
!     ..
!     .. External Subroutines ..
      EXTERNAL XERBLA
!     ..
!     .. Intrinsic Functions ..
      INTRINSIC MAX
!     ..
!     .. Executable Statements ..
!
!     Test the input parameters
!
      info = 0
      IF ( .NOT.(LSAME(Side,'L') .OR. LSAME(Side,'R')) ) THEN
         info = 1
      ELSEIF ( .NOT.(LSAME(Pivot,'V') .OR. LSAME(Pivot,'T') .OR.        &
     &         LSAME(Pivot,'B')) ) THEN
         info = 2
      ELSEIF ( .NOT.(LSAME(Direct,'F') .OR. LSAME(Direct,'B')) ) THEN
         info = 3
      ELSEIF ( M<0 ) THEN
         info = 4
      ELSEIF ( N<0 ) THEN
         info = 5
      ELSEIF ( Lda<MAX(1,M) ) THEN
         info = 9
      ENDIF
      IF ( info/=0 ) THEN
         CALL XERBLA('DLASR ',info)
         RETURN
      ENDIF
!
!     Quick return if possible
!
      IF ( (M==0) .OR. (N==0) ) RETURN
      IF ( LSAME(Side,'L') ) THEN
!
!        Form  P * A
!
         IF ( LSAME(Pivot,'V') ) THEN
            IF ( LSAME(Direct,'F') ) THEN
               DO j = 1 , M - 1
                  ctemp = C(j)
                  stemp = S(j)
                  IF ( (ctemp/=ONE) .OR. (stemp/=ZERO) ) THEN
                     DO i = 1 , N
                        temp = A(j+1,i)
                        A(j+1,i) = ctemp*temp - stemp*A(j,i)
                        A(j,i) = stemp*temp + ctemp*A(j,i)
                     ENDDO
                  ENDIF
               ENDDO
            ELSEIF ( LSAME(Direct,'B') ) THEN
               DO j = M - 1 , 1 , -1
                  ctemp = C(j)
                  stemp = S(j)
                  IF ( (ctemp/=ONE) .OR. (stemp/=ZERO) ) THEN
                     DO i = 1 , N
                        temp = A(j+1,i)
                        A(j+1,i) = ctemp*temp - stemp*A(j,i)
                        A(j,i) = stemp*temp + ctemp*A(j,i)
                     ENDDO
                  ENDIF
               ENDDO
            ENDIF
         ELSEIF ( LSAME(Pivot,'T') ) THEN
            IF ( LSAME(Direct,'F') ) THEN
               DO j = 2 , M
                  ctemp = C(j-1)
                  stemp = S(j-1)
                  IF ( (ctemp/=ONE) .OR. (stemp/=ZERO) ) THEN
                     DO i = 1 , N
                        temp = A(j,i)
                        A(j,i) = ctemp*temp - stemp*A(1,i)
                        A(1,i) = stemp*temp + ctemp*A(1,i)
                     ENDDO
                  ENDIF
               ENDDO
            ELSEIF ( LSAME(Direct,'B') ) THEN
               DO j = M , 2 , -1
                  ctemp = C(j-1)
                  stemp = S(j-1)
                  IF ( (ctemp/=ONE) .OR. (stemp/=ZERO) ) THEN
                     DO i = 1 , N
                        temp = A(j,i)
                        A(j,i) = ctemp*temp - stemp*A(1,i)
                        A(1,i) = stemp*temp + ctemp*A(1,i)
                     ENDDO
                  ENDIF
               ENDDO
            ENDIF
         ELSEIF ( LSAME(Pivot,'B') ) THEN
            IF ( LSAME(Direct,'F') ) THEN
               DO j = 1 , M - 1
                  ctemp = C(j)
                  stemp = S(j)
                  IF ( (ctemp/=ONE) .OR. (stemp/=ZERO) ) THEN
                     DO i = 1 , N
                        temp = A(j,i)
                        A(j,i) = stemp*A(M,i) + ctemp*temp
                        A(M,i) = ctemp*A(M,i) - stemp*temp
                     ENDDO
                  ENDIF
               ENDDO
            ELSEIF ( LSAME(Direct,'B') ) THEN
               DO j = M - 1 , 1 , -1
                  ctemp = C(j)
                  stemp = S(j)
                  IF ( (ctemp/=ONE) .OR. (stemp/=ZERO) ) THEN
                     DO i = 1 , N
                        temp = A(j,i)
                        A(j,i) = stemp*A(M,i) + ctemp*temp
                        A(M,i) = ctemp*A(M,i) - stemp*temp
                     ENDDO
                  ENDIF
               ENDDO
            ENDIF
         ENDIF
      ELSEIF ( LSAME(Side,'R') ) THEN
!
!        Form A * P**T
!
         IF ( LSAME(Pivot,'V') ) THEN
            IF ( LSAME(Direct,'F') ) THEN
               DO j = 1 , N - 1
                  ctemp = C(j)
                  stemp = S(j)
                  IF ( (ctemp/=ONE) .OR. (stemp/=ZERO) ) THEN
                     DO i = 1 , M
                        temp = A(i,j+1)
                        A(i,j+1) = ctemp*temp - stemp*A(i,j)
                        A(i,j) = stemp*temp + ctemp*A(i,j)
                     ENDDO
                  ENDIF
               ENDDO
            ELSEIF ( LSAME(Direct,'B') ) THEN
               DO j = N - 1 , 1 , -1
                  ctemp = C(j)
                  stemp = S(j)
                  IF ( (ctemp/=ONE) .OR. (stemp/=ZERO) ) THEN
                     DO i = 1 , M
                        temp = A(i,j+1)
                        A(i,j+1) = ctemp*temp - stemp*A(i,j)
                        A(i,j) = stemp*temp + ctemp*A(i,j)
                     ENDDO
                  ENDIF
               ENDDO
            ENDIF
         ELSEIF ( LSAME(Pivot,'T') ) THEN
            IF ( LSAME(Direct,'F') ) THEN
               DO j = 2 , N
                  ctemp = C(j-1)
                  stemp = S(j-1)
                  IF ( (ctemp/=ONE) .OR. (stemp/=ZERO) ) THEN
                     DO i = 1 , M
                        temp = A(i,j)
                        A(i,j) = ctemp*temp - stemp*A(i,1)
                        A(i,1) = stemp*temp + ctemp*A(i,1)
                     ENDDO
                  ENDIF
               ENDDO
            ELSEIF ( LSAME(Direct,'B') ) THEN
               DO j = N , 2 , -1
                  ctemp = C(j-1)
                  stemp = S(j-1)
                  IF ( (ctemp/=ONE) .OR. (stemp/=ZERO) ) THEN
                     DO i = 1 , M
                        temp = A(i,j)
                        A(i,j) = ctemp*temp - stemp*A(i,1)
                        A(i,1) = stemp*temp + ctemp*A(i,1)
                     ENDDO
                  ENDIF
               ENDDO
            ENDIF
         ELSEIF ( LSAME(Pivot,'B') ) THEN
            IF ( LSAME(Direct,'F') ) THEN
               DO j = 1 , N - 1
                  ctemp = C(j)
                  stemp = S(j)
                  IF ( (ctemp/=ONE) .OR. (stemp/=ZERO) ) THEN
                     DO i = 1 , M
                        temp = A(i,j)
                        A(i,j) = stemp*A(i,N) + ctemp*temp
                        A(i,N) = ctemp*A(i,N) - stemp*temp
                     ENDDO
                  ENDIF
               ENDDO
            ELSEIF ( LSAME(Direct,'B') ) THEN
               DO j = N - 1 , 1 , -1
                  ctemp = C(j)
                  stemp = S(j)
                  IF ( (ctemp/=ONE) .OR. (stemp/=ZERO) ) THEN
                     DO i = 1 , M
                        temp = A(i,j)
                        A(i,j) = stemp*A(i,N) + ctemp*temp
                        A(i,N) = ctemp*A(i,N) - stemp*temp
                     ENDDO
                  ENDIF
               ENDDO
            ENDIF
         ENDIF
      ENDIF
!
!
!     End of DLASR
!
      END SUBROUTINE DLASR


!	Find posterior distribution for latent vectors.

!	catobsrange is used to relate observed and underlying categories.
!	Responses are in dat.
!	maxita is used to determine the number of iterations to use to find the maximum posterior density.
!	numcat provides ranges for underlying observations,
!	numcatobs provides ranges of observations.
!	mask indicates items to use for computations.
!	normal determines if a normal or multinomial analysis is used.
!	beta is the parameter vector.
!	changemin is used for minimum change.
!	Predictors are in indvar.
!	lintran is the linear transformation for the latent vector.
!	maxdalpha is used to limit changes in alpha approximations.
!	quadpoint provides quadrature points.
!	quadweight provides quadrature weights.
!	tau controls shrinking of step size.
!	tola is used to determine the desired accuracy of the maximization.
!	tolsing is used to determine the tolerance for the Cholesky decomposition.
!	alpha contains locations of maxima.
!	cholnhess contains modified Cholesky decompositions.
!	postdensity gives weights at quadrature points.
!	theta are quadrature points.
subroutine posterior(catobsrange,dat,maxita,numcat,numcatobs,mask,normal,beta,changemin,indvar,lintran,&
	maxdalpha,quadpoint,quadweight,&
	tau,tola,tolsing,alpha,cholnhess,postdensity,theta)
implicit none
interface

	function chol(sym,tolsing)
!	chol is used to compute a modified Cholesky decomposition of the n by n matrix sym.
!	The singularity tolerance is tolsing.
		implicit none
		real(kind=8),intent(in)::sym(:,:),tolsing
	 	real(kind=8)::chol(size(sym,1),size(sym,1))
	end function chol

!	Get normal density for scaled latent vector.  Place in density.
!	Input involves linear component linth, quadrature points quadpoint, quadratic components quadth, quadrature weights quadweight, and scale
!	factor scaletheta.
	function density(linth,quadpoint,quadth,quadweight,scaletheta)
		implicit none
		real(kind=8),intent(in)::linth(:),quadpoint(:,:),quadth(:,:),quadweight(:),scaletheta
		real(kind=8)::density(size(quadweight))
	end function density
!	Get multinomial probabilities for latent vector.  Place in density.
!	Input involves linear component linth, quadrature points quadpoint, quadratic components quadth, and quadrature weights quadweight.
	function densitym(linth,quadpoint,quadth,quadweight)
		real(kind=8),intent(in)::linth(:),quadpoint(:,:),quadth(:,:),quadweight(:)
		real(kind=8)::densitym(size(quadweight))
	end function densitym

!	convert an n by n symmetric matrix from compact to regular form.
	function expandmat(n,compmat)
		implicit none
		integer,intent(in)::n
		real(kind=8),intent(in)::compmat(:)
		real(kind=8)::expandmat(n,n)
	end function expandmat

!	Extract the linear component lintheta and the quadratic component quadtheta from the parameter vector beta.
	subroutine getlinquad(beta,lintheta,quadtheta)
		implicit none
		real(kind=8),intent(in)::beta(:)
		real(kind=8),intent(out)::lintheta(:,:),quadtheta(:,:,:)
	end subroutine getlinquad

!	Extract the linear component linth and the quadratic component quadth from lintheta and quadtheta for a given predictor thpr.
	subroutine getlinquadth(lintheta,quadtheta,thpr,linth,quadth)
		implicit none
		real(kind=8),intent(in)::lintheta(:,:),quadtheta(:,:,:),thpr(:)
		real(kind=8),intent(out)::linth(:),quadth(:,:)
	end subroutine getlinquadth
!
!	Find location of maximum posterior density of the latent vector.
!	catobsrange is used to relate observed and underlying categories.
!	maxita is used to determine the number of iterations to use to find the maximum posterior density.
!	maxitb is used for subiterations.
!	npred is the number of predictors.
!	numcat provides ranges for underlying observations,
!	numcatobs provides ranges of observations, mask indicates items to use for computations.
!	obsmask is an observation mask that can be used to restrict attention to a portion of the response.
!	beta is the parameter vector.
!	changemin is used for minimum change.
!	linth is the linear component of the log density.
!	lintran is the linear transformation of the latent vector.
!	maxdtheta is the maximum permitted change in the approximation of the maximum in one step.
!	quadth is the quadratic component of the log density.
!	resp is the response.
!	tau is used for minimum step sizes.
!	tol is used to determine the desired accuracy of the maximization.
!	tolsing is used to determine the tolerance for the Cholesky decomposition.
!	nhchol is the modified Cholesky decomposition of the negative hessian.
!	theta is the maximum.
	subroutine maxpost(catobsrange,maxita,maxitb,npred,numcat,numcatobs,resp,obsmask,beta,changemin,&
		linth,lintran,maxdtheta,quadth,tau,tol,tolsing,&
		nhchol,theta)
		implicit none
		integer,intent(in)::catobsrange(:,:),maxita,maxitb,npred,numcat(:),numcatobs(:),resp(:)
		logical,intent(in)::obsmask(:)
		real(kind=8),intent(in)::beta(:),changemin,lintran(:,:),linth(:),maxdtheta,quadth(:,:),tau,tol,tolsing
		real(kind=8),intent(inout)::theta(:)
		real(kind=8),intent(out)::nhchol(:,:)
	end subroutine maxpost

!	probvec is used to compute underlying item probabilities.
!   	numcat indicates the number of underlying categories per item.
!	mask indicates which items were presented.
!	locations provides location parameters and scales provides scale factors.
!	theta is the value of the latent vector.
	function probvec(numcat,mask,locations,scales,theta)
		implicit none
		integer,intent(in)::numcat(:)
		logical,intent(in)::mask(:)
		real(kind=8),intent(in)::locations(:),scales(:,:),theta(:)
		real(kind=8)::probvec(size(locations))
	end function probvec

!	probvecobs is used to compute observed item probabilities given the latent vector.
!	catobsrange is the table of ranges of underlying categories per observed category.
!	numcatobs provides the number of observed categories per item.
!	resp is the response vector.
!	mask indicates which items were presented.
!	probcat is the vector of underlying item probabilities.
	function probvecobs(catobsrange,numcatobs,resp,mask,probcat)
		implicit none
		integer,intent(in)::catobsrange(:,:),numcatobs(:),resp(:)
		logical,intent(in)::mask(:)
		real(kind=8),intent(in)::probcat(:)
		real(kind=8)::probvecobs(size(mask))
	end function probvecobs

!	Convert from adaptive quadrature scale to normalized latent vector.
!	Use nhchol for modified Cholesky decomposition of negative hessian.
!	Use oldalpha for location.
!	Use scalez for scaling.
!	Quadrature point is z.
	function scaleadapt(nhchol,oldalpha,scalez,z)
		implicit none
		real(kind=8),intent(in)::nhchol(:,:),oldalpha(:),scalez(:),z(:)
		real(kind=8)::scaleadapt(size(z))
	end function scaleadapt

!		This program uses the modified Cholesky decomposition
!		in lu to solve the equation lusolve = b.
	function solve(lu,b)
		implicit none

		real(kind=8),intent(in)::lu(:,:)
		real(kind=8),intent(in)::b(:)
		real(kind=8)::solve(size(b))
	end function solve
end interface
integer,intent(in)::catobsrange(:,:),dat(:,:),maxita,numcat(:),numcatobs(:)
logical,intent(in)::mask(:),normal
real(kind=8),intent(in)::beta(:),changemin,indvar(:,:),lintran(:,:),&
	maxdalpha,quadpoint(:,:),quadweight(:),&
	tau,tola,tolsing
real(kind=8),intent(inout)::alpha(:,:),cholnhess(:,:,:)
real(kind=8),intent(out)::postdensity(:,:),theta(:,:,:)

!	dimlatin is the dimension of the latent vector.
!	dimlatout is the dimension of the transformed latent vector.

!	ncat is the total number of underlying categories.
!	nitems is the number of items, and item is an item.
!	nobs is number of observations.
!	npred is the number of predictors.
!	nquad is the number of quadrature points.
!	nscale is the pointer for scale parameters.
!	nscale1 points to the end of the scale parameters.
!	obs is observation number.
!	pred is a predictor counter.
!	quad counts quadrature points.
!	resp is a single response.
!	row is a row counter.
integer::dimlatin,dimlatout,item,ncat,nitems,&
	nobs,npred,nquad,nscale,&
	nscale1,obs,pred,quad,&
	resp(size(mask)),row
!	Data mask.
logical::datamask(size(mask))
!	alpha is used to find the maximum posterior density.
!	locations is for location parameters.
!	linth is a linear component of the log density of a latent vector.
!	lintheta is for linear parameters for the log density.
!	meanth is the mean of the latent vector.
!	newtheta is lintran(theta).
!	nhchol is a Cholesky decomposition of -2*quadth.
!	priordensity provides a prior density.
!	probcat is used for underlying marginal conditional probabilities.
!	probcatobs is used for observed marginal conditional probabilities.
!	probobs is the probability of an observation.
!	quadth is a quadratic component of the log density of a latent vector.
!	quadtheta is for quadratic parameters for the log density.
!	scales is for scale parameters.
!	scaletheta is used for scaling.
!	scalez is used for scaling.
!	thpr is the predictor vector for an observation.
!	z is a transformation of a latent vector value.
real(kind=8)::linth(size(quadpoint,1)),&
	lintheta(size(quadpoint,1),size(indvar,1)),locations(sum(numcat)),&
	newtheta(size(lintran,1)),meanth(size(quadpoint,1)),&
	nhchol(size(quadpoint,1),size(quadpoint,1)),&
	priordensity(size(quadweight)),probcat(sum(numcat)),probcatobs(size(mask)),probobs,&
	quadth(size(quadpoint,1),size(quadpoint,1)),&
	quadtheta(size(quadpoint,1),size(quadpoint,1),size(indvar,1)),&
	scales(size(lintran,1),sum(numcat)),scaletheta,&
	scalez(size(quadpoint,1)),&
	thpr(size(indvar,1)),z(size(quadpoint,1))

!	Set up parameter arrays for simplified processing.
dimlatin=size(quadpoint,1)

dimlatout=size(lintran,1)
ncat=sum(numcat)
nitems=size(mask)
npred=size(indvar,1)

nobs=size(dat,2)
!	Item scale parameters.
nscale=ncat+1

nscale1=ncat*(1+dimlatout)
nquad=size(quadweight)
!	Parameters for quadratic component of log density.

locations=beta(1:ncat)
scales=reshape(beta(nscale:nscale1),(/dimlatout,ncat/))

call getlinquad(beta,lintheta,quadtheta)

!	Obtain results for each observation.
do obs=1,nobs

! Observation obs.
	resp=dat(:,obs)

!	Density specification for observation.
	thpr=indvar(:,obs)

	call getlinquadth(lintheta,quadtheta,thpr,linth,quadth)

!	Mask
	datamask=.false.
	do item=1,nitems
		if(resp(item)>=0.and.resp(item)<numcatobs(item).and.mask(item))datamask(item)=.true.
	end do

!	The normal versus multinomial cases differ here.
	if(normal)then
!	Mean and covariance matrix of latent vector.

		nhchol=chol(-2.0_8*quadth,tolsing)
!	Check that covariance matrix is positive definite.
		if(minval((/(nhchol(row,row),row=1,dimlatin)/))<=0.0_8)stop "Unacceptable covariance matrix"
		meanth=solve(nhchol,linth)
		scaletheta=sqrt(product((/(nhchol(row,row),row=1,dimlatin)/)))
		scaletheta=scaletheta*exp(-0.5_8*dot_product(linth,meanth))

!	See if adaptive approach to be used.
		if(maxita>0) then

			call maxpost(catobsrange,maxita,maxita,npred,numcat,numcatobs,resp,datamask,beta,changemin,&
				linth,lintran,maxdalpha,quadth,tau,tola,tolsing,&
				nhchol,alpha(:,obs))
			cholnhess(:,:,obs)=nhchol
		end if
		nhchol=cholnhess(:,:,obs)
		do row=1,dimlatin
			scalez(row)=1.0_8/sqrt(nhchol(row,row))
		end do
		scaletheta=scaletheta*product(scalez)
		do quad=1,nquad
				theta(:,quad,obs)=scaleadapt(nhchol,alpha(:,obs),scalez,quadpoint(:,quad))
		end do
		priordensity=density(linth,theta(:,:,obs),quadth,quadweight,scaletheta)
	else
		theta(:,:,obs)=quadpoint
		priordensity=densitym(linth,quadpoint,quadth,quadweight)
	end if
!	Cycle through quadrature points.

	do quad=1,nquad

! Linear transformation of latent vector.
		newtheta=matmul(lintran,theta(:,quad,obs))

! Probability computation.

		probcat=probvec(numcat,datamask,locations,scales,newtheta)
		probcatobs=probvecobs(catobsrange,numcatobs,resp,datamask,probcat)
		postdensity(quad,obs)=priordensity(quad)
		if(any(datamask))postdensity(quad,obs)=postdensity(quad,obs)*product(probcatobs,datamask)
	end do
	probobs=sum(postdensity(:,obs))
	postdensity(:,obs)=postdensity(:,obs)/probobs
end do
return
end subroutine posterior


module complex_transform_routines

    use auxiliary_routines

    use fftpack_precision, only: &
        wp, & ! working precision
        ip ! integer precision

    ! Explicit typing only
    implicit none

    ! Everything is private unless stated otherwise
    private
    public :: cfft1i, cfft1b, cfft1f, cfft2i, cfft2b, cfft2f, cfftmi, cfftmb, cfftmf

    interface
        !
        !==> 1D complex initialization
        !
        module subroutine cfft1i(n, wsave, lensav, ierror)
            !
            ! cfft1i: initialization for cfft1b and cfft1f.
            !
            !  Purpose:
            !
            !  cfft1i initializes array wsave for use in its companion routines
            !  cfft1b and cfft1f. Routine cfft1i must be called before the first
            !  call to cfft1b or cfft1f, and after whenever the value of integer
            !  n changes.
            !
            !  Parameters:
            !
            !  input,
            !  n, the length of the sequence to be
            !  transformed.  the transform is most efficient when n is a product
            !  of small primes.
            !
            !  input
            !  lensav, the dimension of the wsave array.
            !  lensav must be at least 2*n + int(log(real(n))) + 4.
            !
            !  output,
            !  wsave(lensav), containing the prime factors
            !  of n and  also containing certain trigonometric values which will be used
            !  in routines cfft1b or cfft1f.
            !
            !  output
            !  ierror, error_flag.
            !  0, successful exit;
            !  2, input parameter lensav not big enough.
            !
            !--------------------------------------------------------------
            ! Dummy arguments
            !--------------------------------------------------------------
            integer (ip), intent (in)  :: n
            real (wp),    intent (out) :: wsave(lensav)
            integer (ip), intent (in)  :: lensav
            integer (ip), intent (out) :: ierror
            !--------------------------------------------------------------
        end subroutine cfft1i
        !
        !==> 1D complex backward
        !
        module subroutine cfft1b(n, inc, c, lenc, wsave, lensav, work, lenwrk, ierror)

            use, intrinsic :: ISO_C_binding, only: c_f_pointer, c_loc
            !
            !  input
            !  integer n, the length of the sequence to be
            !  transformed.  the transform is most efficient when n is a product of
            !  small primes.
            !
            !  integer inc, the increment between the locations, in
            !  array c, of two consecutive elements within the sequence to be transformed.
            !
            !  integer lenc, the dimension of the c array.
            !  lenc must be at least inc*(n-1) + 1.
            !
            !  real wsave(lensav). wsave's contents must be initialized with a call
            !  to cfft1i before the first call to routine cfft1f
            !  or cfft1b for a given transform length n.  wsave's contents may be
            !  re-used for subsequent calls to cfft1f and cfft1b with the same n.
            !
            !  integer lensav, the dimension of the wsave array.
            !  lensav must be at least 2*n + int(log(real(n))) + 4.
            !
            !
            !  input lenwrk, the dimension of the work array.
            !  lenwrk must be at least 2*n.
            !
            !  input/output
            !  complex c(lenc) containing the sequence to be
            !  transformed.
            !
            !  real workspace work(lenwrk).
            !
            !  integer ier, error_flag.
            !  0, successful exit;
            !  1, input parameter lenc not big enough;
            !  2, input parameter lensav not big enough;
            !  3, input parameter lenwrk not big enough;
            !  20, input error returned by lower level routine.
            !
            !--------------------------------------------------------------
            ! Dummy arguments
            !--------------------------------------------------------------
            integer (ip), intent (in)             :: n
            integer (ip), intent (in)             :: inc
            complex (wp), intent (in out), target :: c(lenc)
            integer (ip), intent (in)             :: lenc
            integer (ip), intent (in)             :: lensav
            integer (ip), intent (in)             :: lenwrk
            integer (ip), intent (out)            :: ierror
            real (wp),    intent (in out)         :: work(lenwrk)
            real (wp),    intent (in out)         :: wsave(lensav)
            !--------------------------------------------------------------
        end subroutine cfft1b
            !
            !==> 1D complex forward
            !
        module subroutine cfft1f(n, inc, c, lenc, wsave, lensav, work, lenwrk, ierror)
            use, intrinsic :: ISO_C_binding, only: c_f_pointer, c_loc
            !
            ! cfft1f: complex forward fast fourier transform, 1d.
            !
            !  purpose:
            !
            !  cfft1f computes the one-dimensional fourier transform of a single
            !  periodic sequence within a complex array. this transform is referred
            !  to as the forward transform or fourier analysis, transforming the
            !  sequence from physical to spectral space.
            !
            !  this transform is normalized since a call to cfft1f followed
            !  by a call to cfft1b (or vice-versa) reproduces the original
            !  array within roundoff error.
            !
            !  input
            !
            !  integer n, the length of the sequence to be
            !  transformed. the transform is most efficient when
            !  n is a product of small primes.
            !
            !  integer inc, the increment between the locations, in
            !  array c, of two consecutive elements within the sequence to be transformed.
            !
            !  real wsave(lensav).  wsave's contents must be
            !  initialized with a call to cfft1i before the first call to routine cfft1f
            !  or cfft1b for a given transform length n.  wsave's contents may be re-used
            !  for subsequent calls to cfft1f and cfft1b with the same n.
            !
            !  integer lensav, the dimension of the wsave array.
            !  lensav must be at least 2*n + int(log(real(n))) + 4.
            !
            !  integer lenwrk, the dimension of the work array.
            !  lenwrk must be at least 2*n.
            !
            !  input/output
            !  complex c(lenc) containing the sequence to be transformed.
            !
            !  real work(lenwrk), workspace array
            !  integer lenc, the dimension of the c array.
            !  lenc must be at least inc*(n-1) + 1.
            !
            !  output
            !  integer ier, error_flag.
            !  0, successful exit;
            !  1, input parameter lenc not big enough;
            !  2, input parameter lensav not big enough;
            !  3, input parameter lenwrk not big enough;
            !  20, input error returned by lower level routine.
            !
            !------------------------------------------------------------------
            ! Dummy arguments
            !------------------------------------------------------------------
            integer (ip), intent (in)             :: n
            integer (ip), intent (in)             :: inc
            complex (wp), intent (in out), target :: c(lenc)
            integer (ip), intent (in)             :: lenc
            real (wp),    intent (in)             :: wsave(lensav)
            integer (ip), intent (in)             :: lensav
            real (wp),    intent (out)            :: work(lenwrk)
            integer (ip), intent (in)             :: lenwrk
            integer (ip), intent (out)            :: ierror
            !------------------------------------------------------------------
        end subroutine cfft1f
        !
        !==> 2D complex initialization
        !
        module subroutine cfft2i(l, m, wsave, lensav, ierror)
            !
            ! cfft2i: initialization for cfft2b and cfft2f.
            !
            !  purpose:
            !
            !  cfft2i initializes real array wsave for use in its companion
            !  routines cfft2f and cfft2b for computing two-dimensional fast
            !  fourier transforms of complex data.  prime factorizations of l and m,
            !  together with tabulations of the trigonometric functions, are
            !  computed and stored in array wsave.
            !
            !  on 10 may 2010, this code was modified by changing the value
            !  of an index into the wsave array.
            !
            !  input
            !  integer l, the number of elements to be transformed
            !  in the first dimension. the transform is most efficient when l is a
            !  product of small primes.
            !
            !  integer m, the number of elements to be transformed
            !  in the second dimension. the transform is most efficient when m is a
            !  product of small primes.
            !
            !  integer lensav, the dimension of the wsave array.
            !  lensav must be at least
            !
            !  2*(l+m) + int(log(real(l)))+ int(log(real(m))) + 8.
            !
            !  output
            !  real wsave(lensav), contains the prime factors of l
            !  and m, and also certain trigonometric values which will be used in
            !  routines cfft2b or cfft2f.
            !
            !  integer  ierror, error_flag.
            !  0, successful exit;
            !  2, input parameter lensav not big enough;
            !  20, input error returned by lower level routine.
            !
            !--------------------------------------------------------------
            ! Dummy arguments
            !--------------------------------------------------------------
            integer (ip), intent (in)  :: l
            integer (ip), intent (in)  :: m
            real (wp),    intent (out) :: wsave(lensav)
            integer (ip), intent (in)  :: lensav
            integer (ip), intent (out) :: ierror
            !--------------------------------------------------------------
        end subroutine cfft2i
        !
        !==> 2D complex backward
        !
        module subroutine cfft2b(ldim, l, m, c, wsave, lensav, work, lenwrk, ierror)
            !
            ! cfft2b: complex fast ror transform, 2d.
            !
            !  purpose:
            !
            !  cfft2b computes the two-dimensional discrete ror transform of a
            !  complex periodic array.  this transform is known as the backward
            !  transform or synthesis, transforming from spectral to
            !  physical space.  routine cfft2b is normalized, in that a call to
            !  cfft2b followed by a call to cfft2f (or vice-versa) reproduces the
            !  original array within roundoff error.
            !
            !  bug fix
            !  on 10 may 2010, this code was modified by changing the value
            !  of an index into the wsave array.
            !
            !  parameters:
            !
            !  input
            !  integer ldim, the first dimension of c.
            !
            !  integer l, the number of elements to be transformed
            !  in the first dimension of the two-dimensional complex array c.  the value
            !  of l must be less than or equal to that of ldim.  the transform is
            !  most efficient when l is a product of small primes.
            !
            !  integer m, the number of elements to be transformed in
            !  the second dimension of the two-dimensional complex array c.  the transform
            !  is most efficient when m is a product of small primes.
            !
            !  real wsave(lensav). wsave's contents must be initialized with
            !  a call to cfft2i before the first call to routine cfft2f
            !  or cfft2b with transform lengths l and m.  wsave's contents may be
            !  re-used for subsequent calls to cfft2f and cfft2b with the same
            !  transform lengths l and m.
            !
            !  integer lensav, the dimension of the wsave array.
            !  lensav must be at least
            !
            !  2*(l+m) + int(log(real(l))) + int(log(real(m))) + 8.
            !
            !  real work(lenwrk).
            !
            !  integer lenwrk, the dimension of the work array.
            !  lenwrk must be at least 2*l*m.
            !
            !  input/output
            !  complex c(ldim,m), on intput, the array of two dimensions
            !  containing the (l,m) subarray to be transformed.  on
            !  output, the transformed data.
            !
            !
            !  output
            !
            !  integer ierror, the error_flag.
            !
            !  0, successful exit;
            !  2, input parameter lensav not big enough;
            !  3, input parameter lenwrk not big enough;
            !  5, input parameter ldim < l;
            !  20, input error returned by lower level routine.
            !
            !------------------------------------------------------------------
            ! Dummy arguments
            !------------------------------------------------------------------
            integer (ip), intent (in)     :: ldim
            integer (ip), intent (in)     :: l
            integer (ip), intent (in)     :: m
            complex (wp), intent (in out) :: c(ldim,m)
            real (wp),    intent (in)     :: wsave(lensav)
            integer (ip), intent (in)     :: lensav
            real (wp),    intent (in out) :: work(lenwrk)
            integer (ip), intent (in)     :: lenwrk
            integer (ip), intent (out)    :: ierror
            !------------------------------------------------------------------
        end subroutine cfft2b
        !
        !==> 2D complex forward
        !
        module subroutine cfft2f(ldim, l, m, c, wsave, lensav, work, lenwrk, ierror)
            !
            ! cfft2f: complex forward fast fourier transform, 2d.
            !
            !  Purpose:
            !
            !  cfft2f computes the two-dimensional discrete fourier transform of
            !  a complex periodic array. This transform is known as the forward
            !  transform or fourier analysis, transforming from physical to
            !  spectral space. routine cfft2f is normalized, in that a call to
            !  cfft2f followed by a call to cfft2b (or vice-versa) reproduces the
            !  original array within roundoff error.
            !
            !  BUG FIX
            !  On 10 May 2010, this code was modified by changing the value
            !  of an index into the wsave array.
            !
            !  INPUT
            !
            !  integer ldim, the first dimension of the array c.
            !
            !  integer l, the number of elements to be transformed
            !  in the first dimension of the two-dimensional complex array c.  the value
            !  of l must be less than or equal to that of ldim.  the transform is most
            !  efficient when l is a product of small primes.
            !
            !  integer m, the number of elements to be transformed
            !  in the second dimension of the two-dimensional complex array c.  the
            !  transform is most efficient when m is a product of small primes.
            !
            !  real wsave(lensav). wsave's contents must be
            !  initialized with a call to cfft2i before the first call to routine cfft2f
            !  or cfft2b with transform lengths l and m.  wsave's contents may be re-used
            !  for subsequent calls to cfft2f and cfft2b having those same
            !  transform lengths.
            !
            !  integer lensav, the dimension of the wsave array.
            !  lensav must be at least
            !
            !  2*(l+m) + int(log(real(l))) + int(log(real(m))) + 8.
            !
            !
            !  integer lenwrk, the dimension of the work array.
            !  lenwrk must be at least 2*l*m.
            !
            !  INPUT/OUTPUT
            !  complex c(ldim,m), on input, the array of two
            !  dimensions containing the (l,m) subarray to be transformed. On output, the
            !  transformed data.
            !
            !  real work(lenwrk), workspace array
            !
            !  OUTPUT
            !  integer ierror, error_flag.
            !  0, successful exit;
            !  2, input parameter lensav not big enough;
            !  3, input parameter lenwrk not big enough;
            !  5, input parameter ldim < l;
            !  20, input error returned by lower level routine.
            !
            !------------------------------------------------------------------
            ! Dummy arguments
            !------------------------------------------------------------------
            integer (ip), intent (in)     :: ldim
            integer (ip), intent (in)     :: l
            integer (ip), intent (in)     :: m
            complex (wp), intent (in out) :: c(ldim,m)
            real (wp),    intent (in)     :: wsave(lensav)
            integer (ip), intent (in)     :: lensav
            real (wp),    intent (in out) :: work(lenwrk)
            integer (ip), intent (in)     :: lenwrk
            integer (ip), intent (out)    :: ierror
            !------------------------------------------------------------------
        end subroutine cfft2f
        !
        !==> multiple complex initialization
        !
        module subroutine cfftmi(n, wsave, lensav, ierror)
            !
            !  cfftmi: initialization for cfftmb and cfftmf.
            !
            !  purpose:
            !
            !  cfftmi initializes array wsave for use in its companion routines
            !  cfftmb and cfftmf.  cfftmi must be called before the first call
            !  to cfftmb or cfftmf, and after whenever the value of integer n changes.
            !
            !  parameters:
            !
            !  integer n, the length of each sequence to be
            !  transformed.  the transform is most efficient when n is a product of
            !  small primes.
            !
            !  integer lensav, the dimension of the wsave array.
            !  lensav must be at least 2*n + int(log(real(n))) + 4.
            !
            !  real wsave(lensav), containing the prime factors
            !  of n and also containing certain trigonometric values which will be used in
            !  routines cfftmb or cfftmf.
            !
            !  integer ierror, error_flag.
            !  0, successful exit;
            !  2, input parameter lensav not big enough.
            !
            !--------------------------------------------------------------
            ! Dummy arguments
            !--------------------------------------------------------------
            integer (ip), intent (in)  :: n
            real (wp),    intent (out) :: wsave(lensav)
            integer (ip), intent (in)  :: lensav
            integer (ip), intent (out) :: ierror
            !--------------------------------------------------------------
        end subroutine cfftmi
        !
        !==> multiple complex backward
        !
        module subroutine cfftmb(lot, jump, n, inc, c, lenc, wsave, lensav, work, lenwrk, ierror)
            !
            ! cfftmb: complex backward fft, 1d, multiple vectors.
            !
            !  purpose:
            !
            !  cfftmb computes the one-dimensional fourier transform of multiple
            !  periodic sequences within a complex array.  this transform is referred
            !  to as the backward transform or fourier synthesis, transforming the
            !  sequences from spectral to physical space.  this transform is
            !  normalized since a call to cfftmf followed by a call to cfftmb (or
            !  vice-versa) reproduces the original array within roundoff error.
            !
            !  the parameters inc, jump, n and lot are consistent if equality
            !  i1*inc + j1*jump = i2*inc + j2*jump for i1,i2 < n and j1,j2 < lot
            !  implies i1=i2 and j1=j2.  for multiple ffts to execute correctly,
            !  input variables inc, jump, n and lot must be consistent, otherwise
            !  at least one array element mistakenly is transformed more than once.
            !
            !  parameters:
            !
            !  integer lot, the number of sequences to be transformed
            !  within array c.
            !
            !  integer jump, the increment between the locations, in
            !  array c, of the first elements of two consecutive sequences to be
            !  transformed.
            !
            !  integer n, the length of each sequence to be
            !  transformed.  the transform is most efficient when n is a product of
            !  small primes.
            !
            !  integer inc, the increment between the locations, in
            !  array c, of two consecutive elements within the same sequence to be
            !  transformed.
            !
            !  input/output, complex (wp) c(lenc), an array containing lot
            !  sequences, each having length n, to be transformed.  c can have any
            !  number of dimensions, but the total number of locations must be at least
            !  lenc.  on output, c contains the transformed sequences.
            !
            !  integer lenc, the dimension of the c array.
            !  lenc must be at least (lot-1)*jump + inc*(n-1) + 1.
            !
            !  input, real (wp) wsave(lensav).  wsave's contents must be
            !  initialized with a call to cfftmi before the first call to routine cfftmf
            !  or cfftmb for a given transform length n.
            !
            !  integer lensav, the dimension of the wsave array.
            !  lensav must be at least 2*n + int(log(real(n))) + 4.
            !
            !  workspace, real (wp) work(lenwrk).
            !
            !  integer lenwrk, the dimension of the work array.
            !  lenwrk must be at least 2*lot*n.
            !
            !  integer ierror, error_flag.
            !  0, successful exit
            !  1, input parameter lenc not big enough;
            !  2, input parameter lensav not big enough;
            !  3, input parameter lenwrk not big enough;
            !  4, input parameters inc, jump, n, lot are not consistent.
            !
            !--------------------------------------------------
            ! Dummy arguments
            !--------------------------------------------------
            integer (ip) lot
            integer (ip) jump
            integer (ip) n
            integer (ip) inc
            real (wp) c(2,lenc)
            integer (ip) lenc
            real (wp) wsave(lensav)
            integer (ip) lensav
            real (wp) work(lenwrk)
            integer (ip) lenwrk
            integer (ip) ierror
            !--------------------------------------------------
        end subroutine cfftmb
        !
        !==> multiple complex forward
        !
        module subroutine cfftmf(lot, jump, n, inc, c, lenc, wsave, lensav, work, lenwrk, ierror)
            !
            ! cfftmf: complex forward fft, 1d, multiple vectors.
            !
            !  Purpose:
            !
            !  cfftmf computes the one-dimensional fourier transform of multiple
            !  periodic sequences within a complex array. this transform is referred
            !  to as the forward transform or fourier analysis, transforming the
            !  sequences from physical to spectral space. this transform is
            !  normalized since a call to cfftmf followed by a call to cfftmb
            !  (or vice-versa) reproduces the original array within roundoff error.
            !
            !  the parameters integers inc, jump, n and lot are consistent if equality
            !  i1*inc + j1*jump = i2*inc + j2*jump for i1,i2 < n and j1,j2 < lot
            !  implies i1=i2 and j1=j2. for multiple ffts to execute correctly,
            !  input variables inc, jump, n and lot must be consistent, otherwise
            !  at least one array element mistakenly is transformed more than once.
            !
            !  parameters:
            !
            !  integer lot, the number of sequences to be
            !  transformed within array c.
            !
            !  integer jump, the increment between the locations,
            !  in array c, of the first elements of two consecutive sequences to be
            !  transformed.
            !
            !  integer n, the length of each sequence to be
            !  transformed.  the transform is most efficient when n is a product of
            !  small primes.
            !
            !  integer inc, the increment between the locations, in
            !  array c, of two consecutive elements within the same sequence to be
            !  transformed.
            !
            !  input/output, complex (wp) c(lenc), array containing lot sequences,
            !  each having length n, to be transformed.  c can have any number of
            !  dimensions, but the total number of locations must be at least lenc.
            !
            !  integer lenc, the dimension of the c array.
            !  lenc must be at least (lot-1)*jump + inc*(n-1) + 1.
            !
            !  input, real (wp) wsave(lensav).  wsave's contents must be
            !  initialized with a call to cfftmi before the first call to routine cfftmf
            !  or cfftmb for a given transform length n.
            !
            !  integer lensav, the dimension of the wsave array.
            !  lensav must be at least 2*n + int(log(real(n))) + 4.
            !
            !  workspace, real (wp) work(lenwrk).
            !
            !  integer lenwrk, the dimension of the work array.
            !  lenwrk must be at least 2*lot*n.
            !
            !  integer ierror, error_flag.
            !  0 successful exit;
            !  1 input parameter lenc not big enough;
            !  2 input parameter lensav not big enough;
            !  3 input parameter lenwrk not big enough;
            !  4 input parameters inc, jump, n, lot are not consistent.
            !
            !--------------------------------------------------
            ! Dummy arguments
            !--------------------------------------------------
            integer (ip) lot
            integer (ip) jump
            integer (ip) n
            integer (ip) inc
            real(wp) c(2,lenc)
            integer (ip) lenc
            real (wp) wsave(lensav)
            integer (ip) lensav
            real (wp) work(lenwrk)
            integer (ip) lenwrk
            integer (ip) ierror
            !--------------------------------------------------
        end subroutine cfftmf
    end interface

    !---------------------------------------------------------------------------------
    ! Variables confined to the module
    !---------------------------------------------------------------------------------
    real (wp), parameter :: ZERO = 0.0_wp
    real (wp), parameter :: HALF = 0.5_wp
    real (wp), parameter :: ONE = 1.0_wp
    real (wp), parameter :: TWO = 2.0_wp
    real (wp), parameter :: THREE = 3.0_wp
    real (wp), parameter :: FOUR = 4.0_wp
    real (wp), parameter :: FIVE = 5.0_wp
    !---------------------------------------------------------------------------------

end module complex_transform_routines


CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCC
C
C   FFTPACK 5.0
C
C   Authors:  Paul N. Swarztrauber and Richard A. Valent
C
C   $Id: cosqmi.f,v 1.1 2010/02/23 01:10:43 rtoy Exp $
C
CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCC

      SUBROUTINE COSQMI (N, WSAVE, LENSAV, IER)
      INTEGER    N, LENSAV, IER
      REAL       WSAVE(LENSAV)
C
      IER = 0
C
      IF (LENSAV .LT. 2*N + INT(LOG(REAL(N))) +4) THEN
        IER = 2
        CALL XERFFT ('COSQMI', 3)
        GO TO 300
      ENDIF
C
      PIH = 2.*ATAN(1.)
      DT = PIH/FLOAT(N)
      FK = 0.
      DO 101 K=1,N
         FK = FK+1.
         WSAVE(K) = COS(FK*DT)
  101 CONTINUE
      LNSV = N + INT(LOG(REAL(N))) +4
      CALL RFFTMI (N, WSAVE(N+1), LNSV, IER1)
      IF (IER1 .NE. 0) THEN
        IER = 20
        CALL XERFFT ('COSQMI',-5)
      ENDIF
  300 CONTINUE
      RETURN
      END


	!Triangular lattice 2D
    subroutine triangular_nearest( p, p_i_pl, p_i_mi, p_j_pl, p_j_mi, p_i_pl_j_mi, p_i_mi_j_pl, L )
        implicit none
        integer :: p, p_i_pl, p_i_mi, p_j_pl, p_j_mi, p_i_pl_j_mi, p_i_mi_j_pl, L
        integer :: i, j, i_pl, j_pl, i_mi, j_mi

        i = int((p-1)/L)
        j = mod(p-1,L)
        i_pl = mod(i+1,L)
        j_pl = mod(j+1,L)
        if (i==0) then
            i_mi = L-1
        else
            i_mi = i-1
        end if
        if (j==0) then
            j_mi = L-1
        else
            j_mi = j-1
        end if

        p_i_pl = i_pl*L+j+1
        p_i_mi = i_mi*L+j+1
        p_j_pl = i*L+j_pl+1
        p_j_mi = i*L+j_mi+1
        p_i_pl_j_mi = i_pl*L+j_mi+1
        p_i_mi_j_pl = i_mi*L+j_pl+1

    end subroutine Triangular_nearest

	!cubic lattice
    subroutine cubic_nearest( p, p_i_pl, p_i_mi, p_j_pl, p_j_mi, p_k_pl, p_k_mi, L )
        implicit none
        integer :: p, p_i_pl, p_i_mi, p_j_pl, p_j_mi, p_k_pl, p_k_mi, L
        integer :: i, j, k, i_pl, j_pl, k_pl, i_mi, j_mi, k_mi

        i = (p-1)/L/L
        k = mod(p-1,L)
		j = (p-1-i*L*L)/L
        i_pl = mod(i+1,L)
        j_pl = mod(j+1,L)
		k_pl = mod(k+1,L)
        if (i==0) then
            i_mi = L-1
        else
            i_mi = i-1
        end if
        if (j==0) then
            j_mi = L-1
        else
            j_mi = j-1
        end if
        if (k==0) then
            k_mi = L-1
        else
            k_mi = k-1
        end if

        p_i_pl = i_pl*L*L+j*L+k+1
        p_i_mi = i_mi*L*L+j*L+k+1
        p_j_pl = i*L*L+j_pl*L+k+1
        p_j_mi = i*L*L+j_mi*L+k+1
        p_k_pl = i*L*L+j*L+k_pl+1
        p_k_mi = i*L*L+j*L+k_mi+1

    end subroutine cubic_nearest