module wstats_mod

! module containing parameters and routines to generate the "stats.nc" file
! which will contain a mean statistical day of variables extracted at set
! times of day every day of the simulation, and the standard deviations thereof

implicit none

logical,save :: callstats ! global flag to trigger generating a stats.nc
                          ! file or not. Initialized in conf_gcm()
!$OMP THREADPRIVATE(callstats)

integer,save :: istats ! calculate stats every istats physics timestep,
                       ! starting at first call. Initialized by inistats()
!$OMP THREADPRIVATE(istats)

integer,parameter :: istime=12 ! number of time steps per sol to store

integer,save :: count(istime) ! count tab to know the variable record
!$OMP THREADPRIVATE(count)

contains

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

! Main routine to call to trigger writing a given field to the "stats.nc" file 

use mod_phys_lmdz_para, only : is_mpi_root, is_master, klon_mpi_begin
use mod_phys_lmdz_transfert_para, only: gather, bcast
use mod_grid_phy_lmdz, only : klon_glo, Grid1Dto2D_glo, &
                              nbp_lon, nbp_lat, nbp_lev, &
                              grid_type, unstructured
implicit none

include "netcdf.inc"

integer,intent(in) :: ngrid
character (len=*),intent(in) :: nom,titre,unite
integer,intent(in) :: dim
real,intent(in) :: px(ngrid,nbp_lev)
real,allocatable,save :: mean3d(:,:,:),sd3d(:,:,:),dx3(:,:,:)
real,allocatable,save :: mean2d(:,:),sd2d(:,:),dx2(:,:)
character (len=50) :: namebis
character (len=50), save :: firstvar
!$OMP THREADPRIVATE(mean3d,sd3d,dx3,mean2d,sd2d,dx2,firstvar)
integer :: ierr,varid,nbdim,nid
integer :: meanid,sdid
integer, dimension(4)  :: id,start,sizes
logical, save :: firstcall=.TRUE.
integer,save :: indx
integer, save :: step=0
!$OMP THREADPRIVATE(firstcall,indx,step)
integer :: l,i,j,ig0,n

! Added to read an optional stats.def file to select outputs
logical,save :: check_stats_def=.TRUE. ! to keep track of if stats.def was checked
logical,save :: stats_def ! .true. if there is a stats.def file
integer,save :: n_name_def ! number of fields to output in stats.nc
integer,parameter :: n_name_def_max=199 ! max number of fields to output
character(len=120),save :: name_def(n_name_def_max)
logical :: getout ! to trigger an early exit if variable not in output list

!$OMP THREADPRIVATE(check_stats_def,stats_def,n_name_def,name_def)

! Added to work in parallel mode
#ifdef CPP_PARA
real px3_glop(klon_glo,nbp_lev) ! to store a 3D data set on global physics grid
real px3_glo(nbp_lon,nbp_lat,nbp_lev) ! 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(nbp_lon,nbp_lat) ! to store a 2D data set on global lonxlat grid
real px2(ngrid)
real px3(ngrid,nbp_lev)
#else
! When not running in parallel mode:
real px3_glop(ngrid,nbp_lev) ! to store a 3D data set on global physics grid
real px3_glo(nbp_lon,nbp_lat,nbp_lev) ! 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(nbp_lon,nbp_lat) ! to store a 2D data set on global lonxlat grid
#endif

! 0. Preliminary stuff 
if (callstats.eqv..false.) then
  ! exit because creating/writing stats.nc not requested by user
  return 
endif

if (grid_type==unstructured) then
  ! exit because non-structured grid case is not handled
  return
endif

if (check_stats_def) then
  ! this step only needs be made a single unique time
  ! hence the chack_stats_def flag.
  
  if (is_master) then
   ! Only the master reads the file
   ! open stats.def definition file if there is one
   open(99,file="stats.def",status='old',form='formatted',&
        iostat=ierr)
   if (ierr.eq.0) then
    stats_def=.true. ! yes there is a stats.def file
    write(*,*) "*****************"
    write(*,*) "Reading stats.def"
    write(*,*) "*****************"
    do n=1,n_name_def_max
      read(99,fmt='(a)',end=88) name_def(n)
      write(*,*) 'Output in stats: ', trim(name_def(n))
    enddo
