program ml_model_prediction
  implicit none
  integer,parameter  :: r8   = selected_real_kind(12) ! 8 byte real
  integer,parameter  :: r4   = selected_real_kind( 6) ! 4 byte real
  integer,parameter  :: cs   = 80                     ! short char
  integer,parameter  :: cm   = 160                    ! mid-sized char
  integer,parameter  :: cl   = 256                    ! long char
  integer,parameter  :: cx   = 512                    ! extra-long char
  integer,parameter  :: cxx  = 4096                   ! extra-extra-long char

  integer,parameter  :: ncol = 21600 ! number of spatial columns
  integer,parameter  :: nlev = 72    ! number of vertical levels 
  integer,parameter  :: nt   = 5     ! number of time steps 

  real(r8), allocatable :: u(:,:,:),v(:,:,:),t(:,:,:),q(:,:,:)
  real(r8), allocatable :: um(:,:),vm(:,:),tm(:,:),qm(:,:)
  real(r8), allocatable :: lat(:), lon(:), lev(:)

  real(r8), allocatable :: utend(:,:,:),vtend(:,:,:),ttend(:,:,:),qtend(:,:,:)

  real :: tstart, tfinish
  character(len=cxx) :: ufile,vfile,tfile,qfile
  character(len=cxx) :: fu,fv,ft,fq
  character(len=cxx) :: fmlout, fnlout, finout
  character(len=cs)  :: vpred, vinam, vonam,ttstr
  character(len=cxx) :: ptpath, ptgroup 
  character(len=cxx) :: utfile, vtfile, ttfile, qtfile
  character          :: xx,yy,zz
  integer            :: m,n,i,j,k,ii,jj
  integer            :: iit !time step for ml input/output  

  !variable to be predicted
  ptpath  = "/pscratch/sd/z/zhan391/SEACROGS_project/e3sm_model/machine_learning/pt_global"
  ptgroup = "UNet_4var/unified_iunet"
  ufile   = trim(ptpath)//'/'//trim(ptgroup)//"/UVTQ_U_global.pt"
  vfile   = trim(ptpath)//'/'//trim(ptgroup)//"/UVTQ_V_global.pt"
  tfile   = trim(ptpath)//'/'//trim(ptgroup)//"/UVTQ_T_global.pt" 
  qfile   = trim(ptpath)//'/'//trim(ptgroup)//"/UVTQ_Q_global.pt" 

  !read model data
  fu = "../input/ml_input_U_2009.nc"
  fv = "../input/ml_input_V_2009.nc"
  ft = "../input/ml_input_T_2009.nc"
  fq = "../input/ml_input_Q_2009.nc"

  allocate(u(ncol,1,nlev),v(ncol,1,nlev),t(ncol,1,nlev),q(ncol,1,nlev))
  allocate(lat(ncol), lon(ncol), lev(nlev))
  allocate(um(ncol,nlev),vm(ncol,nlev),tm(ncol,nlev),qm(ncol,nlev))

  allocate(utend(ncol,1,nlev),vtend(ncol,1,nlev),ttend(ncol,1,nlev),qtend(ncol,1,nlev))

  do iit = 1, nt
    write(*,*) "extract data at ",iit,"th step:"
    call read_model_data(trim(fu),'U',u(:,1,:),ncol,nlev,iit,lat,lon,lev)
    call read_model_data(trim(fv),'V',v(:,1,:),ncol,nlev,iit,lat,lon,lev)
    call read_model_data(trim(ft),'T',t(:,1,:),ncol,nlev,iit,lat,lon,lev)
    call read_model_data(trim(fq),'Q',q(:,1,:),ncol,nlev,iit,lat,lon,lev)
    write(*,*) 'read model data, u(min/max) = ', minval(u), maxval(u), shape(u)
    write(*,*) 'read model data, v(min/max) = ', minval(v), maxval(v), shape(v)
    write(*,*) 'read model data, t(min/max) = ', minval(t), maxval(t), shape(t)
    write(*,*) 'read model data, q(min/max) = ', minval(q), maxval(q), shape(q)
    write(*,*) 'read model data, lat(min/max) = ', minval(lat), maxval(lat), shape(lat)
    write(*,*) 'read model data, lon(min/max) = ', minval(lon), maxval(lon), shape(lon)

    !output files 
    write(ttstr,"(I0.4)") iit
    finout = "./output/ML_INPUTE_STATE_U.2009-"//trim(ttstr)//".nc"
    call out_model_data(trim(finout),'U',u(:,1,:),ncol,nlev,1,lat,lon,lev,.false.)
    finout = "./output/ML_INPUTE_STATE_V.2009-"//trim(ttstr)//".nc"
    call out_model_data(trim(finout),'V',v(:,1,:),ncol,nlev,1,lat,lon,lev,.false.)
    finout = "./output/ML_INPUTE_STATE_T.2009-"//trim(ttstr)//".nc"
    call out_model_data(trim(finout),'T',t(:,1,:),ncol,nlev,1,lat,lon,lev,.false.)
    finout = "./output/ML_INPUTE_STATE_Q.2009-"//trim(ttstr)//".nc"
    call out_model_data(trim(finout),'Q',q(:,1,:),ncol,nlev,1,lat,lon,lev,.false.)

    !ML model prediction 
    call cpu_time(tstart)
    call advance_ml(ncol,1,nlev,u,v,t,q,lat,lon,ufile,vfile,tfile,qfile,utend,vtend,ttend,qtend)
    call cpu_time(tfinish)
    print '("total run time = ",f6.3," seconds.")',tfinish - tstart
    fmlout = "./output/ML_PREDICT_Nudge_U.2009-"//trim(ttstr)//".nc"
    call out_model_data(trim(fmlout),'Nudge_U',utend(:,1,:),ncol,nlev,1,lat,lon,lev,.false.)
    fmlout = "./output/ML_PREDICT_Nudge_V.2009-"//trim(ttstr)//".nc"
    call out_model_data(trim(fmlout),'Nudge_V',vtend(:,1,:),ncol,nlev,1,lat,lon,lev,.false.)
    fmlout = "./output/ML_PREDICT_Nudge_T.2009-"//trim(ttstr)//".nc"
    call out_model_data(trim(fmlout),'Nudge_T',ttend(:,1,:),ncol,nlev,1,lat,lon,lev,.false.)
    fmlout = "./output/ML_PREDICT_Nudge_Q.2009-"//trim(ttstr)//".nc"
    call out_model_data(trim(fmlout),'Nudge_Q',qtend(:,1,:),ncol,nlev,1,lat,lon,lev,.false.)

  ! derive state after appliying nudging 
    print*,shape(utend)
    do k = 1,nlev
      do i = 1,ncol
        um(i,k) = utend(i,1,k) * 1800.0  + u(i,1,k) 
        vm(i,k) = vtend(i,1,k) * 1800.0  + v(i,1,k)
        tm(i,k) = ttend(i,1,k) * 1800.0  + t(i,1,k)
        qm(i,k) = qtend(i,1,k) * 1800.0  + q(i,1,k)
      end do
    end do
    fnlout = "./output/ML_PREDICT_STATE_U.2009-"//trim(ttstr)//".nc"
    call out_model_data(trim(fnlout),'U',um(:,:),ncol,nlev,1,lat,lon,lev,.false.)
    fnlout = "./output/ML_PREDICT_STATE_V.2009-"//trim(ttstr)//".nc"
    call out_model_data(trim(fnlout),'V',vm(:,:),ncol,nlev,1,lat,lon,lev,.false.)
    fnlout = "./output/ML_PREDICT_STATE_T.2009-"//trim(ttstr)//".nc"
    call out_model_data(trim(fnlout),'T',tm(:,:),ncol,nlev,1,lat,lon,lev,.false.)
    fnlout = "./output/ML_PREDICT_STATE_Q.2009-"//trim(ttstr)//".nc"
    call out_model_data(trim(fnlout),'Q',qm(:,:),ncol,nlev,1,lat,lon,lev,.false.)
  end do 

  deallocate(u,v,t,q,lat,lon,lev)
  deallocate(um,vm,tm,qm,utend,vtend,ttend,qtend)
