!----------------------------------------------------------------------------
! This file is part of UMDLES.
!
! UMDLES is free software; you can redistribute it and/or modify
! it under the terms of the GNU General Public License as published by
! the Free Software Foundation; either version 3 of the License, or
! (at your option) any later version.
!
! UMDLES is distributsed in the hope that it will be useful,
! but WITHOUT ANY WARRANTY; without even the implied warranty of
! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
! GNU General Public License for more details.
!
! You should have received a copy of the GNU General Public License
! along with this program.  If not, see <http://www.gnu.org/licenses/>.
!
! Copyright 2015-2017, David A. New
!----------------------------------------------------------------------------

module modstat_io
  implicit none

  integer, parameter :: ntypes = 11

  character(*), parameter :: time_name   = 'time'
  character(*), parameter :: rank_name   = 'rank'
  character(*), parameter :: zt_name     = 'zt'
  character(*), parameter :: zm_name     = 'zm'
  character(*), parameter :: x_name      = 'x'
  character(*), parameter :: y_name      = 'y'
  character(*), parameter :: type_name   = 'type'
  character(*), parameter :: w_name      = 'w_bin_dim'
  character(*), parameter :: t_name      = 't_bin_dim'
  character(*), parameter :: b_name      = 'b_bin_dim'
  character(*), parameter :: q_name      = 'q_bin_dim'
  character(*), parameter :: z_hist_name = 'z' 

  interface create_stat_var
     module procedure create_stat_var_0d
     module procedure create_stat_var_1d
  end interface create_stat_var

  interface write_stat_var
     module procedure write_stat_real_var_0d
     module procedure write_stat_real_var_1d
     module procedure write_stat_real_var_2d
     module procedure write_stat_int_var_1d
  end interface

  interface write_hist_var
     module procedure write_hist_var2
     module procedure write_hist_var3
  end interface write_hist_var