88  continue
    ! check there is no overflow
    if (n.ge.n_name_def_max) then
      write(*,*) "n_name_def_max too small in wstats:",n
      call abort_physic("wstats","n_name_def_max too small",1)
    endif
    n_name_def=n-1
    close(99)
   else
    stats_def=.false. ! no stats.def file; output all fields sent to wstats
   endif ! of if (ierr.eq.0)

  endif ! of if (is_master)
    
  ! share the info gathered by master with all other cores
  call bcast(stats_def)
  call bcast(n_name_def)
  do n=1,n_name_def
    call bcast(name_def(n))
  enddo
  
  check_stats_def=.false.
endif ! of if (check_stats_def)

! 1. Initialization (creation of stats.nc file)
if (firstcall) then
   firstcall=.false.

   ! identify the first variable name with which wstats is called in
   ! a time step
   if (stats_def) then
     ! check if "nom" is in the wanted output list. If it is then
     ! it is the sought "first variable"; if not then skip all the
     ! rest and reset "firstcall" to true
     getout=.true.
     do n=1,n_name_def
       ! look for the variable's name in the list
       if (trim(name_def(n)).eq.nom) then
         getout=.false.
         ! found it, no need to scan the rest of the list exit loop
         firstvar=trim((nom))
         exit
       endif
     enddo
     if (getout) then
       ! variable not in the list so exit routine now
       firstcall=.true.
       return
     endif
   else
     ! without a stats.def, the very first call to stats contains first
     ! variable name
     firstvar=trim((nom))
   endif ! of if (stats_def)
   
   ! initialize stats.nc file
   call inistats(ierr)
   
   ! allocate arrays that will be needed later
   if (klon_glo>1) then ! general case, 3D GCM
     allocate(mean3d(nbp_lon+1,nbp_lat,nbp_lev))
     allocate(sd3d(nbp_lon+1,nbp_lat,nbp_lev))
     allocate(dx3(nbp_lon+1,nbp_lat,nbp_lev))
     allocate(mean2d(nbp_lon+1,nbp_lat))
     allocate(sd2d(nbp_lon+1,nbp_lat))
     allocate(dx2(nbp_lon+1,nbp_lat))
   else ! 1D model
     allocate(mean3d(1,1,nbp_lev))
     allocate(sd3d(1,1,nbp_lev))
     allocate(dx3(1,1,nbp_lev))
     allocate(mean2d(1,1))
     allocate(sd2d(1,1))
     allocate(dx2(1,1))
   endif

endif ! of if (firstcall)

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

! Exit if there is a stats.def file and the variable is not in the list
if (stats_def) then
  getout=.true.
  do n=1,n_name_def
    ! look for the variable's name in the list
    if (trim(name_def(n)).eq.nom) then
      getout=.false.
      ! found it, no need to scan the rest of the list exit loop
      exit
    endif
  enddo
  if (getout) then
    ! variable not in the list so exit routine now
    return
  endif
endif ! of if (stats_def) 

! collect fields on a global physics grid
#ifdef CPP_PARA
 if (dim.eq.3) then
  px3(1:ngrid,1:nbp_lev)=px(1:ngrid,1:nbp_lev)
  ! 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)
    if (klon_glo>1) then ! general case (unless in 1D)
    ! copy dx3_glo() to dx3(:) and add redundant longitude
    dx3(1:nbp_lon,:,:)=px3_glo(1:nbp_lon,:,:)
    dx3(nbp_lon+1,:,:)=dx3(1,:,:)
    endif
  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)
    if (klon_glo>1) then ! general case (unless in 1D)
      ! copy px2_glo() to dx2(:) and add redundant longitude
      dx2(1:nbp_lon,:)=px2_glo(1:nbp_lon,:)
      dx2(nbp_lon+1,:)=dx2(1,:)
    endif
  endif
!$OMP END MASTER
!$OMP BARRIER
 endif
#else
  if (dim.eq.3) then
    px3_glop(:,1:nbp_lev)=px(:,1:nbp_lev)
