subroutine wstats(ngrid,nom,titre,unite,dim,px)

use mod_phys_lmdz_para, only : is_mpi_root, is_master, gather, klon_mpi_begin
use mod_grid_phy_lmdz, only : klon_glo, Grid1Dto2D_glo

implicit none

#include "dimensions.h"
#include "dimphys.h"
#include "comconst.h"
#include "statto.h"
#include "netcdf.inc"

integer,intent(in) :: ngrid
character (len=*),intent(in) :: nom,titre,unite
integer,intent(in) :: dim
integer,parameter :: iip1=iim+1
integer,parameter :: jjp1=jjm+1
real,intent(in) :: px(ngrid,llm)
real, dimension(iip1,jjp1,llm) :: mean3d,sd3d,dx3
real, dimension(iip1,jjp1) :: mean2d,sd2d,dx2
character (len=50) :: namebis
character (len=50), save :: firstvar
integer :: ierr,varid,nbdim,nid
integer :: meanid,sdid
integer, dimension(4)  :: id,start,sizes
logical, save :: firstcall=.TRUE.
integer :: l,i,j,ig0
integer,save :: indx

integer, save :: step=0

! Added to work in parallel mode
#ifdef CPP_PARA
real px3_glop(klon_glo,llm) ! to store a 3D data set on global physics grid
real px3_glo(iim,jjp1,llm) ! to store a global 3D data set on global lonxlat grid
real px2_glop(klon_glo) ! to store a 2D data set on global physics grid
real px2_glo(iim,jjp1) ! to store a 2D data set on global lonxlat grid
real px2(ngrid)
real px3(ngrid,llm)
#else
! When not running in parallel mode:
real px3_glop(ngrid,llm) ! to store a 3D data set on global physics grid
real px3_glo(iim,jjp1,llm) ! to store a global 3D data set on global lonxlat grid
real px2_glop(ngrid) ! to store a 2D data set on global physics grid
real px2_glo(iim,jjp1) ! to store a 2D data set on global lonxlat grid
#endif

! 1. Initialization (creation of stats.nc file)
if (firstcall) then
   firstcall=.false.
   firstvar=trim((nom))
   call inistats(ierr)
endif

if (firstvar==nom) then ! If we're back to the first variable, increment time counter
      step = step + 1
endif

if (mod(step,istats).ne.0) then
  ! if its not time to write to file, exit
   RETURN
endif

! collect fields on a global physics grid
#ifdef CPP_PARA
 if (dim.eq.3) then
  px3(1:ngrid,1:llm)=px(1:ngrid,1:llm)
  ! Gather fieds on a "global" (without redundant longitude) array
  call Gather(px3,px3_glop)
!$OMP MASTER
  if (is_mpi_root) then
    call Grid1Dto2D_glo(px3_glop,px3_glo)
    ! copy dx3_glo() to dx3(:) and add redundant longitude
    dx3(1:iim,:,:)=px3_glo(1:iim,:,:)
    dx3(iip1,:,:)=dx3(1,:,:)
  endif
!$OMP END MASTER
!$OMP BARRIER
 else ! dim.eq.2
  ! Gather fieds on a "global" (without redundant longitude) array
  px2(:)=px(:,1)
  call Gather(px2,px2_glop)
!$OMP MASTER
          if (is_mpi_root) then
            call Grid1Dto2D_glo(px2_glop,px2_glo)
            ! copy px2_glo() to dx2(:) and add redundant longitude
            dx2(1:iim,:)=px2_glo(1:iim,:)
            dx2(iip1,:)=dx2(1,:)
          endif
!$OMP END MASTER
!$OMP BARRIER
 endif
#else
  if (dim.eq.3) then
    px3_glop(:,1:llm)=px(:,1:llm)
