module RvicRestFile

!-----------------------------------------------------------------------
!BOP
!
! !MODULE: restFileMod
!
! !DESCRIPTION:
! Reads from or writes to/ the RVIC restart file.
!
! !USES:
  use shr_kind_mod  , only : r8 => shr_kind_r8
  use shr_sys_mod   , only : shr_sys_abort
  use RvicSpmd       , only : masterproc
  use RvicHistFile   , only : RvicHistRestart
  use RvicFileUtils  , only : relavu, getavu, opnfil, getfil
  use RvicTimeManager, only : timemgr_restart, get_nstep, get_curr_date, is_last_step
  use RvicIO
  use RvicDateTime
  use rvic_cpl_indices , only : nt_rvic, rvic_tracers
  use RvicVar
  use netcdf

  implicit none

  ! PUBLIC MEMBER FUNCTIONS:
  public :: RvicRestFileWrite
  public :: RvicRestGetfile
  public :: RvicRestFileRead        ! rvic read state routine
  public :: RvicRestart

! !PRIVATE MEMBER FUNCTIONS:
  private :: RvicRestFileName
  private :: restFile_read_pfile
  private :: restFile_write_pfile    ! Writes restart pointer file
  private :: handle_err

  ! PUBLIC DATA MEMBERS:
  ! None

  ! REVISION HISTORY:
  ! Author:  Joe Hamman, University of Washington
  
  contains

!-----------------------------------------------------------------------

!=======================================================================

  subroutine RvicRestFileWrite( rdate )

    !-----------------------------------------------------------------------
    ! !DESCRIPTION:
    ! Read/write RVIC restart file.

    ! !ARGUMENTS:
    implicit none
    !character(len=*) , intent(in) :: file            ! output netcdf restart file
    character(len=*) , intent(in) :: rdate           ! restart file time stamp for name

    ! !LOCAL VARIABLES:
    character(len=256)  :: restfile
    type(file_desc_t) :: ncid ! netcdf id
    integer :: i       ! index
    logical :: ptrfile ! write out the restart pointer file
    character(len=*),parameter :: subname='RvicRestFileWrite'
    !-----------------------------------------------------------------------

    ! Define dimensions and variables

    if (masterproc) then        
       write(iulog,*)
       write(iulog,*)'restFile_open: writing RVIC restart dataset '
       write(iulog,*)
    end if
    restfile = RvicRestFileName( rdate )
    write(iulog,*) subname,' restart filename:',restfile
    call ncd_pio_createfile(ncid, trim(restfile))
    call restFile_dimset( ncid )
    call RvicRestart( ncid, flag='define' )
    call RvicHistRestart ( ncid, flag='define', rdate=rdate )
    call timemgr_restart( ncid, flag='define' )
    call ncd_enddef(ncid)

    ! Write restart file variables
    call RvicRestart( ncid, flag='write' )
    call RvicHistRestart ( ncid, flag='write' )
    call timemgr_restart( ncid, flag='write' )
    call ncd_pio_closefile(ncid)

    if (masterproc) then
       write(iulog,*) 'Successfully wrote local restart file ',trim(restfile)
       write(iulog,'(72a1)') ("-",i=1,60)
       write(iulog,*)
    end if

    ! Write restart pointer file
    call restFile_write_pfile( restfile )
    
    ! Write out diagnostic info

    if (masterproc) then
       write(iulog,*) 'Successfully wrote out restart data at nstep = ',get_nstep()
       write(iulog,'(72a1)') ("-",i=1,60)
    end if
    
  end subroutine RvicRestFileWrite

