module restart_physics

  use shr_kind_mod,       only: r8 => shr_kind_r8
  use spmd_utils,         only: masterproc
  use co2_cycle,          only: co2_transport
  use constituents,       only: pcnst

  use radiation,          only: radiation_define_restart, radiation_write_restart, &
                                radiation_read_restart

  use ioFileMod
  use cam_abortutils,     only: endrun
  use camsrfexch,         only: cam_in_t, cam_out_t
  use cam_control_mod,    only: adiabatic, ideal_phys
  use cam_logfile,        only: iulog
  use pio,                only: file_desc_t, io_desc_t, var_desc_t, &
                                pio_double, pio_int, pio_noerr, &
                                pio_seterrorhandling, pio_bcast_error, &
                                pio_inq_varid, &
                                pio_def_var, pio_def_dim, &
                                pio_put_var, pio_get_var

  implicit none
  private
  save
!
! Public interfaces
!
  public :: write_restart_physics    ! Write the physics restart info out
  public :: read_restart_physics     ! Read the physics restart info in
  public :: init_restart_physics

!
! Private data
!

    type(var_desc_t) :: flwds_desc, &
         solld_desc, co2prog_desc, co2diag_desc, sols_desc, soll_desc, &
         solsd_desc

    type(var_desc_t) :: bcphidry_desc, bcphodry_desc, ocphidry_desc, ocphodry_desc, &
       dstdry1_desc, dstdry2_desc, dstdry3_desc, dstdry4_desc

    type(var_desc_t) :: cflx_desc(pcnst)

    type(var_desc_t) :: wsx_desc
    type(var_desc_t) :: wsy_desc
    type(var_desc_t) :: shf_desc

  CONTAINS
    subroutine init_restart_physics ( File, pbuf2d)
      
    use physics_buffer,      only: pbuf_init_restart, physics_buffer_desc
    use ppgrid,              only: pver, pverp
    use chemistry,           only: chem_init_restart
    use prescribed_ozone,    only: init_prescribed_ozone_restart
    use prescribed_ghg,      only: init_prescribed_ghg_restart
    use prescribed_aero,     only: init_prescribed_aero_restart
    use prescribed_volcaero, only: init_prescribed_volcaero_restart
    use cam_grid_support,    only: cam_grid_write_attr, cam_grid_id
    use cam_grid_support,    only: cam_grid_header_info_t
    use cam_pio_utils,       only: cam_pio_def_dim
    use subcol_utils,        only: is_subcol_on
    use subcol,              only: subcol_init_restart

    type(file_desc_t), intent(inout) :: file
    type(physics_buffer_desc), pointer :: pbuf2d(:,:)

    integer                      :: grid_id
    integer                      :: hdimcnt, ierr, i
    integer                      :: dimids(4)
    integer, allocatable         :: hdimids(:)
    type(cam_grid_header_info_t) :: info
    character(len=4)   :: num

    call pio_seterrorhandling(File, PIO_BCAST_ERROR)
    ! Probably should have the grid write this out.
    grid_id = cam_grid_id('physgrid')
    call cam_grid_write_attr(File, grid_id, info)
    hdimcnt = info%num_hdims()

    do i = 1, hdimcnt
      dimids(i) = info%get_hdimid(i)
    end do
    allocate(hdimids(hdimcnt))
    hdimids(1:hdimcnt) = dimids(1:hdimcnt)

    call pbuf_init_restart(File, pbuf2d)

    if ( .not. adiabatic .and. .not. ideal_phys )then
       
       call chem_init_restart(File)

       call init_prescribed_ozone_restart(File)
       call init_prescribed_ghg_restart(File)
       call init_prescribed_aero_restart(File)
       call init_prescribed_volcaero_restart(File)

       ierr = pio_def_var(File, 'FLWDS', pio_double, hdimids, flwds_desc)
       ierr = pio_def_var(File, 'SOLS', pio_double, hdimids, sols_desc)
       ierr = pio_def_var(File, 'SOLL', pio_double, hdimids, soll_desc)
       ierr = pio_def_var(File, 'SOLSD', pio_double, hdimids, solsd_desc)
       ierr = pio_def_var(File, 'SOLLD', pio_double, hdimids, solld_desc)

       ierr = pio_def_var(File, 'BCPHIDRY', pio_double, hdimids, bcphidry_desc)
       ierr = pio_def_var(File, 'BCPHODRY', pio_double, hdimids, bcphodry_desc)
       ierr = pio_def_var(File, 'OCPHIDRY', pio_double, hdimids, ocphidry_desc)
       ierr = pio_def_var(File, 'OCPHODRY', pio_double, hdimids, ocphodry_desc)
       ierr = pio_def_var(File, 'DSTDRY1',  pio_double, hdimids, dstdry1_desc)
       ierr = pio_def_var(File, 'DSTDRY2',  pio_double, hdimids, dstdry2_desc)
       ierr = pio_def_var(File, 'DSTDRY3',  pio_double, hdimids, dstdry3_desc)
       ierr = pio_def_var(File, 'DSTDRY4',  pio_double, hdimids, dstdry4_desc)

       if(co2_transport()) then
          ierr = pio_def_var(File, 'CO2PROG', pio_double, hdimids, co2prog_desc)
          ierr = pio_def_var(File, 'CO2DIAG', pio_double, hdimids, co2diag_desc)
       end if

       ! cam_import variables -- write the constituent surface fluxes as individual 2D arrays
       ! rather than as a single variable with a pcnst dimension.  Note that the cflx components
       ! are only needed for those constituents that are not passed to the coupler.  The restart
       ! for constituents passed through the coupler are handled by the .rs. restart file.  But
       ! we don't currently have a mechanism to know whether the constituent is handled by the
       ! coupler or not, so we write all of cflx to the CAM restart file.
       do i = 1, pcnst
          write(num,'(i4.4)') i
          ierr = pio_def_var(File, 'CFLX'//num,  pio_double, hdimids, cflx_desc(i))
       end do

       ierr = pio_def_var(File, 'wsx',  pio_double, hdimids, wsx_desc)
       ierr = pio_def_var(File, 'wsy',  pio_double, hdimids, wsy_desc)
       ierr = pio_def_var(File, 'shf',  pio_double, hdimids, shf_desc)
 
       call radiation_define_restart(file)

    end if


    if (is_subcol_on()) then
      call subcol_init_restart(file, hdimids)
    end if
      
  end subroutine init_restart_physics

  subroutine write_restart_physics (File, cam_in, cam_out, pbuf2d)

      !-----------------------------------------------------------------------
      use physics_buffer,      only: physics_buffer_desc, pbuf_write_restart
      use phys_grid,           only: phys_decomp
      
      use ppgrid,              only: begchunk, endchunk, pcols, pverp
      use chemistry,           only: chem_write_restart
      use prescribed_ozone,    only: write_prescribed_ozone_restart
      use prescribed_ghg,      only: write_prescribed_ghg_restart
      use prescribed_aero,     only: write_prescribed_aero_restart
      use prescribed_volcaero, only: write_prescribed_volcaero_restart

      use cam_history_support, only: fillvalue
      use spmd_utils,          only: iam
      use cam_grid_support,    only: cam_grid_write_dist_array, cam_grid_id
      use cam_grid_support,    only: cam_grid_get_decomp, cam_grid_dimensions
      use cam_grid_support,    only: cam_grid_write_var
      use pio,                 only: pio_write_darray
      use subcol_utils,        only: is_subcol_on
      use subcol,              only: subcol_write_restart
      !
      ! Input arguments
      !
      type(file_desc_t), intent(inout) :: File
      type(cam_in_t),    intent(in)    :: cam_in(begchunk:endchunk)
      type(cam_out_t),   intent(in)    :: cam_out(begchunk:endchunk)
      type(physics_buffer_desc), pointer        :: pbuf2d(:,:)
      !
      ! Local workspace
      !
      type(io_desc_t), pointer :: iodesc
      real(r8):: tmpfield(pcols, begchunk:endchunk)
      integer :: i, m          ! loop index
      integer :: ncol          ! number of vertical columns
      integer :: ierr
      integer :: physgrid
      integer :: dims(3), gdims(3)
      integer :: nhdims
      !-----------------------------------------------------------------------

      ! Write grid vars
      call cam_grid_write_var(File, phys_decomp)

      ! Physics buffer
      if (is_subcol_on()) then
         call subcol_write_restart(File)
      end if

      call pbuf_write_restart(File, pbuf2d)

      physgrid = cam_grid_id('physgrid')
      call cam_grid_dimensions(physgrid, gdims(1:2), nhdims)

      if ( .not. adiabatic .and. .not. ideal_phys )then

         ! data for chemistry
         call chem_write_restart(File)

         call write_prescribed_ozone_restart(File)
         call write_prescribed_ghg_restart(File)
         call write_prescribed_aero_restart(File)
         call write_prescribed_volcaero_restart(File)

         ! cam_in/out variables
         ! This is a group of surface variables so can reuse dims
         dims(1) = pcols
         dims(2) = endchunk - begchunk + 1
         call cam_grid_get_decomp(physgrid, dims(1:2), gdims(1:nhdims), &
              pio_double, iodesc)

         do i = begchunk, endchunk
            ncol = cam_out(i)%ncol
            tmpfield(:ncol, i) = cam_out(i)%flwds(:ncol)
            ! Only have to do this once (cam_in/out vars all same shape)
            if (ncol < pcols) then
               tmpfield(ncol+1:, i) = fillvalue
            end if
         end do
         call pio_write_darray(File, flwds_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_out(i)%ncol
            tmpfield(:ncol, i) = cam_out(i)%sols(:ncol)
         end do
         call pio_write_darray(File, sols_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_out(i)%ncol
            tmpfield(:ncol, i) = cam_out(i)%soll(:ncol)
         end do
         call pio_write_darray(File, soll_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_out(i)%ncol
            tmpfield(:ncol, i) = cam_out(i)%solsd(:ncol)
         end do
         call pio_write_darray(File, solsd_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_out(i)%ncol
            tmpfield(:ncol, i) = cam_out(i)%solld(:ncol)
         end do
         call pio_write_darray(File, solld_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_out(i)%ncol
            tmpfield(:ncol, i) = cam_out(i)%bcphidry(:ncol)
         end do
         call pio_write_darray(File, bcphidry_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_out(i)%ncol
            tmpfield(:ncol, i) = cam_out(i)%bcphodry(:ncol)
         end do
         call pio_write_darray(File, bcphodry_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_out(i)%ncol
            tmpfield(:ncol, i) = cam_out(i)%ocphidry(:ncol)
         end do
         call pio_write_darray(File, ocphidry_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_out(i)%ncol
            tmpfield(:ncol, i) = cam_out(i)%ocphodry(:ncol)
         end do
         call pio_write_darray(File, ocphodry_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_out(i)%ncol
            tmpfield(:ncol, i) = cam_out(i)%dstdry1(:ncol)
         end do
         call pio_write_darray(File, dstdry1_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_out(i)%ncol
            tmpfield(:ncol, i) = cam_out(i)%dstdry2(:ncol)
         end do
         call pio_write_darray(File, dstdry2_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_out(i)%ncol
            tmpfield(:ncol, i) = cam_out(i)%dstdry3(:ncol)
         end do
         call pio_write_darray(File, dstdry3_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_out(i)%ncol
            tmpfield(:ncol, i) = cam_out(i)%dstdry4(:ncol)
         end do
         call pio_write_darray(File, dstdry4_desc, iodesc, tmpfield, ierr)

         if (co2_transport()) then
            do i = begchunk, endchunk
               ncol = cam_out(i)%ncol
               tmpfield(:ncol, i) = cam_out(i)%co2prog(:ncol)
            end do
            call pio_write_darray(File, co2prog_desc, iodesc, tmpfield, ierr)

            do i = begchunk, endchunk
               ncol = cam_out(i)%ncol
               tmpfield(:ncol, i) = cam_out(i)%co2diag(:ncol)
            end do
            call pio_write_darray(File, co2diag_desc, iodesc, tmpfield, ierr)
         end if

         ! cam_in components
         do m = 1, pcnst
            do i = begchunk, endchunk
               ncol = cam_in(i)%ncol
               tmpfield(:ncol, i) = cam_in(i)%cflx(:ncol, m)
            end do
            call pio_write_darray(File, cflx_desc(m), iodesc, tmpfield, ierr)
         end do

         do i = begchunk, endchunk
            ncol = cam_in(i)%ncol
            tmpfield(:ncol,i) = cam_in(i)%wsx(:ncol)
         end do
         call pio_write_darray(File, wsx_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_in(i)%ncol
            tmpfield(:ncol,i) = cam_in(i)%wsy(:ncol)
         end do
         call pio_write_darray(File, wsy_desc, iodesc, tmpfield, ierr)

         do i = begchunk, endchunk
            ncol = cam_in(i)%ncol
            tmpfield(:ncol,i) = cam_in(i)%shf(:ncol)
         end do
         call pio_write_darray(File, shf_desc, iodesc, tmpfield, ierr)

         call radiation_write_restart(file)

      end if

      
    end subroutine write_restart_physics

!#######################################################################

    subroutine read_restart_physics(File, cam_in, cam_out, pbuf2d)

     !-----------------------------------------------------------------------
     use physics_buffer,      only: physics_buffer_desc, pbuf_read_restart
     
     use ppgrid,              only: begchunk, endchunk, pcols, pver, pverp
     use chemistry,           only: chem_read_restart
     use cam_grid_support,    only: cam_grid_read_dist_array, cam_grid_id
     use cam_grid_support,    only: cam_grid_get_decomp, cam_grid_dimensions
     use cam_history_support, only: fillvalue

     use prescribed_ozone,    only: read_prescribed_ozone_restart
     use prescribed_ghg,      only: read_prescribed_ghg_restart
     use prescribed_aero,     only: read_prescribed_aero_restart
     use prescribed_volcaero, only: read_prescribed_volcaero_restart
     use subcol_utils,        only: is_subcol_on
     use subcol,              only: subcol_read_restart
     use pio,                 only: pio_read_darray
     !
     ! Arguments
     !
     type(file_desc_t),   intent(inout) :: File
     type(cam_in_t),            pointer :: cam_in(:)
     type(cam_out_t),           pointer :: cam_out(:)
     type(physics_buffer_desc), pointer :: pbuf2d(:,:)
     !
     ! Local workspace
     !
     real(r8), allocatable :: tmpfield2(:,:)
     integer :: i, c, m           ! loop index
     integer :: ierr             ! I/O status
     type(io_desc_t), pointer :: iodesc
     type(var_desc_t)         :: vardesc
     integer                  :: csize, vsize
     character(len=4)         :: num
     integer                  :: dims(3), gdims(3), nhdims
     integer                  :: err_handling
     integer                  :: physgrid
     !-----------------------------------------------------------------------

     ! subcol_read_restart must be called before pbuf_read_restart
     if (is_subcol_on()) then
        call subcol_read_restart(File)
     end if

     call pbuf_read_restart(File, pbuf2d)

     csize=endchunk-begchunk+1
     dims(1) = pcols
     dims(2) = csize

     physgrid = cam_grid_id('physgrid')

     call cam_grid_dimensions(physgrid, gdims(1:2))

     if (gdims(2) == 1) then
       nhdims = 1
     else
       nhdims = 2
     end if
     call cam_grid_get_decomp(physgrid, dims(1:2), gdims(1:nhdims), pio_double, &
          iodesc)
     
     camphys: if ( .not. adiabatic .and. .not. ideal_phys )then

        ! data for chemistry
        call chem_read_restart(File)

        call read_prescribed_ozone_restart(File)
        call read_prescribed_ghg_restart(File)
        call read_prescribed_aero_restart(File)
        call read_prescribed_volcaero_restart(File)

        allocate(tmpfield2(pcols, begchunk:endchunk))
        tmpfield2 = fillvalue

        ierr = pio_inq_varid(File, 'FLWDS', vardesc)
        call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
        do c=begchunk,endchunk
           do i=1,pcols
              cam_out(c)%flwds(i) = tmpfield2(i, c)
           end do
        end do

        ierr = pio_inq_varid(File, 'SOLS', vardesc)
        call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
        do c=begchunk,endchunk
           do i=1,pcols
              cam_out(c)%sols(i) = tmpfield2(i, c)
           end do
        end do

        ierr = pio_inq_varid(File, 'SOLL', vardesc)
        call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
        do c=begchunk,endchunk
           do i=1,pcols
              cam_out(c)%soll(i) = tmpfield2(i, c)
           end do
        end do

        ierr = pio_inq_varid(File, 'SOLSD', vardesc)
        call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
        do c=begchunk,endchunk
           do i=1,pcols
              cam_out(c)%solsd(i) = tmpfield2(i, c)
           end do
        end do

        ierr = pio_inq_varid(File, 'SOLLD', vardesc)
        call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
        do c=begchunk,endchunk
           do i=1,pcols
              cam_out(c)%solld(i) = tmpfield2(i, c)
           end do
        end do

        ierr = pio_inq_varid(File, 'BCPHIDRY', vardesc)
        call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
        do c=begchunk,endchunk
           do i=1,pcols
              cam_out(c)%bcphidry(i) = tmpfield2(i, c)
           end do
        end do

        ierr = pio_inq_varid(File, 'BCPHODRY', vardesc)
        call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
        do c=begchunk,endchunk
           do i=1,pcols
              cam_out(c)%bcphodry(i) = tmpfield2(i, c)
           end do
        end do

        ierr = pio_inq_varid(File, 'OCPHIDRY', vardesc)
        call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
        do c=begchunk,endchunk
           do i=1,pcols
              cam_out(c)%ocphidry(i) = tmpfield2(i, c)
           end do
        end do

        ierr = pio_inq_varid(File, 'OCPHODRY', vardesc)
        call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
        do c=begchunk,endchunk
           do i=1,pcols
              cam_out(c)%ocphodry(i) = tmpfield2(i, c)
           end do
        end do

        ierr = pio_inq_varid(File, 'DSTDRY1', vardesc)
        call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
        do c=begchunk,endchunk
           do i=1,pcols
              cam_out(c)%dstdry1(i) = tmpfield2(i, c)
           end do
        end do

        ierr = pio_inq_varid(File, 'DSTDRY2', vardesc)
        call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
        do c=begchunk,endchunk
           do i=1,pcols
              cam_out(c)%dstdry2(i) = tmpfield2(i, c)
           end do
        end do

        ierr = pio_inq_varid(File, 'DSTDRY3', vardesc)
        call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
        do c=begchunk,endchunk
           do i=1,pcols
              cam_out(c)%dstdry3(i) = tmpfield2(i, c)
           end do
        end do

        ierr = pio_inq_varid(File, 'DSTDRY4', vardesc)
        call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
        do c=begchunk,endchunk
           do i=1,pcols
              cam_out(c)%dstdry4(i) = tmpfield2(i, c)
           end do
        end do

        if (co2_transport()) then
           ierr = pio_inq_varid(File, 'CO2PROG', vardesc)
           call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
           do c=begchunk,endchunk
              do i=1,pcols
                 cam_out(c)%co2prog(i) = tmpfield2(i, c)
              end do
           end do

           ierr = pio_inq_varid(File, 'CO2DIAG', vardesc)
           call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
           do c=begchunk,endchunk
              do i=1,pcols
                 cam_out(c)%co2diag(i) = tmpfield2(i, c)
              end do
           end do
        end if

        ! Reading the CFLX* components from the restart is optional for
        ! backwards compatibility.  These fields were not needed for an
        ! exact restart until the UNICON scheme was added.  More generally,
        ! these components are only needed if they are not handled by the
        ! coupling layer restart (the ".rs." file), and if the values are
        ! used in the tphysbc physics before the tphysac code has a chance
        ! to update the values that are coming from boundary datasets.
        do m = 1, pcnst

           write(num,'(i4.4)') m

           call pio_seterrorhandling(File, PIO_BCAST_ERROR, err_handling)
           ierr = pio_inq_varid(File, 'CFLX'//num, vardesc)
           call pio_seterrorhandling(File, err_handling)

           if (ierr == PIO_NOERR) then ! CFLX variable found on restart file
              call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
              do c= begchunk, endchunk
                 do i = 1, pcols
                    cam_in(c)%cflx(i,m) = tmpfield2(i, c)
                 end do
              end do
           end if

        end do

        call pio_seterrorhandling(File, PIO_BCAST_ERROR, err_handling)
        ierr = pio_inq_varid(File, 'wsx', vardesc)
        if (ierr == PIO_NOERR) then ! variable found on restart file
           call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
           do c= begchunk, endchunk
              do i = 1, pcols
                 cam_in(c)%wsx(i) = tmpfield2(i, c)
              end do
           end do
        endif
        ierr = pio_inq_varid(File, 'wsy', vardesc)
        if (ierr == PIO_NOERR) then ! variable found on restart file
           call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
           do c= begchunk, endchunk
              do i = 1, pcols
                 cam_in(c)%wsy(i) = tmpfield2(i, c)
              end do
           end do
        endif
        ierr = pio_inq_varid(File, 'shf', vardesc)
        if (ierr == PIO_NOERR) then ! variable found on restart file
           call pio_read_darray(File, vardesc, iodesc, tmpfield2, ierr)
           do c= begchunk, endchunk
              do i = 1, pcols
                 cam_in(c)%shf(i) = tmpfield2(i, c)
              end do
           end do
        endif
        call pio_seterrorhandling(File, err_handling)

        deallocate(tmpfield2)

        call radiation_read_restart(file)

     end if camphys


   end subroutine read_restart_physics

 end module restart_physics