end program ml_model_prediction

subroutine advance_ml(nx,ny,nz,u,v,t,q,lat,lon,ufile,vfile,tfile,qfile, &
                      utend,vtend,ttend,qtend)
  use torch_ftn
  
  implicit none
  integer,parameter  :: r8   = selected_real_kind(12) ! 8 byte real
  integer,parameter  :: r4   = selected_real_kind( 6) ! 4 byte real

  character(len=*), intent(in)  :: ufile,vfile,tfile,qfile
  integer,          intent(in)  :: nx,ny,nz

  real(r8),         intent(in)  :: u(nx,ny,nz)
  real(r8),         intent(in)  :: v(nx,ny,nz)
  real(r8),         intent(in)  :: t(nx,ny,nz)
  real(r8),         intent(in)  :: q(nx,ny,nz)
  real(r8),         intent(in)  :: lat(nx)
  real(r8),         intent(in)  :: lon(nx)

  real(r8),         intent(out) :: utend(nx,ny,nz)
  real(r8),         intent(out) :: vtend(nx,ny,nz)
  real(r8),         intent(out) :: ttend(nx,ny,nz)
  real(r8),         intent(out) :: qtend(nx,ny,nz)

  ! Arguments
  !-----------
  type(torch_tensor_wrap)      :: input_tensors
  type(torch_module)           :: torch_umod
  type(torch_module)           :: torch_vmod
  type(torch_module)           :: torch_tmod
  type(torch_module)           :: torch_qmod
  type(torch_tensor)           :: u_tensor
  type(torch_tensor)           :: v_tensor
  type(torch_tensor)           :: t_tensor
  type(torch_tensor)           :: q_tensor

  ! Local values
  !----------------
  logical, parameter           :: l_print_diag = .true.
  integer                      :: i,j,k,n,m
  integer                      :: k1,k2
  logical                      :: l_file_exist
  real(r8), pointer            :: uout(:,:,:)
  real(r8), pointer            :: vout(:,:,:)
  real(r8), pointer            :: tout(:,:,:)
  real(r8), pointer            :: qout(:,:,:)

  inquire(file=trim(ufile),exist=l_file_exist)
  if ( .not. l_file_exist) then
    write(*,*) "ERROR: "//trim(ufile)//" not found!"
  end if

  inquire(file=trim(vfile),exist=l_file_exist)
  if ( .not. l_file_exist) then
    write(*,*) "ERROR: "//trim(vfile)//" not found!"
  end if

  inquire(file=trim(tfile),exist=l_file_exist)
  if ( .not. l_file_exist) then
    write(*,*) "ERROR: "//trim(tfile)//" not found!"
  end if

  inquire(file=trim(qfile),exist=l_file_exist)
  if ( .not. l_file_exist) then
    write(*,*) "ERROR: "//trim(qfile)//" not found!"
  end if

  utend(:,:,:) = 0.0_r8
  vtend(:,:,:) = 0.0_r8
  ttend(:,:,:) = 0.0_r8
  qtend(:,:,:) = 0.0_r8

  k1 = 1
  !if (present(kbot)) k1 = kbot
  k2 = nz
  !if (present(ktop)) k2 = ktop

  if (l_print_diag) then
    write(*,*) 'shape of ml input (ncol,nlev,nt) = ',shape(u),shape(v),shape(t),shape(q)
    write(*,*) 'ml input min(u,v,t,q) =',minval(u),minval(v),minval(t),minval(q)
    write(*,*) 'ml input max(u,v,t,q) =',maxval(u),maxval(v),maxval(t),maxval(q)
  end if

  call input_tensors%create
  call input_tensors%add_array(u(:,:,k1:k2))
  call input_tensors%add_array(v(:,:,k1:k2))
  call input_tensors%add_array(t(:,:,k1:k2))
  call input_tensors%add_array(q(:,:,k1:k2))
  !call input_tensors%add_array(lat(:))
  !call input_tensors%add_array(lon(:))

  call torch_umod%load(trim(ufile),0)
  call torch_umod%forward(input_tensors,u_tensor)
  call u_tensor%to_array(uout)
  write(*,*) 'shape of ml output (ncol,nlev,nt) = ',shape(uout)
  
  utend(:,:,k1:k2) = uout(:,:,:)

  call torch_vmod%load(trim(vfile),0)
  call torch_vmod%forward(input_tensors,v_tensor)
  call v_tensor%to_array(vout)
  vtend(:,:,k1:k2) = vout(:,:,:)

  call torch_tmod%load(trim(tfile),0)
  call torch_tmod%forward(input_tensors,t_tensor)
  call t_tensor%to_array(tout)
  ttend(:,:,k1:k2) = tout(:,:,:)

  call torch_qmod%load(trim(qfile),0)
  call torch_qmod%forward(input_tensors,q_tensor)
  call q_tensor%to_array(qout)
  qtend(:,:,k1:k2) = qout(:,:,:)

  !return output data
  !utend(:,:,k1:k2) = uout(:,:,:)
  !vtend(:,:,k1:k2) = vout(:,:,:)
  !ttend(:,:,k1:k2) = tout(:,:,:)
  !qtend(:,:,k1:k2) = qout(:,:,:)

  if (l_print_diag) then
    write(*,*) 'tendput shape,min,max = ',shape(utend),minval(utend),maxval(utend)
    write(*,*) 'tendput shape,min,max = ',shape(vtend),minval(vtend),maxval(vtend)
    write(*,*) 'tendput shape,min,max = ',shape(ttend),minval(ttend),maxval(ttend)
    write(*,*) 'tendput shape,min,max = ',shape(qtend),minval(qtend),maxval(qtend)
  end if

 return