!-----------------------------------------------------------------------

  subroutine restFile_dimset( ncid )

    !----------------------------------------------------------------
    ! !DESCRIPTION:
    ! Read/Write initial data from/to netCDF instantaneous initial data file

    ! !ARGUMENTS:
    implicit none
    type(file_desc_t), intent(inout) :: ncid

    ! !LOCAL VARIABLES:
    integer :: dimid               ! netCDF dimension id
    integer :: ier                 ! error status
    character(len=  8) :: curdate  ! current date
    character(len=  8) :: curtime  ! current time
    character(len=256) :: str
    character(len= 32) :: subname='restFile_dimset' ! subroutine name
    !----------------------------------------------------------------

    ! Define dimensions
    
    call ncd_defdim(ncid, 'time', 1, dimid)
    call ncd_defdim(ncid, 'timesteps', Rvar%full_time_length, dimid)
    call ncd_defdim(ncid, 'outlets', Rvar%nOutlets, dimid)
    call ncd_defdim(ncid, 'tracer', nt_rvic, dimid)
    !call ncd_defdim(ncid, 'string_length', 64        , dimid)
       
    ! Define global attributes
    
    call ncd_putatt(ncid, NCD_GLOBAL, 'Conventions', trim(conventions))
    call getdatetime(curdate, curtime)
    str = 'created on ' // curdate // ' ' // curtime
    call ncd_putatt(ncid, NCD_GLOBAL, 'history' , trim(str))
    call ncd_putatt(ncid, NCD_GLOBAL, 'username', trim(username))
    call ncd_putatt(ncid, NCD_GLOBAL, 'host'    , trim(hostname))
    call ncd_putatt(ncid, NCD_GLOBAL, 'version' , trim(version))
    call ncd_putatt(ncid, NCD_GLOBAL, 'source'  , trim(source))
    str = 'RvicRestFile Mod - jhamman'
    call ncd_putatt(ncid, NCD_GLOBAL, 'revision_id'    , trim(str))
    call ncd_putatt(ncid, NCD_GLOBAL, 'case_title'     , trim(ctitle))
    call ncd_putatt(ncid, NCD_GLOBAL, 'case_id'        , trim(caseid))
    call ncd_putatt(ncid, NCD_GLOBAL, 'title', &
          'RVIC Restart information, required to continue a simulation' )

  end subroutine restFile_dimset
  
!-----------------------------------------------------------------------

  subroutine RvicRestart(ncid, flag)

    !-----------------------------------------------------------------------
    ! DESCRIPTION:
    ! Read/write RVIC restart data.
    !
    ! ARGUMENTS:
    implicit none
    type(file_desc_t), intent(inout)  :: ncid ! netcdf id
    character(len=*) , intent(in) :: flag   ! 'read' or 'write'
    ! LOCAL VARIABLES:
    logical :: readvar          ! determine if variable is on initial file
    integer :: nt,nv,n,i,j      ! indices
    integer :: xtype
    real(r8) , pointer :: drfld(:,:) ! temporary 2d real array
    integer , pointer :: difld(:) ! temporary array
    character(len=32)  :: vname,uname, dim1name, dim2name, dim3name
    character(len=128) :: lname
    character(len=*),parameter :: subname='RvicRestart'
    !-----------------------------------------------------------------------
    
    do nv = 1,5
       if (nv == 1) then
          !vname = 'ring'
          lname = 'Convolution Ring'
          uname = 'kg m-2 s-1'
          xtype = ncd_double
          !drfld  => rvar%ring(:,:,1)
       elseif (nv == 2) then
          vname = 'outlet_x_ind'
          lname = 'x grid coordinate of outlet grid cell'
          uname = 'unitless'
          xtype = ncd_int
          difld  => rvar%outlet_x_ind
       elseif (nv == 3) then
          vname = 'outlet_y_ind'
          lname = 'y grid coordinate of outlet grid cell'
          uname = 'unitless'
          xtype = ncd_int
          difld  => rvar%outlet_y_ind
       elseif (nv == 4) then
          vname = 'outlet_decomp_ind'
          lname = '1d grid location of outlet grid cell'
          uname = 'unitless'
          xtype = ncd_int
          difld  => rvar%outlet_decomp_ind
       elseif (nv == 5) then
          vname = 'ring_timesteps'
          lname = 'Series of timesteps'
          uname = 'unitless'
          xtype = ncd_int
          difld  => rvar%ring_timesteps
       else
          write(iulog,*) 'Rvic ERROR: illegal nv value a ',nv
          call shr_sys_abort()
       endif

       if (flag == 'define') then
         if (nv == 1) then
           do nt = 1, nt_rvic
             vname = rvic_tracers(nt)//'_ring'
             call ncd_defvar(ncid=ncid, varname=trim(vname), &
                  xtype=xtype,  dim1name='outlets', dim2name='timesteps', &
                  long_name=trim(lname), units=trim(uname))
           end do
         elseif (nv == 5) then
           call ncd_defvar(ncid=ncid, varname=trim(vname), &
                xtype=xtype,  dim1name='timesteps', &
                long_name=trim(lname), units=trim(uname))
         else
           call ncd_defvar(ncid=ncid, varname=trim(vname), &
                xtype=xtype,  dim1name='outlets', &
                long_name=trim(lname), units=trim(uname))
         endif
        else if (flag == 'read' .or. flag == 'write') then
          if (nv == 1) then
           allocate(drfld(Rvar%nOutlets,Rvar%full_time_length))
           do nt = 1, nt_rvic
             vname = rvic_tracers(nt)//'_ring'
             if (flag == 'write') then
                do i = 1,Rvar%full_time_length
                do j = 1,Rvar%nOutlets
                   drfld(j,i) = rvar%ring(i,j,nt)
                enddo
                enddo
             endif
             call ncd_io(varname=trim(vname), data=drfld, &
                 ncid=ncid, flag=flag, readvar=readvar)
             if (flag == 'read') then
                do i = 1,Rvar%full_time_length
                do j = 1,Rvar%nOutlets
                   rvar%ring(i,j,nt) = drfld(j,i)
                enddo
                enddo
             endif
           end do
           deallocate(drfld)
          elseif (nv == 5) then
            call ncd_io(varname=trim(vname), data=difld, dim1name='timesteps', &
                 ncid=ncid, flag=flag, readvar=readvar)
          else
            call ncd_io(varname=trim(vname), data=difld, dim1name='outlets', &
                 ncid=ncid, flag=flag, readvar=readvar)
          end if
          if (flag=='read' .and. .not. readvar) then
             if (nsrest == nsrContinue) then
                call shr_sys_abort()
             else
                drfld = 0._r8
             end if
          end if
       end if
    end do
  end subroutine RvicRestart
   
