module output
!
! Save netcdf history file.
! Fields on either geographic or magnetic grids may be added to the file 
! at edynamo runtime by calling sub addfld (addfld_mod.F90).
!
  use netcdf
  use shr_kind_mod, only: r8 => shr_kind_r8            ! 8-byte reals
  use params,only: max_fields,fillvalue,iulog
  use my_mpi,only: mytid
  use fields,only: field,fields_out,field_init,nfields
  use read_ncfile,only: mtime,tlbc,ulbc,vlbc,lbc,year,iday,ut,nsteph,&
                        date,datesec,ndcur,nscur
  use geogrid,only: nlat,nlev,nilev,nlon,glat,glon,zlev,zilev
  use maggrid,only: nmlat,nmlev,nmlon,nmlonp1,gmlon,gmlat,zmlev
  use namelist,only: model_name,ctpoten
  implicit none
  save
  private
!
! Dimension and variable id's on netcdf file:
! (not including fields added with sub addfld)
!
  integer :: & 
    idlon,idlat,idlev,idilev,             & ! geographic dimension ids
    idmlon,idmlat,idmlev,                 & ! magnetic dimension ids
    idunlim,idmtime                         ! time dimension ids

  integer :: & 
    idv_lon,idv_lat,idv_lev,idv_ilev,           & ! geographic coord var ids
    idv_mlon,idv_mlat,idv_mlev,                 & ! magnetic coord var ids
    idv_lbc,idv_tlbc,idv_ulbc,idv_vlbc,         & ! lower boundaries
    idv_time,idv_mtime,idv_year,idv_day,idv_ut, & ! time variable ids
    idv_nsteph,idv_ctpoten,idv_date,idv_datesec,&
    idv_ndcur,idv_nscur
!
! if waccm_format=true,  then waccm lon =    0->360 and lev top->bottom
! if waccm_format=false, then waccm lon = -180->180 and lev bottom->top
!
  logical :: waccm_format = .true.  ! waccm lon =    0->360, lev=top2bottom
! logical :: waccm_format = .false. ! waccm lon = -180->180, lev=bottom2top

  character(len=1024) :: msg      ! error message

  public write_output,define_ncfile,output_init

  contains
!-----------------------------------------------------------------------
  subroutine write_output(ncfile,istep)
!
! Write a history for current timestep to output file:
!
! Args:
    integer,intent(in) :: istep
    character(len=*),intent(in) :: ncfile ! output file path
!
! Local:
    integer :: istat,n
    integer,save :: ncid            ! netcdf file id
    logical,save :: first = .true.  ! true if first call
!
! If first call, open new dataset and define dimension, variables, etc.
! If not first call, reopen existing dataset for appending.
!
    if (first) then
      istat = nf90_create(ncfile,NF90_CLOBBER,ncid)
      if (istat /= NF90_NOERR) then
        write(msg,"('Error opening output ncfile ',a)") trim(ncfile)
        call handle_ncerr(istat,msg,1)
      else
        write(iulog,'(/,2a)') 'Created output file ',trim(ncfile)
      endif
      call define_ncfile(ncid)
    else                        ! reopen for appending new timestep
      istat = nf90_open(ncfile,NF90_WRITE,ncid)
    endif
!
! Gather fields to root task and write to file:
!
    call gather_fields
    if (mytid==0) call write_fields(ncid,istep)

    istat = nf90_close(ncid)
    first = .false.
  end subroutine write_output
!-----------------------------------------------------------------------
  subroutine define_ncfile(ncid)
!
! Define dimensions and coordinate variables on netcdf output file.
! This is called once per run from write_output.
!
! Args:
    integer,intent(in) :: ncid
!
! Local:
    integer :: n,istat
!
! Define dimensions:
!
    istat = nf90_def_dim(ncid,'time',NF90_UNLIMITED,idunlim)

    istat = nf90_def_dim(ncid,'lon' ,nlon,  idlon)
    istat = nf90_def_dim(ncid,'lat' ,nlat,  idlat)

    istat = nf90_def_dim(ncid,'lev' ,nlev  ,idlev)
    istat = nf90_def_dim(ncid,'ilev',nilev ,idilev)

    istat = nf90_def_dim(ncid,'mlon' ,nmlon, idmlon)
    istat = nf90_def_dim(ncid,'mlat' ,nmlat, idmlat)
    istat = nf90_def_dim(ncid,'mlev' ,nmlev, idmlev)

    istat = nf90_def_dim(ncid,'mtimedim',3,  idmtime)