contains

  !
  ! create_stat_file
  !
  subroutine create_stat_file(name, ncid, irec, cntlat)
    use netcdf_interface, only : create_file, create_dim, create_unlimited_dim, create_var, &
                                 my_nf90_real, my_nf90_int, write_var
    use mpi_interface, only    : pecount 
    use grid, only             : zt, zm, dn0, dthldtls, dqtdtls, level, wfls, pi0, pi1, u0, v0, nzp
    use defs, only             : cp, R, alvl
    implicit none

    character(*), intent(in) :: name
    integer, intent(out)     :: ncid, irec
    real, intent(in)         :: cntlat

    integer              :: i
    integer, allocatable :: rank_values(:), type_values(:)
    real, allocatable    :: exner(:)
    real                 :: cp_copy, alvl_copy, R_copy, cntlat_copy

    ! Create file
    call create_file(name, ncid)

    ! Create time dimension (unlimited)
    call create_unlimited_dim(ncid, time_name, my_nf90_real, 's')

    ! Create rank dimension if using MPI
    if (pecount > 1) then
       allocate(rank_values(pecount))
       do i = 1,pecount
          rank_values(i) = i
       end do
       call create_dim(ncid, rank_name, '1', rank_values)
    end if

    ! Define altitude dimensions
    call create_dim(ncid, zt_name, 'm', zt)
    call create_dim(ncid, zm_name, 'm', zm)
    
    !
    allocate(type_values(ntypes))
    do i = 1,ntypes
       type_values(i) = i
    end do
    call create_dim(ncid, type_name, '1', type_values)

    ! Create non-time-dependant parameters for simulation parameters
    call create_var(ncid, 'cp',   my_nf90_real, 'JK-1kg-1')
    call create_var(ncid, 'alvl', my_nf90_real, 'Jkg-1')
    call create_var(ncid, 'R',    my_nf90_real, 'J+1K-1kg-1')
    call create_var(ncid, 'cntlat', my_nf90_real, 'degrees')

    call create_var(ncid, 'dn0',  my_nf90_real, 'kgm-3', (/zt_name/))
    call create_var(ncid, 'pi0',  my_nf90_real, '1', (/zt_name/))

    ! Create variable for first and last sample times
    call create_var(ncid, 'nsmp',  my_nf90_int, '1', (/time_name/))
    call create_var(ncid, 'fsttm', my_nf90_real, 's', (/time_name/))
    call create_var(ncid, 'lsttm', my_nf90_real, 's', (/time_name/))

    ! Needed for intent(inout) with parameter
    cp_copy     = cp
    alvl_copy   = alvl
    R_copy      = R
    cntlat_copy = cntlat
    allocate(exner(nzp))
    exner(:) = pi0(:)/cp

    ! Write parameters
    call write_var(ncid, 'cp',   cp_copy)
    call write_var(ncid, 'alvl', alvl_copy)
    call write_var(ncid, 'R',    R_copy)
    call write_var(ncid, 'cntlat', cntlat_copy)

    call write_var(ncid, 'dn0', dn0)
    call write_var(ncid, 'pi0', exner)

    ! Initialize record index
    irec = 0

  end subroutine create_stat_file

  !
  ! create_hist_file
  !
  subroutine create_hist_file(name, nbins, nz_hist, k_hist, ncid, irec)
    use grid, only : level, zt
    use netcdf_interface, only : create_file, create_dim, create_unlimited_dim, create_var, &
                                 my_nf90_real, my_nf90_int, write_var
    use mpi_interface, only    : pecount 
    implicit none

    character(*), intent(in)    :: name
    integer, intent(out)        :: ncid, irec
    integer, intent(in)         :: nbins, nz_hist
    integer, dimension(nz_hist) :: k_hist

    integer                   :: i, ibin, iz_hist
    integer, allocatable      :: rank_values(:)
    integer, dimension(nbins) :: bins 
    real, dimension(nz_hist)  :: z_hist

    ! Create file
    call create_file(name, ncid)

    ! Create time dimension (unlimited)
    call create_unlimited_dim(ncid, time_name, my_nf90_real, 's')

    ! Create rank dimension if using MPI
    if ( pecount > 1 ) then
       allocate(rank_values(pecount))
       do i = 1,pecount
          rank_values(i) = i
       end do
       call create_dim(ncid, rank_name, '1', rank_values)
    end if

    ! Create height dimension
    do iz_hist = 1,nz_hist
       z_hist(iz_hist) = zt(k_hist(iz_hist)) 
    end do
    call create_dim(ncid, z_hist_name, 'm', z_hist)

    ! Create individual bin dimensions
    do ibin = 1,nbins
       bins(ibin) = ibin
    end do
    call create_dim(ncid, w_name, '1', bins)
    call create_dim(ncid, t_name, '1', bins)
    call create_dim(ncid, b_name, '1', bins)
    if ( level > 0 ) then
       call create_dim(ncid, q_name, '1', bins)
    end if

    ! Create variable for first and last sample times
    call create_var(ncid, 'nsmp',  my_nf90_int, '1', (/time_name/))
    call create_var(ncid, 'fsttm', my_nf90_real, 's', (/time_name/))
    call create_var(ncid, 'lsttm', my_nf90_real, 's', (/time_name/))

    ! Initialize record index
    irec = 0

  end subroutine create_hist_file

  !
  ! create_analysis_file
  !
  subroutine create_analysis_file(name, ncid, irec)
    use grid, only : level, zt, xm, ym, nxp, nyp
    use netcdf_interface, only : create_file, create_dim, create_unlimited_dim, create_var, &
                                 my_nf90_real, my_nf90_int, write_var
    use mpi_interface, only    : pecount, nxprocs, nyprocs
    implicit none

    character(*), intent(in) :: name
    integer, intent(out)     :: ncid, irec

    integer :: i
    integer, allocatable, dimension(:) :: rank_values
    real, allocatable, dimension(:) :: xm_global, ym_global

    allocate ( xm_global(nxprocs*( nxp - 4 )), ym_global(nyprocs*( nyp - 4 )) )
    

    ! Create file
    call create_file(name, ncid)

    ! Create time dimension (unlimited)
    call create_unlimited_dim(ncid, time_name, my_nf90_real, 's')

    ! Create rank dimension if using MPI
    if ( pecount > 1 ) then
       allocate(rank_values(pecount))
       do i = 1,pecount
          rank_values(i) = i
       end do
       call create_dim(ncid, rank_name, '1', rank_values)
    end if

    ! Create spatial dimensions
    call create_dim(ncid, x_name, '1',  xm_global)
    call create_dim(ncid, y_name, '1',  ym_global)
    call create_dim(ncid, zt_name, '1', zt)

    ! Initialize record index
    irec = 0

  end subroutine create_analysis_file

  !
  ! create_analysis_var
  !
  subroutine create_analysis_var(ncid, var_name, data_type, units)
    use netcdf_interface, only : create_var, create_dim
    use mpi_interface, only    : pecount
    implicit none

    integer, intent(in)      :: ncid, data_type
    character(*), intent(in) :: var_name, units

    call create_var(ncid, var_name, data_type, units, (/ zt_name, x_name, y_name, time_name /))

  end subroutine create_analysis_var

  !
  ! write_analysis_var
  !
  subroutine write_analysis_var(ncid, name, var, irec)
    use mpi_interface, only    : myid, wrxid, wryid, xoffset, yoffset, nxprocs, nyprocs
    use grid, only             : nzp, nxp, nyp
    use netcdf_interface, only : write_var, get_var_ndims
    implicit none
    
    integer, intent(in)      :: ncid, irec
    character(*), intent(in) :: name
    real, intent(in)         :: var(:,:,:)

    real, allocatable :: var_copy(:,:,:)

    ! Because Parallel NetCDF does not allow intent(in) variables
    allocate( var_copy(nzp,nxp-4,nyp-4) ) 
    var_copy = var(:,3:nxp-2,3:nyp-2)

    call write_var(ncid, name, var_copy(:,:,:), irec, start_in=(/ 1, xoffset(wrxid)+1, yoffset(wryid)+1 /), stride_in=(/ nzp, nxp-4, nyp-4 /))

  end subroutine write_analysis_var

  !
  ! create_stat_var_0d
  !
  subroutine create_stat_var_0d(ncid, var_name, data_type, units, multicore_flag_in)
    use netcdf_interface, only : create_var, create_dim
    use mpi_interface, only    : pecount
    implicit none

    integer, intent(in)           :: ncid, data_type
    character(*), intent(in)      :: var_name, units
    logical, intent(in), optional :: multicore_flag_in

    logical :: multicore_flag

    if ( present(multicore_flag_in) ) then
       multicore_flag = multicore_flag_in
    else
       multicore_flag = .false.
    end if

    if ( pecount == 1 .or. .not. multicore_flag ) then
       call create_var(ncid, var_name, data_type, units, (/time_name/))
    else
       call create_var(ncid, var_name, data_type, units, (/rank_name, time_name/))
    end if

  end subroutine create_stat_var_0d

  !
  ! create_stat_var_1d
  !
  subroutine create_stat_var_1d(ncid, var_name, z_dim_name, data_type, units, multicore_flag_in)
    use netcdf_interface, only : create_var, create_dim
    use mpi_interface, only    : pecount
    implicit none

    integer, intent(in)           :: ncid, data_type
    character(*), intent(in)      :: var_name, z_dim_name, units
    logical, intent(in), optional :: multicore_flag_in

    logical :: multicore_flag

    if ( present(multicore_flag_in) ) then
       multicore_flag = multicore_flag_in
    else
       multicore_flag = .false.
    end if

    if ( pecount == 1 .or. .not. multicore_flag ) then
       call create_var(ncid, var_name, data_type, units, (/z_dim_name, time_name/))
    else
       call create_var(ncid, var_name, data_type, units, (/z_dim_name, rank_name, time_name/))
    end if

  end subroutine create_stat_var_1d

  !
  ! create_stat_var_2d
  !
  subroutine create_stat_var_2d(ncid, var_name, z_dim_name, data_type, units, multicore_flag_in)
    use netcdf_interface, only : create_var, create_dim
    use mpi_interface, only    : pecount
    implicit none

    integer, intent(in)           :: ncid, data_type
    character(*), intent(in)      :: var_name, z_dim_name, units
    logical, intent(in), optional :: multicore_flag_in

    logical :: multicore_flag

    if ( present(multicore_flag_in) ) then
       multicore_flag = multicore_flag_in
    else
       multicore_flag = .false.
    end if

    if ( pecount == 1 .or. .not. multicore_flag ) then
       call create_var(ncid, var_name, data_type, units, (/ z_dim_name, type_name, time_name /))
    else
       call create_var(ncid, var_name, data_type, units, (/ z_dim_name, type_name, rank_name, time_name /))
    end if

  end subroutine create_stat_var_2d

  !
  ! create_hist_bins
  !
  subroutine create_hist_bins(ncid, bins_name, bins_dim_name, data_type, units)
    use netcdf_interface, only : create_var
    implicit none

    integer, intent(in)      :: ncid, data_type
    character(*), intent(in) :: bins_name, bins_dim_name, units

    call create_var(ncid, bins_name, data_type, units, (/ bins_dim_name, z_hist_name /))

  end subroutine create_hist_bins

  !
  ! create_hist_var_2d
  !
  subroutine create_hist_var_2d(ncid, var_name, data_type, multicore_flag_in)
    use netcdf_interface, only : create_var
    use mpi_interface, only    : pecount
    use grid, only             : level
    implicit none

    integer, intent(in)           :: ncid, data_type
    character(*), intent(in)      :: var_name
    logical, intent(in), optional :: multicore_flag_in

    logical :: multicore_flag

    if ( present(multicore_flag_in) ) then
       multicore_flag = multicore_flag_in
    else
       multicore_flag = .false.
    end if

    if ( pecount == 1 .or. .not. multicore_flag ) then
       call create_var(ncid, var_name, data_type, '1', (/ w_name, b_name, z_hist_name, time_name /))
    else
       call create_var(ncid, var_name, data_type, '1', (/ w_name, b_name, z_hist_name, rank_name, time_name /))
    end if

  end subroutine create_hist_var_2d

  !
  ! create_hist_var_3d
  !
  subroutine create_hist_var_3d(ncid, var_name, data_type, multicore_flag_in)
    use netcdf_interface, only : create_var
    use mpi_interface, only    : pecount
    use grid, only             : level
    implicit none

    integer, intent(in)           :: ncid, data_type
    character(*), intent(in)      :: var_name
    logical, intent(in), optional :: multicore_flag_in

    logical :: multicore_flag

    if ( present(multicore_flag_in) ) then
       multicore_flag = multicore_flag_in
    else
       multicore_flag = .false.
    end if

    if ( pecount == 1 .or. .not. multicore_flag ) then
       if ( level == 0 ) then
          call create_var(ncid, var_name, data_type, '1', (/ w_name, t_name,         z_hist_name, time_name /))
       else
          call create_var(ncid, var_name, data_type, '1', (/ w_name, t_name, q_name, z_hist_name, time_name /))
       end if
    else
       if ( level == 0 ) then
          call create_var(ncid, var_name, data_type, '1', (/ w_name, t_name,         z_hist_name, rank_name, time_name /))
       else
          call create_var(ncid, var_name, data_type, '1', (/ w_name, t_name, q_name, z_hist_name, rank_name, time_name /))
       end if
    end if

  end subroutine create_hist_var_3d

  !
  ! write_hist_var2
  !
  subroutine write_hist_var2(ncid, name, nbins, nz_hist, var, irec)
    use mpi_interface, only    : pecount, myid
    use grid, only             : nzp, level
    use netcdf_interface, only : write_var, get_var_ndims
    implicit none
    
    integer, intent(in)      :: ncid, irec, nbins, nz_hist
    character(*), intent(in) :: name
    real, intent(in)         :: var(:,:,:)

    integer :: ndims
    real, allocatable :: var_copy(:,:,:)

    ! Because Parallel NetCDF does not allow intent(in) variables
    allocate(var_copy(nbins,nbins,nz_hist))
    var_copy = var

    if ( pecount == 1 ) then
       call write_var(ncid, name, var_copy, irec)
    else
       call get_var_ndims(ncid, name, ndims)

       if ( ndims == 4 ) then
          call write_var(ncid, name, var_copy, irec)
       else
          call write_var(ncid, name, var_copy(:,:,:), irec, start_in=(/1, 1, 1, myid+1/), stride_in=(/nbins, nbins, nz_hist, 1/))
       end if
    end if

  end subroutine write_hist_var2

  !
  ! write_hist_var3
  !
  subroutine write_hist_var3(ncid, name, nbins, nz_hist, var, irec)
    use mpi_interface, only    : pecount, myid
    use grid, only             : nzp, level
    use netcdf_interface, only : write_var, get_var_ndims
    implicit none
    
    integer, intent(in)      :: ncid, irec, nbins, nz_hist
    character(*), intent(in) :: name
    real, intent(in)         :: var(:,:,:,:)

    integer :: ndims
    real, allocatable :: var_copy(:,:,:,:)

    ! Because Parallel NetCDF does not allow intent(in) variables
    if ( level == 0 ) then
       allocate(var_copy(nbins,nbins,1,nz_hist))
    else
       allocate(var_copy(nbins,nbins,nbins,nz_hist))
    end if
    var_copy = var

    if ( pecount == 1 ) then
       call write_var(ncid, name, var_copy, irec)
    else
       call get_var_ndims(ncid, name, ndims)

       if ( level == 0 ) then
          if ( ndims == 4 ) then
             call write_var(ncid, name, var_copy(:,:,1,:), irec)
          else
             call write_var(ncid, name, var_copy(:,:,1,:), irec, start_in=(/1, 1, 1, myid+1/), stride_in=(/nbins, nbins, nz_hist, 1/))
          end if
       else
          if ( ndims == 5 ) then
             call write_var(ncid, name, var_copy, irec)
          else
             call write_var(ncid, name, var_copy, irec, start_in=(/1, 1, 1, 1, myid+1/), stride_in=(/nbins, nbins, nbins, nz_hist, 1/))
          end if
       end if
    end if

  end subroutine write_hist_var3

  !
  ! write_hist_bins
  !
  subroutine write_hist_bins(ncid, name, nbins, nz_hist, bins, irec)
    use grid, only             : level
    use netcdf_interface, only : write_var
    implicit none
    
    integer, intent(in)      :: ncid, irec, nbins, nz_hist
    character(*), intent(in) :: name
    real, intent(in)         :: bins(nbins,nz_hist)

    integer :: ndims
    real, dimension(nbins,nz_hist) :: bins_copy

    ! Because Parallel NetCDF does not allow intent(in) binsiables
    bins_copy = bins

    call write_var(ncid, name, bins_copy, irec)

  end subroutine write_hist_bins

  !
  ! advance_stat_time
  !
  subroutine advance_stat_time(time, fsttm, nsmp, ncid, irec)
    use netcdf_interface, only : write_var
    use grid, only             : pi1, nzp
    use defs, only             : cp
    implicit none

    real, intent(in)     :: time, fsttm
    integer, intent(in)  :: nsmp, ncid
    integer, intent(out) :: irec

    real    :: time_copy, fsttm_copy, stat_time
    integer :: nsmp_copy

    ! Because Parallel NetCDF does not allow intent(in) variables
    time_copy  = time
    nsmp_copy  = nsmp
    fsttm_copy = fsttm

    !
    stat_time = 0.5*( fsttm + time )

    ! Increment record counter
    irec = irec + 1

    ! Write to disk
    call write_var(ncid, time_name, stat_time, irec)
    call write_var(ncid, 'nsmp',  nsmp_copy,   irec)
    call write_var(ncid, 'fsttm', fsttm_copy,  irec)
    call write_var(ncid, 'lsttm', time_copy,   irec)

  end subroutine advance_stat_time

  !
  ! advance_analysis_time
  !
  subroutine advance_analysis_time(time, ncid, irec)
    use netcdf_interface, only : write_var
    implicit none

    real, intent(in)     :: time
    integer, intent(in)  :: ncid
    integer, intent(out) :: irec

    real :: time_copy 

    time_copy = time

    ! Increment record counter
    irec = irec + 1

    ! Write to disk
    call write_var(ncid, time_name, time_copy, irec)

  end subroutine advance_analysis_time

  !
  ! write_stat_real_var_0d
  !
  subroutine write_stat_real_var_0d(ncid, name, var, irec)
    use mpi_interface, only    : pecount, myid
    use grid, only             : nzp
    use netcdf_interface, only : write_var, get_var_ndims
    implicit none
    
    integer, intent(in)      :: ncid, irec
    character(*), intent(in) :: name
    real, intent(in)         :: var

    integer :: ndims
    real :: var_copy(1)

    ! Because Parallel NetCDF does not allow intent(in) variables
    var_copy(1) = var

    if ( pecount == 1 ) then
       call write_var(ncid, name, var_copy, irec)
    else
       call get_var_ndims(ncid, name, ndims)

       if ( ndims == 1 ) then
          call write_var(ncid, name, var_copy, irec)
       else
          call write_var(ncid, name, var_copy, irec, start_in=(/myid+1/), stride_in=(/1/))
       end if
    end if

  end subroutine write_stat_real_var_0d

  !
  ! write_stat_real_var_1d
  !
  subroutine write_stat_real_var_1d(ncid, name, var, irec)
    use mpi_interface, only    : pecount, myid
    use grid, only             : nzp
    use netcdf_interface, only : write_var, get_var_ndims
    implicit none
    
    integer, intent(in)      :: ncid, irec
    character(*), intent(in) :: name
    real, intent(in)         :: var(:)

    integer :: ndims
    real, allocatable :: var_copy(:)

    ! Because Parallel NetCDF does not allow intent(in) variables
    allocate( var_copy(nzp) )
    var_copy(:) = var(:)

    if ( pecount == 1 ) then
       call write_var(ncid, name, var_copy, irec)
    else
       call get_var_ndims(ncid, name, ndims)

       if ( ndims == 2 ) then
          call write_var(ncid, name, var_copy, irec)
       else
          call write_var(ncid, name, var_copy, irec, start_in=(/1, myid+1/), stride_in=(/nzp, 1/))
       end if
    end if

  end subroutine write_stat_real_var_1d

  !
  ! write_stat_real_var_2d
  !
  subroutine write_stat_real_var_2d(ncid, name, var, irec)
    use mpi_interface, only    : pecount, myid
    use grid, only             : nzp
    use netcdf_interface, only : write_var, get_var_ndims
    implicit none
    
    integer, intent(in)      :: ncid, irec
    character(*), intent(in) :: name
    real, intent(in)         :: var(:,:)

    integer :: ndims
    real, allocatable :: var_copy(:,:)

    ! Because Parallel NetCDF does not allow intent(in) variables
    allocate(var_copy(nzp,ntypes))
    var_copy(:,:) = var(:,:)

    if ( pecount == 1 ) then
       call write_var(ncid, name, var_copy, irec)
    else
       call get_var_ndims(ncid, name, ndims)
       
       if ( ndims == 3 ) then
          call write_var(ncid, name, var_copy, irec)
       else
          call write_var(ncid, name, var_copy, irec, start_in=(/1, 1, myid+1/), stride_in=(/nzp, ntypes, 1/))
       end if
    end if

  end subroutine write_stat_real_var_2d

  !
  ! write_stat_int_var_1d
  !
  subroutine write_stat_int_var_1d(ncid, name, var, irec)
    use mpi_interface, only    : pecount, myid
    use grid, only             : nzp
    use netcdf_interface, only : write_var, get_var_ndims
    implicit none
    
    integer, intent(in)      :: ncid, irec
    character(*), intent(in) :: name
    integer, intent(in)      :: var(:)

    integer :: ndims
    integer, allocatable :: var_copy(:)

    ! Because Parallel NetCDF does not allow intent(in) variables
    allocate(var_copy(nzp))
    var_copy(:) = var(:)

    if ( pecount == 1 ) then
       call write_var(ncid, name, var_copy, irec)
    else
       call get_var_ndims(ncid, name, ndims)

       if ( ndims == 2 ) then
          call write_var(ncid, name, var_copy, irec)
       else
          call write_var(ncid, name, var_copy, irec, start_in=(/1, myid+1/), stride_in=(/nzp, 1/))
       end if
    end if

  end subroutine write_stat_int_var_1d
  
end module modstat_io