!  Passage variable physique -->  variable dynamique
    DO l=1,llm
      DO i=1,iim
         px3_glo(i,1,l)=px(1,l)
         px3_glo(i,jjp1,l)=px(ngrid,l)
      ENDDO
      DO j=2,jjm
         ig0= 1+(j-2)*iim
         DO i=1,iim
            px3_glo(i,j,l)=px(ig0+i,l)
         ENDDO
      ENDDO
    ENDDO
  else ! dim.eq.2
    px2_glop(:)=px(:,1)
!    Passage variable physique -->  physique dynamique
   DO i=1,iim
     px2_glo(i,1)=px(1,1)
     px2_glo(i,jjp1)=px(ngrid,1)
   ENDDO
   DO j=2,jjm
     ig0= 1+(j-2)*iim
     DO i=1,iim
        px2_glo(i,j)=px(ig0+i,1)
     ENDDO
   ENDDO
  endif
#endif

! 2. Write field to file

if (is_master) then
! only master needs do this

ierr = NF_OPEN("stats.nc",NF_WRITE,nid)

namebis=trim(nom)

! test: check if that variable already exists in file
ierr= NF_INQ_VARID(nid,namebis,meanid)

if (ierr.ne.NF_NOERR) then
  ! variable not in file, create/define it
   if (firstvar==nom) then 
      indx=1
      count(:)=0
   endif


!declaration de la variable