!
! Define coordinates and other vars:
!
    istat = nf90_def_var(ncid,'lon',NF90_REAL8,(/idlon/),idv_lon)
    istat = nf90_put_att(ncid,idv_lon,"long_name","geographic longitude (-west, +east)")
    istat = nf90_put_att(ncid,idv_lon,"units","degrees_east")

    istat = nf90_def_var(ncid,'lat',NF90_REAL8,(/idlat/),idv_lat)
    istat = nf90_put_att(ncid,idv_lat,"long_name","geographic latitude (-south, +north)")
    istat = nf90_put_att(ncid,idv_lat,"units","degrees_north")

    istat = nf90_def_var(ncid,'lev',NF90_REAL8,(/idlev/),idv_lev)
    istat = nf90_put_att(ncid,idv_lev,"long_name","midpoint levels")
    istat = nf90_put_att(ncid,idv_lev,"units","")
    istat = nf90_put_att(ncid,idv_lev,"positive","up")
    istat = nf90_put_att(ncid,idv_lev,"standard_name","atmosphere_ln_pressure_coordinate")
    istat = nf90_put_att(ncid,idv_lev,"formula_terms","p0:p0 lev:lev")
    istat = nf90_put_att(ncid,idv_lev,"formula","p(k) = p0 * exp(-lev(k))")

    istat = nf90_def_var(ncid,'ilev',NF90_REAL8,(/idilev/),idv_ilev)
    istat = nf90_put_att(ncid,idv_ilev,"long_name","interface levels")
    istat = nf90_put_att(ncid,idv_ilev,"units","")
    istat = nf90_put_att(ncid,idv_ilev,"positive","up")
    istat = nf90_put_att(ncid,idv_ilev,"standard_name","atmosphere_ln_pressure_coordinate")
    istat = nf90_put_att(ncid,idv_ilev,"formula_terms","p0:p0 lev:lev")
    istat = nf90_put_att(ncid,idv_ilev,"formula","p(k) = p0 * exp(-lev(k))")

    istat = nf90_def_var(ncid,'mlon',NF90_REAL8,(/idmlon/),idv_mlon)
    istat = nf90_put_att(ncid,idv_mlon,"long_name","magnetic longitude (-west, +east)")
    istat = nf90_put_att(ncid,idv_mlon,"units","degrees_east")

    istat = nf90_def_var(ncid,'mlat',NF90_REAL8,(/idmlat/),idv_mlat)
    istat = nf90_put_att(ncid,idv_mlat,"long_name","magnetic latitude (-south, +north)")
    istat = nf90_put_att(ncid,idv_mlat,"units","degrees_north")

    istat = nf90_def_var(ncid,'mlev',NF90_REAL8,(/idmlev/),idv_mlev)
    istat = nf90_put_att(ncid,idv_mlev,"long_name","magnetic midpoint levels")
    istat = nf90_put_att(ncid,idv_mlev,"units","")
    istat = nf90_put_att(ncid,idv_mlev,"positive","up")
    istat = nf90_put_att(ncid,idv_mlev,"standard_name","atmosphere_ln_pressure_coordinate")
    istat = nf90_put_att(ncid,idv_mlev,"formula_terms","p0:p0 mlev:mlev")
    istat = nf90_put_att(ncid,idv_mlev,"formula","p(k) = p0 * exp(-mlev(k))")

    if (trim(model_name)=='TIMEGCM') then

      istat = nf90_def_var(ncid,'LBC',NF90_REAL8,idv_lbc)
      istat = nf90_put_att(ncid,idv_lbc,"long_name",&
        "Interface level of t,u,v lower boundary condition")

      istat = nf90_def_var(ncid,'TLBC',NF90_REAL8,(/idlon,idlat,idunlim/),idv_tlbc)
      istat = nf90_put_att(ncid,idv_tlbc,"long_name","Lower boundary condition of TN")
      istat = nf90_put_att(ncid,idv_tlbc,"units","K")
      istat = nf90_put_att(ncid,idv_tlbc,"coordinates","LBC")

      istat = nf90_def_var(ncid,'ULBC',NF90_REAL8,(/idlon,idlat,idunlim/),idv_ulbc)
      istat = nf90_put_att(ncid,idv_ulbc,"long_name","Lower boundary condition of UN")
      istat = nf90_put_att(ncid,idv_ulbc,"units","cm/s")
      istat = nf90_put_att(ncid,idv_ulbc,"coordinates","LBC")

      istat = nf90_def_var(ncid,'VLBC',NF90_REAL8,(/idlon,idlat,idunlim/),idv_vlbc)
      istat = nf90_put_att(ncid,idv_vlbc,"long_name","Lower boundary condition of VN")
      istat = nf90_put_att(ncid,idv_vlbc,"units","cm/s")
      istat = nf90_put_att(ncid,idv_vlbc,"coordinates","LBC")

    endif ! TIMEGCM

    istat = nf90_def_var(ncid,'mtime',NF90_INT,(/idmtime,idunlim/),idv_mtime)
    istat = nf90_put_att(ncid,idv_mtime,"long_name","model times (day, hour, minute)")
    istat = nf90_put_att(ncid,idv_mtime,"units","day, hour, minute")

    istat = nf90_def_var(ncid,'time',NF90_REAL8,(/idunlim/),idv_time)
    istat = nf90_put_att(ncid,idv_time,"long_name","model time")
    istat = nf90_put_att(ncid,idv_time,"units","decimal julian day")

    istat = nf90_def_var(ncid,'year',NF90_INT,(/idunlim/),idv_year)
    istat = nf90_put_att(ncid,idv_year,"long_name","calendar year")

    istat = nf90_def_var(ncid,'day',NF90_INT,(/idunlim/),idv_day)
    istat = nf90_put_att(ncid,idv_day,"long_name","calendar day")

    istat = nf90_def_var(ncid,'ut',NF90_REAL8,(/idunlim/),idv_ut)
    istat = nf90_put_att(ncid,idv_ut,"long_name",&
      "universal time (from model time hour and minute)")
    istat = nf90_put_att(ncid,idv_ut,"units","hours")

    istat = nf90_def_var(ncid,'ctpoten',NF90_REAL8,(/idunlim/),idv_ctpoten)
    istat = nf90_put_att(ncid,idv_ctpoten,"long_name","Cross-tail potential")
    istat = nf90_put_att(ncid,idv_ctpoten,"units","Volts")

    if (trim(model_name)=='WACCM') then
      istat = nf90_def_var(ncid,'nsteph' ,NF90_INT,(/idunlim/),idv_nsteph)
      istat = nf90_put_att(ncid,idv_nsteph,'long_name','current timestep')

      istat = nf90_def_var(ncid,'date'   ,NF90_INT,(/idunlim/),idv_date)
      istat = nf90_put_att(ncid,idv_date,'long_name','current date (YYYYMMDD)')

      istat = nf90_def_var(ncid,'datesec',NF90_INT,(/idunlim/),idv_datesec)
      istat = nf90_put_att(ncid,idv_datesec,'long_name','current seconds of current date')

      istat = nf90_def_var(ncid,'ndcur'  ,NF90_INT,(/idunlim/),idv_ndcur)
      istat = nf90_put_att(ncid,idv_ndcur,'long_name','current day (from base day)')

      istat = nf90_def_var(ncid,'nscur'  ,NF90_INT,(/idunlim/),idv_nscur)
      istat = nf90_put_att(ncid,idv_nscur,'long_name','current seconds of current day')
    endif