!-----------------------------------------------------------------------

  subroutine restFile_read_pfile( pnamer )

        ! !DESCRIPTION:
        ! Setup restart file and perform necessary consistency checks

        ! !ARGUMENTS:
        implicit none
        character(len=*), intent(out) :: pnamer ! full path of restart file

    ! !LOCAL VARIABLES:
    integer :: i                  ! indices
    integer :: nio                ! restart unit
    integer :: status             ! substring check status
    character(len=256) :: locfn   ! Restart pointer file name
    !--------------------------------------------------------

    ! Obtain the restart file from the restart pointer file. 
    ! For restart runs, the restart pointer file contains the full pathname 
    ! of the restart file. For branch runs, the namelist variable 
    ! [nrevsn_rvic] contains the full pathname of the restart file. 
    ! New history files are always created for branch runs.

    if (masterproc) then
       write(iulog,*) 'Reading restart pointer file....'
    endif

    nio = getavu()
    locfn = './'// trim(rpntfil)//trim(inst_suffix)
    call opnfil (locfn, nio, 'f')
    read (nio,'(a256)') pnamer
    call relavu (nio)

    if (masterproc) then
       write(iulog,*) 'Reading restart data.....'
       write(iulog,'(72a1)') ("-",i=1,60)
    end if

  end subroutine restFile_read_pfile

!-----------------------------------------------------------------------

!-----------------------------------------------------------------------

  subroutine RvicRestGetfile( file, path )

    !---------------------------------------------------
    ! DESCRIPTION:
    ! Determine and obtain netcdf restart file

    ! ARGUMENTS:
    implicit none
    character(len=*), intent(out) :: file  ! name of netcdf restart file
    character(len=*), intent(out) :: path  ! full pathname of netcdf restart file

    ! LOCAL VARIABLES:
    integer :: status                      ! return status
    integer :: length                      ! temporary          
    character(len=256) :: ftest,ctest      ! temporaries
    !---------------------------------------------------

    ! Continue run:
    ! Restart file pathname is read restart pointer file 
    if (nsrest==nsrContinue) then
       call restFile_read_pfile( path )
       call getfil( path, file, 0 )
    end if

    ! Branch run: 
    ! Restart file pathname is obtained from namelist "nrevsn_rvic"
    if (nsrest==nsrBranch) then
       length = len_trim(nrevsn_rvic)
       if (nrevsn_rvic(length-2:length) == '.nc') then
          path = trim(nrevsn_rvic)
       else
          path = trim(nrevsn_rvic) // '.nc'
       end if
       call getfil( path, file, 0 )

       ! Check case name consistency (case name must be different 
       ! for branch run, unless brnch_retain_casename is set)
       ctest = 'xx.'//trim(caseid)//'.rvic'
       ftest = 'xx.'//trim(file)
       status = index(trim(ftest),trim(ctest))
       if (status /= 0 .and. .not.(brnch_retain_casename)) then
          write(iulog,*) 'Must change case name on branch run if ',&
               'brnch_retain_casename namelist is not set'
          write(iulog,*) 'previous case filename= ',trim(file),&
               ' current case = ',trim(caseid), ' ctest = ',trim(ctest), &
               ' ftest = ',trim(ftest)
          call shr_sys_abort()
       end if
    end if

    ! Initial run 
    if (nsrest==nsrStartup) then
       call getfil( path , file, 0 )
    end if

  end subroutine RvicRestGetfile