! choix du nom des coordonnees
   ierr= NF_INQ_DIMID(nid,"longitude",id(1))
   ierr= NF_INQ_DIMID(nid,"latitude",id(2))
   if (dim.eq.3) then
      ierr= NF_INQ_DIMID(nid,"altitude",id(3))
      ierr= NF_INQ_DIMID(nid,"Time",id(4))
      nbdim=4
   else if (dim==2) then
      ierr= NF_INQ_DIMID(nid,"Time",id(3))
      nbdim=3
   endif

   write (*,*) "====================="
   write (*,*) "STATS: creation de ",nom
   namebis=trim(nom)
   call def_var(nid,namebis,titre,unite,nbdim,id,meanid,ierr)
   if (dim.eq.3) then
     call inivar(nid,meanid,size(px3_glop,1),dim,indx,px3_glop,ierr)
   else ! dim.eq.2
     call inivar(nid,meanid,size(px2_glop,1),dim,indx,px2_glop,ierr)
   endif
   namebis=trim(nom)//"_sd"
   call def_var(nid,namebis,trim(titre)//" total standard deviation over the season",unite,nbdim,id,sdid,ierr)
   if (dim.eq.3) then
     call inivar(nid,sdid,size(px3_glop,1),dim,indx,px3_glop,ierr)
   else ! dim.eq.2
     call inivar(nid,sdid,size(px2_glop,1),dim,indx,px2_glop,ierr)
   endif

   ierr= NF_CLOSE(nid)
   return

else
   ! variable found in file
   namebis=trim(nom)//"_sd"
   ierr= NF_INQ_VARID(nid,namebis,sdid)

endif

if (firstvar==nom) then 
   count(indx)=count(int(indx))+1
   indx=indx+1
   if (indx>istime) then
      indx=1
   endif
endif

if (count(indx)==0) then
   ! very first time we write the variable to file
   if (dim.eq.3) then
      start=(/1,1,1,indx/)
      sizes=(/iip1,jjp1,llm,1/)
      mean3d(:,:,:)=0
      sd3d(:,:,:)=0
   else if (dim.eq.2) then
      start=(/1,1,indx,0/)
      sizes=(/iip1,jjp1,1,0/)
      mean2d(:,:)=0
      sd2d(:,:)=0
   endif
else
   ! load values from file
   if (dim.eq.3) then
      start=(/1,1,1,indx/)
      sizes=(/iip1,jjp1,llm,1/)
#ifdef NC_DOUBLE
      ierr = NF_GET_VARA_DOUBLE(nid,meanid,start,sizes,mean3d)
      ierr = NF_GET_VARA_DOUBLE(nid,sdid,start,sizes,sd3d)
#else
      ierr = NF_GET_VARA_REAL(nid,meanid,start,sizes,mean3d)
      ierr = NF_GET_VARA_REAL(nid,sdid,start,sizes,sd3d)
#endif
      if (ierr.ne.NF_NOERR) then
         write (*,*) NF_STRERROR(ierr)
         stop ""
      endif

   else if (dim.eq.2) then
      start=(/1,1,indx,0/)
      sizes=(/iip1,jjp1,1,0/)
#ifdef NC_DOUBLE
      ierr = NF_GET_VARA_DOUBLE(nid,meanid,start,sizes,mean2d)
      ierr = NF_GET_VARA_DOUBLE(nid,sdid,start,sizes,sd2d)
#else
      ierr = NF_GET_VARA_REAL(nid,meanid,start,sizes,mean2d)
      ierr = NF_GET_VARA_REAL(nid,sdid,start,sizes,sd2d)
#endif
      if (ierr.ne.NF_NOERR) then
         write (*,*) NF_STRERROR(ierr)
         stop ""
      endif
   endif
endif ! of if (count(indx)==0)


! 2.1. Build dx* (data on lon-lat grid, with redundant longitude)

if (dim.eq.3) then
  dx3(1:iim,:,:)=px3_glo(:,:,:)
  dx3(iip1,:,:)=dx3(1,:,:)
else ! dim.eq.2
  dx2(1:iim,:)=px2_glo(:,:)
  dx2(iip1,:)=dx2(1,:)
endif


! 2.2. Add current values to previously stored sums

if (dim.eq.3) then

   mean3d(:,:,:)=mean3d(:,:,:)+dx3(:,:,:)
   sd3d(:,:,:)=sd3d(:,:,:)+dx3(:,:,:)**2

#ifdef NC_DOUBLE
   ierr = NF_PUT_VARA_DOUBLE(nid,meanid,start,sizes,mean3d)
   ierr = NF_PUT_VARA_DOUBLE(nid,sdid,start,sizes,sd3d)
#else
   ierr = NF_PUT_VARA_REAL(nid,meanid,start,sizes,mean3d)
   ierr = NF_PUT_VARA_REAL(nid,sdid,start,sizes,sd3d)
#endif

else if (dim.eq.2) then

   mean2d(:,:)= mean2d(:,:)+dx2(:,:)
   sd2d(:,:)=sd2d(:,:)+dx2(:,:)**2

#ifdef NC_DOUBLE
   ierr = NF_PUT_VARA_DOUBLE(nid,meanid,start,sizes,mean2d)
   ierr = NF_PUT_VARA_DOUBLE(nid,sdid,start,sizes,sd2d)
#else
   ierr = NF_PUT_VARA_REAL(nid,meanid,start,sizes,mean2d)
   ierr = NF_PUT_VARA_REAL(nid,sdid,start,sizes,sd2d)
#endif

endif ! of if (dim.eq.3) elseif (dim.eq.2)

  ierr= NF_CLOSE(nid)
endif ! of if (is_master)

end subroutine wstats

!======================================================
subroutine inivar(nid,varid,ngrid,dim,indx,px,ierr)

implicit none

include "dimensions.h"
include "dimphys.h"
include "netcdf.inc"

integer, intent(in) :: nid,varid,dim,indx,ngrid
real, dimension(ngrid,llm), intent(in) :: px
integer, intent(out) :: ierr

integer,parameter :: iip1=iim+1
integer,parameter :: jjp1=jjm+1

integer :: l,i,j,ig0
integer, dimension(4) :: start,sizes
real, dimension(iip1,jjp1,llm) :: dx3
real, dimension(iip1,jjp1) :: dx2

if (dim.eq.3) then

   start=(/1,1,1,indx/)
   sizes=(/iip1,jjp1,llm,1/)

!  Passage variable physique -->  variable dynamique

   DO l=1,llm
      DO i=1,iip1
         dx3(i,1,l)=px(1,l)
         dx3(i,jjp1,l)=px(ngrid,l)
      ENDDO
      DO j=2,jjm
         ig0= 1+(j-2)*iim
         DO i=1,iim
            dx3(i,j,l)=px(ig0+i,l)
         ENDDO
         dx3(iip1,j,l)=dx3(1,j,l)
      ENDDO
   ENDDO

#ifdef NC_DOUBLE
   ierr = NF_PUT_VARA_DOUBLE(nid,varid,start,sizes,dx3)
#else
   ierr = NF_PUT_VARA_REAL(nid,varid,start,sizes,dx3)
#endif

else if (dim.eq.2) then

      start=(/1,1,indx,0/)
      sizes=(/iip1,jjp1,1,0/)

!    Passage variable physique -->  physique dynamique

  DO i=1,iip1
     dx2(i,1)=px(1,1)
     dx2(i,jjp1)=px(ngrid,1)
  ENDDO
  DO j=2,jjm
     ig0= 1+(j-2)*iim
     DO i=1,iim
        dx2(i,j)=px(ig0+i,1)
     ENDDO
     dx2(iip1,j)=dx2(1,j)
  ENDDO

#ifdef NC_DOUBLE
   ierr = NF_PUT_VARA_DOUBLE(nid,varid,start,sizes,dx2)
#else
   ierr = NF_PUT_VARA_REAL(nid,varid,start,sizes,dx2)
#endif

endif

end subroutine inivar

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

subroutine def_var_stats(nid,name,title,units,nbdim,dimids,nvarid,ierr)

! This subroutine defines variable 'name' in a (pre-existing and opened)
! NetCDF file (known from its NetCDF ID 'nid').
! The number of dimensions 'nbdim' of the variable, as well as the IDs of
! corresponding dimensions must be set (in array 'dimids').
! Upon successfull definition of the variable, 'nvarid' contains the
! NetCDF ID of the variable.
! The variables' attributes 'title' (Note that 'long_name' would be more
! appropriate) and 'units' are also set. 