!
! Take out of define mode:
!
    istat = nf90_enddef(ncid)
!
! Write coordinate values to the file:
!
! Shift longitude and level coordinates back to native waccm format:
! (do not shift back if waccm_format==.false.)
!
    if (trim(model_name)=='WACCM'.and.waccm_format) then
      call shift_lon(glon,nlon,'zeroto360',.true.)
      call reverse_vec(zlev,nlev)
      call reverse_vec(zilev,nilev)
!
! Shift mag coordinates to waccm format:
      call shift_lon(gmlon(1:nmlon),nmlon,'zeroto360',.true.)
      call reverse_vec(zmlev,nmlev)
    endif
!
    istat = nf90_put_var(ncid,idv_lon,glon)
    istat = nf90_put_var(ncid,idv_lat,glat)
    istat = nf90_put_var(ncid,idv_lev,zlev)
    istat = nf90_put_var(ncid,idv_ilev,zilev)
    istat = nf90_put_var(ncid,idv_mlon,gmlon(1:nmlon))
    istat = nf90_put_var(ncid,idv_mlat,gmlat)
    istat = nf90_put_var(ncid,idv_mlev,zmlev)

  end subroutine define_ncfile
!-----------------------------------------------------------------------
  subroutine gather_fields
!
! Gather output fields to the root task prior to writing to output file.
! Fields may be 2d (lon,lat), or 3d (lon,lat,lev), and on geo or mag grids.
!
  use my_mpi ,only: lon0,lon1,lat0,lat1,mlon0,mlon1,mlat0,mlat1,mp_gather2root