!-----------------------------------------------------------------------

!-----------------------------------------------------------------------

  subroutine restFile_write_pfile( fnamer )

    ! !DESCRIPTION:
    ! Open restart pointer file. Write names of current netcdf restart file.
    !
    ! !ARGUMENTS:
    implicit none
    character(len=*), intent(in) :: fnamer
    !
    ! !LOCAL VARIABLES:
    integer :: m                    ! index
    integer :: nio                  ! restart pointer file
    character(len=256) :: filename  ! local file name

    if (masterproc) then
       nio = getavu()
       filename= './'// trim(rpntfil)//trim(inst_suffix)
       call opnfil( filename, nio, 'f' )

       write(nio,'(a)') fnamer
       call relavu( nio )
       write(iulog,*)'Successfully wrote local restart pointer file'
    end if

  end subroutine restFile_write_pfile


!-----------------------------------------------------------------------

  character(len=256) function RvicRestFileName( rdate )

    implicit none
    character(len=*), intent(in) :: rdate   ! input date for restart file name 

    RvicRestFileName =trim(caseid)//".rvic"//trim(inst_suffix)//".r."//trim(rdate)//".nc"
    if (masterproc) then
       write(iulog,*)'writing restart file ',trim(RvicRestFileName),' for model date = ',rdate
    end if

  end function RvicRestFileName
!------------------------------------------------------------------------

  subroutine RvicRestTimeManager( file )

    ! !DESCRIPTION:
    ! Read a RTM restart file.
    !
    ! !ARGUMENTS:
    implicit none
    character(len=*), intent(in) :: file  ! output netcdf restart file
    !
    ! !LOCAL VARIABLES:
    type(file_desc_t) :: ncid ! netcdf id
    integer :: i              ! index
    !-------------------------------------

    ! Read file
    if (masterproc) write(iulog,*) 'Reading restart Timemanger'
    call ncd_pio_openfile (ncid, trim(file), 0)
    call timemgr_restart(ncid, flag='read')
    call ncd_pio_closefile(ncid)

    ! Write out diagnostic info
    if (masterproc) then
       write(iulog,'(72a1)') ("-",i=1,60)
       write(iulog,*) 'Successfully read restart data for restart run'
       write(iulog,*)
    end if

  end subroutine RvicRestTimeManager

!------------------------------------------------------------------------

  subroutine RvicRestFileRead( file )

    ! !DESCRIPTION:
    ! Read a RTM restart file.
    !
    ! !ARGUMENTS:
    implicit none
    character(len=*), intent(in) :: file  ! output netcdf restart file
    !
    ! !LOCAL VARIABLES:
    type(file_desc_t) :: ncid ! netcdf id
    integer :: i              ! index
    !-------------------------------------

    ! Read file
    if (masterproc) write(iulog,*) 'Reading restart dataset'
    call ncd_pio_openfile (ncid, trim(file), 0)
    call RvicRestart( ncid, flag='read' )
    call RvicHistRestart(ncid, flag='read')
    call ncd_pio_closefile(ncid)

    ! Write out diagnostic info
    if (masterproc) then
       write(iulog,'(72a1)') ("-",i=1,60)
       write(iulog,*) 'Successfully read restart data for restart run'
       write(iulog,*)
    end if

  end subroutine RvicRestFileRead


  subroutine handle_err(status)

    implicit none
    
    integer, intent ( in) :: status
    
    if(status /= nf90_noerr) then
      write(iulog,*) nf90_strerror(status)
      write(iulog, *) "Stopped due to error status" 
      call shr_sys_abort()
    end if
  end subroutine handle_err 

end module RvicRestFile