implicit none

#include "netcdf.inc"

integer,intent(in) :: nid ! NetCDF file ID
character(len=*),intent(in) :: name ! the variable's name
character(len=*),intent(in) :: title ! 'title' attribute of variable
character(len=*),intent(in) :: units ! 'units' attribute of variable
integer,intent(in) :: nbdim ! number of dimensions of the variable
integer,dimension(nbdim),intent(in) :: dimids ! NetCDF IDs of the dimensions
                                              ! the variable is defined along
integer,intent(out) :: nvarid ! NetCDF ID of the variable
integer,intent(out) :: ierr ! returned NetCDF staus code

! 1. Switch to NetCDF define mode 
ierr=NF_REDEF(nid)

! 2. Define the variable
#ifdef NC_DOUBLE
ierr = NF_DEF_VAR (nid,adjustl(name),NF_DOUBLE,nbdim,dimids,nvarid)
#else
ierr = NF_DEF_VAR (nid,adjustl(name),NF_FLOAT,nbdim,dimids,nvarid)
#endif
if(ierr/=NF_NOERR) then
   write(*,*) "def_var_stats: Failed defining variable "//trim(name)
   write(*,*) NF_STRERROR(ierr)
   stop ""
endif

! 3. Write attributes
ierr=NF_PUT_ATT_TEXT(nid,nvarid,"title",&
                     len_trim(adjustl(title)),adjustl(title))
if(ierr/=NF_NOERR) then
   write(*,*) "def_var_stats: Failed writing title attribute for "//trim(name)
   write(*,*) NF_STRERROR(ierr)
   stop ""
endif

ierr=NF_PUT_ATT_TEXT(nid,nvarid,"units",&
                     len_trim(adjustl(units)),adjustl(units))
if(ierr/=NF_NOERR) then
   write(*,*) "def_var_stats: Failed writing units attribute for "//trim(name)
   write(*,*) NF_STRERROR(ierr)
   stop ""
endif

! 4. Switch out of NetCDF define mode
ierr = NF_ENDDEF(nid)

end subroutine def_var_stats