!
! Local
    integer :: n,k,j
    real(r8) :: fsub(lon0:lon1,lat0:lat1,nlev)
    integer :: nlonp4

    nlonp4 = nlon
    if (trim(model_name)=='TIMEGCM') nlonp4 = nlon+4

    do n=1,nfields
!     write(iulog,"('Enter gather_fields: ',a,' ndims=',i3,' geo=',l1)") &
!       trim(fields_out(n)%name),fields_out(n)%ndims,fields_out(n)%geo
!
! If field is already global (addfld was called by root task only),
! then the gather is not necessary.
!
      if (fields_out(n)%task0_only) cycle

      if (fields_out(n)%ndims == 2) then ! 2d field
        if (fields_out(n)%geo) then      ! 2d field on geo grid
          call mp_gather2root(fields_out(n)%data(lon0:lon1,lat0:lat1,1),&
            lon0,lon1,lat0,lat1,fields_out(n)%data,nlonp4,nlat,1,  &
            fields_out(n)%geo)
        else                             ! 2d field on mag grid
          call mp_gather2root(fields_out(n)%data(mlon0:mlon1,mlat0:mlat1,1),&
            mlon0,mlon1,mlat0,mlat1,fields_out(n)%data,nmlonp1,nmlat,1,  &
            fields_out(n)%geo)
        endif
      else                               ! 3d field
        if (fields_out(n)%geo) then      ! 3d field on geo grid
          fsub(lon0:lon1,lat0:lat1,:) = fields_out(n)%data(lon0:lon1,lat0:lat1,:)
          call mp_gather2root(fsub, &
            lon0,lon1,lat0,lat1,fields_out(n)%data,nlonp4,nlat,nlev,  &
            fields_out(n)%geo)
        else                             ! 3d field on mag grid
          call mp_gather2root(fields_out(n)%data(mlon0:mlon1,mlat0:mlat1,:), &
            mlon0,mlon1,mlat0,mlat1,fields_out(n)%data,nmlonp1,nmlat,nmlev,  &
            fields_out(n)%geo)
        endif
      endif

!     write(iulog,"('End gather_fields: ',a,' ndims=',i3,' geo=',l1)") &
!       trim(fields_out(n)%name),fields_out(n)%ndims,fields_out(n)%geo

    enddo ! n=1,nfields
  end subroutine gather_fields
!-----------------------------------------------------------------------
  subroutine write_fields(ncid,istep)
!
! Define and write fields to an open netcdf output file. These are the 2d
! or 3d fields added with calls to sub addfld. This routine is called at 
! every timestep by the root task only. If this is first call, the fields
! are defined on the file before writing the values.
!
! Field data has been gathered to the root task prior to this call.
! Fields will be 3d or 4d on the file (including unlimited time dimension),
! and on either geographic or magnetic grids.
!
! Args:
    integer,intent(in) :: ncid,istep
!
! Local:
    integer :: i,i0,i1,j,k,n,istat,iddims(4),varid
    logical,save :: first=.true.
    integer :: start(4),count(4),iyear1(1),iday1(1),ivar1(1)
    real(r8) :: time(1),ut1(1),ctpoten1(1)