!  Passage variable physique -->  variable dynamique
    DO l=1,nbp_lev
      DO i=1,nbp_lon
         px3_glo(i,1,l)=px(1,l)
         px3_glo(i,nbp_lat,l)=px(ngrid,l)
      ENDDO
      DO j=2,nbp_lat-1
         ig0= 1+(j-2)*nbp_lon
         DO i=1,nbp_lon
            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,nbp_lon
     px2_glo(i,1)=px(1,1)
     px2_glo(i,nbp_lat)=px(ngrid,1)
   ENDDO
   DO j=2,nbp_lat-1
     ig0= 1+(j-2)*nbp_lon
     DO i=1,nbp_lon
        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: creating ",nom
   namebis=trim(nom)
   call def_var_stats(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_stats(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/)
      if (klon_glo>1) then !general case
        sizes=(/nbp_lon+1,nbp_lat,nbp_lev,1/)
      else
        sizes=(/1,1,nbp_lev,1/)
      endif
      mean3d(:,:,:)=0
      sd3d(:,:,:)=0
   else if (dim.eq.2) then
      start=(/1,1,indx,0/)
      if (klon_glo>1) then !general case
        sizes=(/nbp_lon+1,nbp_lat,1,0/)
      else
        sizes=(/1,1,1,0/)
      endif
      mean2d(:,:)=0
      sd2d(:,:)=0
   endif
else
   ! load values from file
   if (dim.eq.3) then
      start=(/1,1,1,indx/)
      if (klon_glo>1) then !general case
        sizes=(/nbp_lon+1,nbp_lat,nbp_lev,1/)
      else ! 1D model case
        sizes=(/1,1,nbp_lev,1/)
      endif
#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 (*,*) "wstats error reading :",trim(nom)
         write (*,*) NF_STRERROR(ierr)
         call abort_physic("wstats","Failed reading "//trim(nom),1)
      endif

   else if (dim.eq.2) then
      start=(/1,1,indx,0/)
      if (klon_glo>1) then ! general case
        sizes=(/nbp_lon+1,nbp_lat,1,0/)
      else
        sizes=(/1,1,1,0/)
      endif
#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 (*,*) "wstats error reading :",trim(nom)
         write (*,*) NF_STRERROR(ierr)
         call abort_physic("wstats","Failed reading "//trim(nom),1)
      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:nbp_lon,:,:)=px3_glo(:,:,:)
  IF (klon_glo>1) THEN ! in 3D, add redundant longitude point
    dx3(nbp_lon+1,:,:)=dx3(1,:,:)
  ENDIF
else ! dim.eq.2
  dx2(1:nbp_lon,:)=px2_glo(:,:)
  IF (klon_glo>1) THEN ! in 3D, add redundant longitude point
    dx2(nbp_lon+1,:)=dx2(1,:)
  ENDIF
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
  if (ierr.ne.NF_NOERR) then
     write (*,*) "wstats error writing :",trim(nom)
     write (*,*) NF_STRERROR(ierr)
     call abort_physic("wstats","Failed writing "//trim(nom),1)
  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
  if (ierr.ne.NF_NOERR) then
     write (*,*) "wstats error writing :",trim(nom)
     write(*,*) "start:",start
     write(*,*) "sizes:",sizes
     write(*,*) "mean2d:",mean2d
     write(*,*) "sd2d:",sd2d
     write (*,*) NF_STRERROR(ierr)
     call abort_physic("wstats","Failed writing "//trim(nom),1)
  endif

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

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

end subroutine wstats

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

subroutine inistats(ierr)

! routine to initialize the stats file (i.e. create file, write
! time independent coordinates, etc.)

use mod_phys_lmdz_para, only : is_master
USE vertical_layers_mod, ONLY: ap,bp,aps,bps,preff,pseudoalt,presnivs
USE nrtype, ONLY: pi
USE time_phylmdz_mod, ONLY: daysec,dtphys
USE regular_lonlat_mod, ONLY: lon_reg, lat_reg
USE mod_grid_phy_lmdz, ONLY: nbp_lon, nbp_lat, nbp_lev
implicit none

include "netcdf.inc"

integer,intent(out) :: ierr

integer :: nid
integer :: l,nsteppd
real, dimension(nbp_lev) ::  sig_s
real,allocatable :: lon_reg_ext(:) ! extended longitudes
integer :: idim_lat,idim_lon,idim_llm,idim_llmp1,idim_time
real, dimension(istime) :: lt
integer :: nvarid


IF (nbp_lon*nbp_lat==1) THEN
  ! 1D model
  ALLOCATE(lon_reg_ext(1))
ELSE
  ! 3D model
  ALLOCATE(lon_reg_ext(nbp_lon+1))
ENDIF

write (*,*) 
write (*,*) '                        || STATS ||'
write (*,*) 
write (*,*) 'daysec',daysec
write (*,*) 'dtphys',dtphys
nsteppd=nint(daysec/dtphys)
write (*,*) 'nsteppd=',nsteppd

if (abs(float(nsteppd)-daysec/dtphys).gt.1.e-8*daysec) then
  call abort_physic("inistats",'1 sol .ne. n physics steps',1)
endif

if(mod(nsteppd,istime).ne.0) then
  call abort_physic("inistats",'1 sol .ne. n*istime physics steps',1)
endif

istats=nsteppd/istime
write (*,*) 'istats=',istats
write (*,*) 'Storing ',istime,'times per day'
write (*,*) 'thus every ',istats,'physical timestep '
write (*,*) 

do l= 1, nbp_lev
  sig_s(l)=((ap(l)+ap(l+1))/preff+bp(l)+bp(l+1))/2.
  pseudoalt(l)=-10.*log(presnivs(l)/preff)   
enddo

lon_reg_ext(1:nbp_lon)=lon_reg(1:nbp_lon)
IF (nbp_lon*nbp_lat/=1) THEN
  ! In 3D, add extra redundant point (180 degrees,
  ! since lon_reg starts at -180)
  lon_reg_ext(nbp_lon+1)=-lon_reg_ext(1)
ENDIF

if (is_master) then
  ! only the master needs do this
  ierr = NF_CREATE("stats.nc",IOR(NF_CLOBBER,NF_64BIT_OFFSET),nid)
  if (ierr.ne.NF_NOERR) then
    write (*,*) NF_STRERROR(ierr)
    call abort_physic("inistats",'failed creating stats.nc',1)
  endif

  ierr = NF_DEF_DIM (nid, "latitude", nbp_lat, idim_lat)
  IF (nbp_lon*nbp_lat==1) THEN
    ierr = NF_DEF_DIM (nid, "longitude", 1, idim_lon)
  ELSE
    ierr = NF_DEF_DIM (nid, "longitude", nbp_lon+1, idim_lon)
  ENDIF
  ierr = NF_DEF_DIM (nid, "altitude", nbp_lev, idim_llm)
  ierr = NF_DEF_DIM (nid, "llmp1", nbp_lev+1, idim_llmp1)
  ierr = NF_DEF_DIM (nid, "Time", NF_UNLIMITED, idim_time)

  ierr = NF_ENDDEF(nid)

  call def_var_stats(nid,"Time","Time", &
                     "days since 0000-00-0 00:00:00",1, &
                     [idim_time],nvarid,ierr)
  ! Time is initialised later by mkstats subroutine

  call def_var_stats(nid,"latitude","latitude", &
                     "degrees_north",1,[idim_lat],nvarid,ierr)
#ifdef NC_DOUBLE
  ierr = NF_PUT_VAR_DOUBLE (nid,nvarid,lat_reg/pi*180)
#else
  ierr = NF_PUT_VAR_REAL (nid,nvarid,lat_reg/pi*180)
#endif

  call def_var_stats(nid,"longitude","East longitude", &
                     "degrees_east",1,[idim_lon],nvarid,ierr)
#ifdef NC_DOUBLE
  ierr = NF_PUT_VAR_DOUBLE (nid,nvarid,lon_reg_ext/pi*180)
#else
  ierr = NF_PUT_VAR_REAL (nid,nvarid,lon_reg_ext/pi*180)
#endif

  ! Niveaux verticaux, aps et bps
  ierr = NF_REDEF (nid)
#ifdef NC_DOUBLE
  ierr = NF_DEF_VAR (nid,"altitude", NF_DOUBLE, 1,[idim_llm],nvarid)
#else
  ierr = NF_DEF_VAR (nid,"altitude", NF_FLOAT, 1,[idim_llm],nvarid)
#endif
  ierr = NF_PUT_ATT_TEXT (nid,nvarid,"long_name",8,"altitude")
  ierr = NF_PUT_ATT_TEXT (nid,nvarid,'units',2,"km")
  ierr = NF_PUT_ATT_TEXT (nid,nvarid,'positive',2,"up")
  ierr = NF_ENDDEF(nid)
#ifdef NC_DOUBLE
  ierr = NF_PUT_VAR_DOUBLE (nid,nvarid,pseudoalt)
#else
  ierr = NF_PUT_VAR_REAL (nid,nvarid,pseudoalt)
#endif 
  call def_var_stats(nid,"aps","hybrid pressure at midlayers", &
                     " ",1,[idim_llm],nvarid,ierr)
#ifdef NC_DOUBLE
  ierr = NF_PUT_VAR_DOUBLE (nid,nvarid,aps)
#else
  ierr = NF_PUT_VAR_REAL (nid,nvarid,aps)
#endif

  call def_var_stats(nid,"bps","hybrid sigma at midlayers", &
                     " ",1,[idim_llm],nvarid,ierr)
#ifdef NC_DOUBLE
  ierr = NF_PUT_VAR_DOUBLE (nid,nvarid,bps)
#else 
  ierr = NF_PUT_VAR_REAL (nid,nvarid,bps)
#endif

  ierr=NF_CLOSE(nid)

endif ! of if (is_master)

end subroutine inistats

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

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

! routine to initialize writing a variable in the stats file

use mod_grid_phy_lmdz, only : nbp_lon, nbp_lat, nbp_lev, klon_glo

implicit none

include "netcdf.inc"

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

integer :: l,i,j,ig0
integer, dimension(4) :: start,sizes
real, dimension(nbp_lon+1,nbp_lat,nbp_lev) :: dx3
real, dimension(nbp_lon+1,nbp_lat) :: dx2
real :: dx3_1d(nbp_lev) ! for 1D outputs
real :: dx2_1d ! for 1D outputs

if (dim.eq.3) then

   start=(/1,1,1,indx/)
   if (klon_glo>1) then ! general 3D case
     sizes=(/nbp_lon+1,nbp_lat,nbp_lev,1/)
   else
     sizes=(/1,1,nbp_lev,1/)
   endif

!  Passage variable physique -->  variable dynamique

   if (klon_glo>1) then ! general case
    DO l=1,nbp_lev
      DO i=1,nbp_lon+1
         dx3(i,1,l)=px(1,l)
         dx3(i,nbp_lat,l)=px(ngrid,l)
      ENDDO
      DO j=2,nbp_lat-1
         ig0= 1+(j-2)*nbp_lon
         DO i=1,nbp_lon
            dx3(i,j,l)=px(ig0+i,l)
         ENDDO
         dx3(nbp_lon+1,j,l)=dx3(1,j,l)
      ENDDO
    ENDDO
   else ! 1D model case
     dx3_1d(1:nbp_lev)=px(1,1:nbp_lev)
   endif

#ifdef NC_DOUBLE
   if (klon_glo>1) then
     ierr = NF_PUT_VARA_DOUBLE(nid,varid,start,sizes,dx3)
   else
     ierr = NF_PUT_VARA_DOUBLE(nid,varid,start,sizes,dx3_1d)
   endif
#else
   if (klon_glo>1) then
     ierr = NF_PUT_VARA_REAL(nid,varid,start,sizes,dx3)
   else
     ierr = NF_PUT_VARA_REAL(nid,varid,start,sizes,dx3_1d)
   endif
#endif
  if (ierr.ne.NF_NOERR) then
     write (*,*) "inivar error writing variable"
     write (*,*) NF_STRERROR(ierr)
     call abort_physic("inivar","error writing variable",1)
  endif

else if (dim.eq.2) then

      start=(/1,1,indx,0/)
      if (klon_glo>1) then ! general 3D case
        sizes=(/nbp_lon+1,nbp_lat,1,0/)
      else
        sizes=(/1,1,1,0/)
      endif

!    Passage variable physique -->  physique dynamique

  if (klon_glo>1) then ! general case
  DO i=1,nbp_lon+1
     dx2(i,1)=px(1,1)
     dx2(i,nbp_lat)=px(ngrid,1)
  ENDDO
  DO j=2,nbp_lat-1
     ig0= 1+(j-2)*nbp_lon
     DO i=1,nbp_lon
        dx2(i,j)=px(ig0+i,1)
     ENDDO
     dx2(nbp_lon+1,j)=dx2(1,j)
  ENDDO
  else ! 1D model case
    dx2_1d=px(1,1)
  endif
  
#ifdef NC_DOUBLE
   if (klon_glo>1) then
     ierr = NF_PUT_VARA_DOUBLE(nid,varid,start,sizes,dx2)
   else
     ierr = NF_PUT_VARA_DOUBLE(nid,varid,start,sizes,dx2_1d)
   endif
#else
   if (klon_glo>1) then
     ierr = NF_PUT_VARA_REAL(nid,varid,start,sizes,dx2)
   else
     ierr = NF_PUT_VARA_REAL(nid,varid,start,sizes,dx2_1d)
   endif
#endif
  if (ierr.ne.NF_NOERR) then
     write (*,*) "inivar error writing variable"
     write (*,*) NF_STRERROR(ierr)
     call abort_physic("inivar","error writing variable",1)
  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)
   call abort_physic("def_var_stats","Failed defining "//trim(name),1)
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)
   call abort_physic("def_var_stats","Failed writing title for "//trim(name),1)
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)
   call abort_physic("def_var_stats","Failed writing units for "//trim(name),1)
endif

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

end subroutine def_var_stats

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

subroutine mkstats(ierr)

!  This routine finalizes tha stats.nc file from sums and sums of squares
!  to means and standard deviations. The data file is overwritten in place.

use mod_phys_lmdz_para, only : is_master
use mod_grid_phy_lmdz, only : nbp_lon, nbp_lat, nbp_lev, klon_glo

implicit none

include "netcdf.inc"

integer,intent(out) :: ierr ! netCDF return error code
integer :: nid,nbvar,i,ndims,lt,nvarid
integer, dimension(4) :: id,varid,start,size
integer, dimension(5) :: dimids
character (len=50) :: name,nameout,units,title
real,allocatable :: sum3d(:,:,:),square3d(:,:,:),mean3d(:,:,:),sd3d(:,:,:)
real,allocatable :: sum2d(:,:),square2d(:,:),mean2d(:,:),sd2d(:,:)
real, dimension(istime) :: time
real, dimension(nbp_lat) :: lat
real,allocatable :: lon(:)
real, dimension(nbp_lev) :: alt
logical :: lcopy=.true.
!integer :: latid,lonid,altid,timeid
integer :: meanid,sdid
!integer, dimension(4) :: dimout

if (is_master) then
! only the master needs do all this

! Incrementation of count for the last step, which is not done in wstats
count(istime)=count(istime)+1

if (klon_glo>1) then
  allocate(lon(nbp_lon+1))
  allocate(sum3d(nbp_lon+1,nbp_lat,nbp_lev))
  allocate(square3d(nbp_lon+1,nbp_lat,nbp_lev))
  allocate(mean3d(nbp_lon+1,nbp_lat,nbp_lev))
  allocate(sd3d(nbp_lon+1,nbp_lat,nbp_lev))
  allocate(sum2d(nbp_lon+1,nbp_lat))
  allocate(square2d(nbp_lon+1,nbp_lat))
  allocate(mean2d(nbp_lon+1,nbp_lat))
  allocate(sd2d(nbp_lon+1,nbp_lat))
else ! 1D model case
  allocate(lon(1))
  allocate(sum3d(1,1,nbp_lev))
  allocate(square3d(1,1,nbp_lev))
  allocate(mean3d(1,1,nbp_lev))
  allocate(sd3d(1,1,nbp_lev))
  allocate(sum2d(1,1))
  allocate(square2d(1,1))
  allocate(mean2d(1,1))
  allocate(sd2d(1,1))
endif

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

! We catch the id of dimensions of the stats file

ierr= NF_INQ_DIMID(nid,"latitude",id(1))
ierr= NF_INQ_DIMID(nid,"longitude",id(2))
ierr= NF_INQ_DIMID(nid,"altitude",id(3))
ierr= NF_INQ_DIMID(nid,"Time",id(4))

ierr= NF_INQ_VARID(nid,"latitude",varid(1))
ierr= NF_INQ_VARID(nid,"longitude",varid(2))
ierr= NF_INQ_VARID(nid,"altitude",varid(3))
ierr= NF_INQ_VARID(nid,"Time",varid(4))

! Time initialisation

do i=1,istime
   time(i)=i*24./istime
#ifdef NC_DOUBLE
   ierr= NF_PUT_VARA_DOUBLE(nid,varid(4),[i],[1],time(i))
#else
   ierr= NF_PUT_VARA_REAL(nid,varid(4),[i],[1],time(i))
#endif
enddo

! We catche the values of the variables

#ifdef NC_DOUBLE
         ierr = NF_GET_VAR_DOUBLE(nid,varid(1),lat)
         ierr = NF_GET_VAR_DOUBLE(nid,varid(2),lon)
         ierr = NF_GET_VAR_DOUBLE(nid,varid(3),alt)
#else
         ierr = NF_GET_VAR_REAL(nid,varid(1),lat)
         ierr = NF_GET_VAR_REAL(nid,varid(2),lon)
         ierr = NF_GET_VAR_REAL(nid,varid(3),alt)
#endif

! We catch the number of variables in the stats file
ierr = NF_INQ_NVARS(nid,nbvar)

! to catche the "real" number of variables (without the "additionnal variables")
nbvar=(nbvar-4)/2 

do i=1,nbvar
   nvarid=(i-1)*2+5

   ! What's the variable's name?
   ierr=NF_INQ_VARNAME(nid,nvarid,name)
   write(*,*) "OK variable ",name
   ! Its units?
   units=" "
   ierr=NF_GET_ATT_TEXT(nid,nvarid,"units",units)
   ! Its title?
   title=" "
   ierr=NF_GET_ATT_TEXT(nid,nvarid,"title",title)
   ! Its number of dimensions?   
   ierr=NF_INQ_VARNDIMS(nid,nvarid,ndims)
   ! Its values?

   if(ndims==4) then ! lat, lon, alt & time

!      dimout(1)=lonid
!      dimout(2)=latid
!      dimout(3)=altid
!      dimout(4)=timeid

      if (klon_glo>1) then ! general case
        size=(/nbp_lon+1,nbp_lat,nbp_lev,1/)
      else ! 1D model
        size=(/1,1,nbp_lev,1/)
      endif
      do lt=1,istime
         start=(/1,1,1,lt/)
         ! Extraction of the "source" variables
#ifdef NC_DOUBLE
         ierr = NF_GET_VARA_DOUBLE(nid,nvarid,start,size,sum3d)
         ierr = NF_GET_VARA_DOUBLE(nid,nvarid+1,start,size,square3d)
#else
         ierr = NF_GET_VARA_REAL(nid,nvarid,start,size,sum3d)
         ierr = NF_GET_VARA_REAL(nid,nvarid+1,start,size,square3d)
#endif
         ! Calculation of these nvariables
         mean3d=sum3d/count(lt)
         sd3d=sqrt(max(0.,square3d/count(lt)-mean3d**2))
         ! Writing of the variables
#ifdef NC_DOUBLE
         ierr = NF_PUT_VARA_DOUBLE(nid,nvarid,start,size,mean3d)
         ierr = NF_PUT_VARA_DOUBLE(nid,nvarid+1,start,size,sd3d)
#else
         ierr = NF_PUT_VARA_REAL(nid,nvarid,start,size,mean3d)
         ierr = NF_PUT_VARA_REAL(nid,nvarid+1,start,size,sd3d)
#endif
      enddo

    else if (ndims.eq.3) then

!      dimout(1)=lonid
!      dimout(2)=latid
!      dimout(3)=timeid

      if (klon_glo>1) then ! general case
        size=(/nbp_lon+1,nbp_lat,1,0/)
      else
        size=(/1,1,1,0/)
      endif
      do lt=1,istime
         start=(/1,1,lt,0/)
         ! Extraction of the "source" variables
#ifdef NC_DOUBLE
         ierr = NF_GET_VARA_DOUBLE(nid,nvarid,start,size,sum2d)
         ierr = NF_GET_VARA_DOUBLE(nid,nvarid+1,start,size,square2d)
#else
         ierr = NF_GET_VARA_REAL(nid,nvarid,start,size,sum2d)
         ierr = NF_GET_VARA_REAL(nid,nvarid+1,start,size,square2d)
#endif
         ! Calculation of these variables
         mean2d=sum2d/count(lt)
         sd2d=sqrt(max(0.,square2d/count(lt)-mean2d**2))
         ! Writing of the variables
#ifdef NC_DOUBLE
         ierr = NF_PUT_VARA_DOUBLE(nid,nvarid,start,size,mean2d)
         ierr = NF_PUT_VARA_DOUBLE(nid,nvarid+1,start,size,sd2d)
#else
         ierr = NF_PUT_VARA_REAL(nid,nvarid,start,size,mean2d)
         ierr = NF_PUT_VARA_REAL(nid,nvarid+1,start,size,sd2d)
#endif
      enddo

    endif 
enddo

ierr= NF_CLOSE(nid)

endif ! of if (is_master)

end subroutine mkstats

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

end module wstats_mod