end subroutine advance_ml

!-----------------------
! open a new netcdf file
!-----------------------
subroutine read_model_data (file_name, vname, varout, ncol, nz, it, lat, lon, lev)
  use netcdf            

  implicit none
  integer,parameter :: r8  = selected_real_kind(12) ! 8 byte real
  integer,parameter :: r4  = selected_real_kind( 6) ! 4 byte real
  integer,parameter :: cs  = 80                     ! short char
  integer,parameter :: cm  = 160                    ! mid-sized char
  integer,parameter :: cl  = 256                    ! long char
  integer,parameter :: cx  = 512                    ! extra-long char
  integer,parameter :: cxx = 4096                   ! extra-extra-long char

  real(r8), parameter :: fillvalue = 1.e36_r8  
  
  character(len=*), intent(in)  :: file_name
  character(len=*), intent(in)  :: vname
  integer,          intent(in)  :: ncol, nz, it 
  real(r8),         intent(out) :: varout(ncol,nz)
  real(r8),         intent(out) :: lat(ncol)
  real(r8),         intent(out) :: lon(ncol)
  real(r8),         intent(out) :: lev(nz)

  real(r8), allocatable :: var(:,:,:)
  real(r8), allocatable :: time(:)

  integer  :: istat, ncid, varid, varid1, ilat, ilon, ilev
  integer  :: strt3(3), cnt3(3)  
  integer  :: i,j,k,ii,jj,kk
  real(r8) :: xlat,xlon 
  
  !dimension for netcdf file 
  strt3(1:3) = 1

  istat=nf90_open(trim(file_name),NF90_NOWRITE,ncid)
  if (istat .ne. NF90_NOERR) then
    write(*,*) 'NF90_OPEN: failed for file',trim(file_name)
    write(*,*) nf90_strerror(istat)
    stop
  endif

  istat=nf90_inq_dimid(ncid,'ncol',varid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
  else 
    istat=nf90_inquire_dimension(ncid,varid,len=cnt3(1))
    if(istat.ne.NF90_NOERR) then
      write(*,*) nf90_strerror(istat)
      stop
    endif 
  endif

  istat=nf90_inq_dimid(ncid,'lev',varid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
  else
    istat=nf90_inquire_dimension(ncid,varid,len=cnt3(2))
    if(istat.ne.NF90_NOERR) then
      write(*,*) nf90_strerror(istat)
      stop
    endif
  endif

  istat=nf90_inq_dimid(ncid,'time',varid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
  else
    istat=nf90_inquire_dimension(ncid,varid,len=cnt3(3))
    if(istat.ne.NF90_NOERR) then
      write(*,*) nf90_strerror(istat)
      stop
    endif
  endif

  !read the data
  allocate(var(cnt3(1),cnt3(2),cnt3(3)))
  allocate(time(cnt3(3)))

  istat=nf90_inq_varid(ncid,'lon',varid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
  else 
    istat=nf90_get_var(ncid,varid,lon)
    if(istat.ne.NF90_NOERR) then
      write(*,*) nf90_strerror(istat)
      stop
    endif
  endif 
  istat=nf90_inq_varid(ncid,'lat',varid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
  else
    istat=nf90_get_var(ncid,varid,lat)
    if(istat.ne.NF90_NOERR) then
      write(*,*) nf90_strerror(istat)
      stop
    endif
  endif
  istat=nf90_inq_varid(ncid,'lev',varid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
  else
    istat=nf90_get_var(ncid,varid,lev)
    if(istat.ne.NF90_NOERR) then
      write(*,*) nf90_strerror(istat)
      stop
    endif
  endif
  istat=nf90_inq_varid(ncid,'time',varid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
  else
    istat=nf90_get_var(ncid,varid,time)
    if(istat.ne.NF90_NOERR) then
      write(*,*) nf90_strerror(istat)
      stop
    endif
  endif

  istat = nf90_inq_varid(ncid,trim(vname),varid)
  if (istat .ne. NF90_NOERR) then
    write(*,*) 'NF90_OPEN: failed for variable',trim(vname)
    write(*,*) nf90_strerror(istat)
  else 
    istat=nf90_get_var(ncid,varid,var,strt3,cnt3)
    if(istat.ne.NF90_NOERR) then
      write(*,*) nf90_strerror(istat)
      stop
    endif
  end if 

  write(*,*) 'reading data from nc file, data dimension: nlon=',cnt3(1),' nlat=',cnt3(1), ' nlev=',cnt3(2), ' ntime=',cnt3(3)
  write(*,*) 'reading data from nc file, current time slice is: ',time(1), ', strt3(3) = ', strt3(3), ', reading variable: ', vname
  write(*,*) 'reading data from nc file, data shape is: ', shape(var)
 
  varout(:,:) = fillvalue 
  varout(:,:) = var(:,:,it) 

  if (any(varout == fillvalue)) then 
    write(*,*) 'processing data from select patch, varout(min/max) = ', minval(varout), maxval(varout) 
    write(*,*) 'processing data from select patch, data contains missing value'
    stop
  end if

return
end subroutine !open_netcdf

subroutine out_model_data (file_name, vname, varin, ncol, nz, nt, lat, lon, lev, l_switch)
  use netcdf

  implicit none
  integer,parameter :: r8  = selected_real_kind(12) ! 8 byte real
  integer,parameter :: r4  = selected_real_kind( 6) ! 4 byte real
  integer,parameter :: cs  = 80                     ! short char
  integer,parameter :: cm  = 160                    ! mid-sized char
  integer,parameter :: cl  = 256                    ! long char
  integer,parameter :: cx  = 512                    ! extra-long char
  integer,parameter :: cxx = 4096                   ! extra-extra-long char

  real(r8), parameter :: fillvalue = 1.e36_r8

  character(len=*), intent(in) :: file_name
  character(len=*), intent(in) :: vname
  logical,          intent(in) :: l_switch
  integer,          intent(in) :: ncol, nz, nt
  real(r8),         intent(in) :: varin(ncol,nz)
  real(r8),         intent(in) :: lon(ncol)
  real(r8),         intent(in) :: lat(ncol)
  real(r8),         intent(in) :: lev(nz)

  real(r8) :: time(nt)
  integer  :: istat, ncid
  integer  :: colmnid,levdid, timdid
  integer  :: colmnvid, levvid, timvid
  integer  :: vardid(3), varvid
  integer  :: strt3(3), cnt3(3)
  integer  :: i,j,k,ii,jj,kk
  real(r8), allocatable :: varout(:,:,:) 

  !dimension for netcdf file 
  strt3(1:3) = 1
  cnt3(1)    = ncol !ncol
  cnt3(2)    = nz   !lev 
  cnt3(3)    = nt
  time(:)    = 0.0_r8

  if (l_switch) then 
    allocate(varout(ncol,nz,nt))
    do i = 1, ncol
      do k = 1, nz 
        do ii = 1, nt 
          varout(i,k,ii) = varin(i,k)
        end do
      end do
    end do 
  else
    allocate(varout(ncol,nz,nt))
    do ii = 1,nt 
      varout(:,:,ii) = varin(:,:)
    end do 
  end if 

  istat=nf90_create(file_name, NF90_CLOBBER, ncid)
  if (istat .ne. NF90_CLOBBER) then
    write(*,*) 'NF90_CREATE: failed for file',trim(file_name)
    write(*,*) nf90_strerror(istat)
    stop
  endif

  !Define the dimensions.
  istat=nf90_def_dim(ncid, "ncol",  cnt3(1), colmnid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if 
  istat=nf90_def_dim(ncid, "lev",  cnt3(2), levdid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if
  istat=nf90_def_dim(ncid, "time", cnt3(3), timdid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if

!Define coordinate variables
  istat=nf90_def_var(ncid, "lon", NF90_DOUBLE, colmnid, colmnvid) 
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if
  istat=nf90_def_var(ncid, "lat", NF90_DOUBLE, colmnid, colmnvid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if
  istat=nf90_def_var(ncid, "lev", NF90_DOUBLE, levdid, levvid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if
  istat=nf90_def_var(ncid, "time", NF90_DOUBLE, timdid, timvid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if

!Define variable
  vardid = (/colmnid,levdid,timdid/)
  istat=nf90_def_var(ncid, trim(vname), NF90_DOUBLE, vardid, varvid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if
  istat=nf90_enddef(ncid) !End Definitions 
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if

!Write Data
  istat=nf90_put_var(ncid, colmnvid, lat)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if
  istat=nf90_put_var(ncid, colmnvid, lon)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if
  istat=nf90_put_var(ncid, levvid, lev)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if
  istat=nf90_put_var(ncid, timvid, time)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if
  istat=nf90_put_var(ncid, varvid, varout)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if
  istat=nf90_close(ncid)
  if(istat.ne.NF90_NOERR) then
    write(*,*) nf90_strerror(istat)
    stop
  end if
  deallocate(varout)
  return 
end subroutine !out_netcdf