!
! If this is first call, define the fields on the file:
!
    if (first) then
      istat = nf90_redef(ncid)
      do n=1,nfields
        do i=1,3
          if (index(fields_out(n)%dimnames(i),'lon' ) > 0) iddims(i) = idlon
          if (index(fields_out(n)%dimnames(i),'lat' ) > 0) iddims(i) = idlat
          if (index(fields_out(n)%dimnames(i),'lev' ) > 0) iddims(i) = idlev
          if (index(fields_out(n)%dimnames(i),'ilev') > 0) iddims(i) = idilev

          if (index(fields_out(n)%dimnames(i),'mlon')  > 0) iddims(i) = idmlon
          if (index(fields_out(n)%dimnames(i),'mlat')  > 0) iddims(i) = idmlat
          if (index(fields_out(n)%dimnames(i),'mlev')  > 0) iddims(i) = idmlev
        enddo

!       write(iulog,"('Define field: n=',i3,' field ',a,' dimnames=',3a6,' dimsizes=',3i4)") &
!         n,fields_out(n)%name(1:16),fields_out(n)%dimnames,fields_out(n)%dimsizes

        if (fields_out(n)%ndims == 2) then ! 2d field (3d with time on the file)
          iddims(3) = idunlim
          istat = nf90_def_var(ncid,fields_out(n)%name,NF90_DOUBLE,iddims(1:3),varid)
          if (istat /= NF90_NOERR) then
            write(msg,"('Defining 2d var ',a)") trim(fields_out(n)%name)
            call handle_ncerr(istat,msg,1)
          endif
        else                               ! 3d field (4d with time on the file)
          iddims(4) = idunlim
          istat = nf90_def_var(ncid,fields_out(n)%name,NF90_DOUBLE,iddims,varid)
          if (istat /= NF90_NOERR) then
            write(msg,"('Defining 3d var ',a)") trim(fields_out(n)%name)
            call handle_ncerr(istat,msg,1)
          endif
        endif
!
! Add long_name and units attributes for each field:
!
        istat = nf90_put_att(ncid,varid,"long_name",fields_out(n)%long_name)
        istat = nf90_put_att(ncid,varid,"units",fields_out(n)%units)
      enddo ! n=1,nfields
!
! Add global file attributes:
!
      istat = nf90_put_att(ncid,NF90_GLOBAL,"model_version","timegcm_edyn") ! place-holder
      istat = nf90_put_att(ncid,NF90_GLOBAL,"model_name"   ,"timegcm_edyn") ! place-holder
!
! Format (e.g., lon and lev coord ordering) will be in native format of model "model_name"
! (currently model_name is either 'TIMEGCM' or 'WACCM')
!
      istat = nf90_put_att(ncid,NF90_GLOBAL,"input_model_name",trim(model_name))
      istat = nf90_put_att(ncid,NF90_GLOBAL,"missing_value",fillvalue)
!
! Take out of define mode:
!
      istat = nf90_enddef(ncid)
    endif ! first

    if (trim(model_name)=='TIMEGCM') then
      i0 = 3 ; i1 = nlon+2 ! data is in 3:nlon+2
    else
      i0 = 1 ; i1 = nlon
    endif
!
! Write field values to the file for current time step:
!
    do n=1,nfields
      istat = nf90_inq_varid(ncid,fields_out(n)%name,varid)
      if (istat /= NF90_NOERR) call handle_ncerr(istat,'nf90_inq_varid',1)
      count = 0 ; start = 0
!
! Write 2d fields (i,j):
      if (fields_out(n)%ndims == 2) then ! 2d field (3d with time on the file)
!
! Reshift 2d fields to native waccm longitude sequence:
        if (trim(model_name)=='WACCM'.and.waccm_format) then
          do j=1,fields_out(n)%dimsizes(2)
            call shift_lon(fields_out(n)%data(:,j,1),&
                           fields_out(n)%dimsizes(1),'zeroto360',.false.)
          enddo
        endif

        start(1:2) = 1
        start(3) = istep
        count(2) = fields_out(n)%dimsizes(2)
        count(3) = 1
!
! Write 2d field on geographic grid (lon,lat):
        if (fields_out(n)%geo) then
          count(1) = nlon
          istat = nf90_put_var(ncid,varid,fields_out(n)%data(i0:i1,:,1), &
                               start(1:3),count(1:3))
!
! Write 2d field on magnetic grid:
        else
          count(1) = nmlon
          istat = nf90_put_var(ncid,varid,fields_out(n)%data(1:nmlon,:,1), &
                               start(1:3),count(1:3))
        endif
!
! Write 3d fields (lon,lat,lev):
      else                               ! 3d field (4d with time on the file)
!
! Reset to native WACCM lon and lev ordering:
!
        if (trim(model_name)=='WACCM'.and.waccm_format) then
          do k=1,fields_out(n)%dimsizes(3)
            do j=1,fields_out(n)%dimsizes(2)
              call shift_lon(fields_out(n)%data(:,j,k),&
                             fields_out(n)%dimsizes(1),'zeroto360',.false.)
            enddo
          enddo
          do j=1,fields_out(n)%dimsizes(2)
            do i=1,fields_out(n)%dimsizes(1)
              call reverse_vec(fields_out(n)%data(i,j,:),fields_out(n)%dimsizes(3))
            enddo
          enddo
        endif

        start(1:3) = 1
        start(4) = istep
        count(2) = fields_out(n)%dimsizes(2)
        count(3) = fields_out(n)%dimsizes(3)
        count(4) = 1
!
! Write 3d field on geographic grid:
        if (fields_out(n)%geo) then
          count(1) = nlon
          istat = nf90_put_var(ncid,varid,fields_out(n)%data(i0:i1,:,:),start,count)
!
! Write 3d field on magnetic grid:
        else
          count(1) = nmlon
          istat = nf90_put_var(ncid,varid,fields_out(n)%data(1:nmlon,:,:),start,count)
        endif
        if (istat /= NF90_NOERR) then
          write(msg,"('Writing 3d var ',a,' start=',4i3,' count=',4i3)") &
            trim(fields_out(n)%name),start,count
          call handle_ncerr(istat,msg,1)
        endif
      endif ! 2d or 3d field
    enddo ! n=1,nfields
!
! Lower boundaries (these were read from input file and are echoed here):
!
    if (trim(model_name)=='TIMEGCM') then
      istat = nf90_put_var(ncid,idv_lbc ,lbc)   ! bottom interface level
      istat = nf90_put_var(ncid,idv_tlbc,tlbc,(/1,1,istep/),(/nlon,nlat,1/))
      istat = nf90_put_var(ncid,idv_ulbc,ulbc,(/1,1,istep/),(/nlon,nlat,1/))
      istat = nf90_put_var(ncid,idv_vlbc,vlbc,(/1,1,istep/),(/nlon,nlat,1/))
    endif
!
! mtime is integer triplet model time (day,hour,minute)
!
    istat = nf90_put_var(ncid,idv_mtime,mtime,(/1,istep/),(/3,1/))
!
! Time is decimal day, calculated from mtime (day,hour,minute).
!
    time(1) = real(mtime(1))+real(mtime(2))/24._r8+real(mtime(3))/(60._r8*24._r8) 
    istat = nf90_put_var(ncid,idv_time,time,(/istep/),(/1/))

    iyear1(1) = int(year)
    istat = nf90_put_var(ncid,idv_year,iyear1,(/istep/),(/1/))
    iday1(1) = iday
    istat = nf90_put_var(ncid,idv_day ,iday1 ,(/istep/),(/1/))
    ut1(1) = ut
    istat = nf90_put_var(ncid,idv_ut  ,ut1   ,(/istep/),(/1/))
    ctpoten1(1) = ctpoten
    istat = nf90_put_var(ncid,idv_ctpoten,ctpoten1,(/istep/),(/1/))

    if (trim(model_name)=='WACCM') then
      ivar1(1) = nsteph
      istat = nf90_put_var(ncid,idv_nsteph ,ivar1,(/istep/),(/1/))
      ivar1(1) = date
      istat = nf90_put_var(ncid,idv_date   ,ivar1,(/istep/),(/1/))
      ivar1(1) = datesec
      istat = nf90_put_var(ncid,idv_datesec,ivar1,(/istep/),(/1/))
      ivar1(1) = ndcur
      istat = nf90_put_var(ncid,idv_ndcur,ivar1,(/istep/),(/1/))
      ivar1(1) = nscur
      istat = nf90_put_var(ncid,idv_nscur,ivar1,(/istep/),(/1/))
    endif
    
    first = .false.
  end subroutine write_fields
!-----------------------------------------------------------------------
  subroutine output_init
!
! Local:
    integer :: n

    do n=1,max_fields
      call field_init(fields_out(n))
    enddo
  end subroutine output_init
!-----------------------------------------------------------------------
end module output