PROGRAM driver

  USE IOIPSL, ONLY: getin, ymds2ju, ioconf_calendar
  USE mod_const_mpi, ONLY: COMM_LMDZ
  USE control_mod, ONLY: planet_type, config_inca, offline
  USE control_mod, ONLY: nday, day_step, iphysiq, dayref, anneeref
  USE infotrac, ONLY: type_trac, nqtot, infotrac_init
  USE temps_mod, ONLY: calend,start_time,annee_ref,day_ref
  USE temps_mod, ONLY: itau_dyn,itau_phy,day_ini,jD_ref,jH_ref,day_end,year_len
  USE comconst_mod, ONLY: cpp, kappa, daysec, dtphys, dtvr, g, r, rad
  USE comconst_mod, ONLY: daylen, year_day, pi, omeg
  USE logic_mod, ONLY: iflag_phys
  USE comvert_mod, ONLY: preff, pa
  USE inigeomphy_mod, ONLY: inigeomphy
  USE iniphysiq_mod, ONLY: iniphysiq

  IMPLICIT NONE

  include "dimensions.h"
  include "paramet.h"
  include "comgeom.h"

!  INTEGER :: step_per_day = 24 ! # of physics steps per day

  INTEGER :: ngrid ! # of grid points on physics grid = 2+(jjm-1)*iim-1/jjm
  INTEGER :: ig,l

!  REAL, PARAMETER :: unjours=86400.,    & ! solar day in seconds
!       &             radius=6.4e6,      & ! planetary radius
!       &             g=9.8,             & ! gravity
!       &             cpp=1004.,         & ! Cp
!       &             kappa=296.945007/cpp ! R/Cp
!  REAL :: timestep !=unjours/step_per_day ! physics time step (s)

!  REAL :: psurf=1e5, ptop=1e4, Temp=250. ! initial values of surface pressure, temperature

  ! dynamics (to load data from start.nc)
  REAL,ALLOCATABLE :: ucov(:,:) ! covariant zonal wind
  REAL,ALLOCATABLE :: vcov(:,:) ! covariant meridional wind
  REAL,ALLOCATABLE :: teta(:,:) ! potential temperature
  REAL,ALLOCATABLE :: q(:,:,:) ! tracers
  REAL,ALLOCATABLE :: masse(:,:) ! air mass
  REAL,ALLOCATABLE :: ps(:) ! surface pressure
  REAL,ALLOCATABLE :: phis(:) ! surface geopotential
  REAL :: time_0
  
  ! physics
  REAL,ALLOCATABLE :: pplev(:,:)  ! pressure at interfaces
  REAL,ALLOCATABLE :: pplay(:,:) ! pressure at full levels
  REAL,ALLOCATABLE :: pphi(:,:) ! geopotential at full levels
  REAL,ALLOCATABLE :: pphis(:) ! surface geopotential
  REAL,ALLOCATABLE :: pt(:,:) ! temperature at full levels
  REAL,ALLOCATABLE :: pu(:,:) ! zonal wind at full levels
  REAL,ALLOCATABLE :: pv(:,:) ! meridional wind at full levels
  REAL,ALLOCATABLE :: pq(:,:,:)! tracers
  REAL,ALLOCATABLE :: pr(:,:) ! vorticity


! 0. Preliminary stuff: read in parameters from run.def
  call conf_gcm( 99, .TRUE.)
!  call conf_planete
  
  daysec=86400. ! day length(s)
  daylen=daysec
  year_day=360 ! # days per year
  rad=6.4e6 ! planet radius (m)
  g=9.81 ! gravity
  cpp=1004. ! Cp
  r=287.05967 ! R
  kappa=287.05967/cpp ! R/Cp
  preff=101325.0
  pa=preff/2.
  pi=2.*asin(1.)
  omeg=2.*pi/daysec*(1./daylen+1./year_day)
  
!  nday = 1 ! simulation lenght (days)
!  CALL getin("nday",nday)
  WRITE(*,*)"driver: nday=",nday
  
!  CALL getin("day_step",day_step)
  WRITE(*,*)"driver: day_step=",day_step
  dtphys=daysec/(day_step/iphysiq) ! physics time step (s)
  dtvr=daysec/day_step
  WRITE(*,*)"driver: timestep= dtphys=",dtphys
  
  
  
! 1. Initialize setup (geometry, physics)
! 1.0. Parameters
  planet_type="earth"
  COMM_LMDZ=1
  iflag_phys=1
  

! 1.1. Tracers-related
  type_trac="lmdz"
  config_inca="none"
  offline=.false.
  
  call infotrac_init
  
  ! 1.2. Geometry 

!  CALL init_latlon(lat, lon)
!  WRITE(*,*)"driver: LAT = " , minval(lat), maxval(lat)
!  WRITE(*,*)"driver: LON = " , minval(lon), maxval(lon)
  ! initialize dyn constants and geometry
  CALL iniconst
  CALL inigeom

  ! 1.3. Initial conditions
  ! dynamics
  ALLOCATE(ucov(ip1jmp1,llm))
  ALLOCATE(vcov(ip1jm,llm))
  ALLOCATE(teta(ip1jmp1,llm))
  ALLOCATE(q(ip1jmp1,llm,nqtot))
  ALLOCATE(masse(ip1jmp1,llm))
  ALLOCATE(ps(ip1jmp1))
  ALLOCATE(phis(ip1jmp1))
  
  CALL dynetat0("start.nc",vcov,ucov,teta,q,masse,ps,phis,time_0)
  

  ! 1.3. Calendar
!  annee_ref=2000
!  dayref=1
  if (calend == 'earth_360d' ) then
    call ioconf_calendar('360d')
  else
    write(*,*) "potential calendar problem with IOIPSL!!"
  endif
  year_len=360

!  annee_ref = anneeref
!  day_ref = dayref
!  day_ini = dayref
  itau_dyn = 0
  itau_phy = 0

  ! Calendar
!  call ymds2ju(annee_ref, mois, day_ref, heure, jD_ref)
  call ymds2ju(annee_ref, 1, day_ref, 0., jD_ref)
  jH_ref = jD_ref - int(jD_ref)
  jD_ref = int(jD_ref)
!  call ioconf_startdate(INT(jD_ref), jH_ref)
  

  day_end = day_ini + nday

  ! 1.4. Initialize physics
  ngrid=2+(jjm-1)*iim-1/jjm
  write(*,*) "driver: ngrid=",ngrid
  CALL iniphysiq(iim,jjm,llm, &
          ngrid,comm_lmdz, &
          daysec,day_ini,dtphys, &
          rlatu,rlatv,rlonu,rlonv,aire,cu,cv,rad,g,r,cpp, &
          iflag_phys)

  
! 2. Initialize fields on physics grid

  ALLOCATE(pplev(ngrid,llm+1)) ! pressure at interfaces
  ALLOCATE(pplay(ngrid,llm)) ! pressure at full levels
  ALLOCATE(pphi(ngrid,llm)) ! geopotential at full levels
  ALLOCATE(pphis(ngrid)) ! surface geopotential
  ALLOCATE(pt(ngrid,llm)) ! temperature (K) at full levels
  ALLOCATE(pu(ngrid,llm)) ! zonal wind (m/s) at full levels
  ALLOCATE(pr(ngrid,llm)) ! relative vorticity
  ALLOCATE(pv(ngrid,llm)) ! meridional wind at full levels
  ALLOCATE(pq(ngrid,llm,nqtot)) ! tracers

  CALL init_temperature_pressure_geopot(ps,masse,teta,phis, &
                                        ngrid,pplev,pplay,pphi,pphis,pt)
  CALL init_winds_tracers(ucov,vcov,q,nqtot, &
                          ngrid,pu,pv,pr,pq)

! 3. Temporal loop  

  ! First time on GPU
  !$acc data create(pplev, pplay, pphi, pphis, pt, pu, pv, pq, pr)

  CALL timeloop(ngrid,llm,nqtot,nday,day_step/iphysiq,dtphys,&
                pplev,pplay,pphi,pphis,&
                pu,pv,pr,pt,pq)

  !$acc end data

  ! output final state as plain text file
  open(10,file="driver.dat")
  do ig=1,ngrid
   do l=1,llm
    write(10,"(I5,I3,X,20(1PE15.8,X))") &
         ig,l,pu(ig,l),pv(ig,l),pt(ig,l)!,pq(ig,l,1:nqtot)
   enddo
  enddo
  close(10)

  write(*,*)"driver: Everything is cool :-)"

CONTAINS

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

   SUBROUTINE timeloop(ngrid,llm,nqtot,nday,calls_per_day,dtphys,&
                pplev,pplay,pphi,pphis,&
                pu,pv,pr,pt,pq)

   USE comvert_mod, ONLY: presnivs
   USE physiq_mod, ONLY: physiq
   USE comconst_mod, ONLY: r
   IMPLICIT NONE

   INTEGER,INTENT(IN) :: ngrid ! # number of atmospheric columns
   INTEGER,INTENT(IN) :: llm ! # number of vertical levels
   INTEGER,INTENT(IN) :: nqtot ! # of tracers
   INTEGER,INTENT(IN) :: nday ! # of days to run
   INTEGER,INTENT(IN) :: calls_per_day ! # number of calls to physics per day
   REAL,INTENT(IN) :: dtphys ! physics time step (s)
   REAL,INTENT(INOUT) :: pplev(ngrid,llm+1) ! pressure at interfaces
   REAL,INTENT(INOUT) :: pplay(ngrid,llm) ! pressure at mid-layer
   REAL,INTENT(INOUT) :: pphi(ngrid,llm) ! geopotential at mid-layer
   REAL,INTENT(INOUT) :: pphis(ngrid) ! surface geopotential
   REAL,INTENT(INOUT) :: pu(ngrid,llm) ! zonal wind (m/s)
   REAL,INTENT(INOUT) :: pv(ngrid,llm) ! meridional wind (m/s)
   REAL,INTENT(INOUT) :: pr(ngrid,llm) ! relative vorticity
   REAL,INTENT(INOUT) :: pt(ngrid,llm) ! temperature (K)
   REAL,INTENT(INOUT) :: pq(ngrid,llm,nqtot) ! tracers
   
   LOGICAL,SAVE :: firstcall=.true.
   LOGICAL,SAVE :: lastcall=.false.
   REAL :: JD_cur ! Julian day
   REAL :: JH_cur ! Julian hour (fraction of day)
   ! tendencies from the physics
   REAL :: pdu(ngrid,llm)
   REAL :: pdv(ngrid,llm)
   REAL :: pdt(ngrid,llm)
   REAL :: pdq(ngrid,llm,nqtot)
   REAL :: pdps(ngrid)
   
   INTEGER :: iday,istep
   INTEGER :: ig,l
   REAL :: rho,gdz

   REAL :: flxw(ngrid,llm) ! vertical mass flux ! set to zero here.
   REAL :: phi_top(ngrid)
   !$acc data present(pplev,pplay,pphi,pphis,pu,pv,pt,pq) &
   !$acc & create(phi_top,flxw,pdu,pdv,pdt,pdq,pdps)
   
   flxw(1:ngrid,1:llm)=0.
   
   DO iday=1,nday
     JD_cur=iday
     DO istep = 0, calls_per_day-1
       JH_cur=istep/calls_per_day
       if ((iday==nday).and.(istep==calls_per_day-1)) lastcall=.true.
       CALL physiq(ngrid,llm,firstcall,lastcall,dtphys, &
                   pplev,pplay,pphi,pphis, &
                   presnivs, &
                   pu,pv,pr,pt,pq,flxw, &
                   pdu,pdv,pdt,pdq,pdps)
       firstcall=.false.
       
       ! add increments sent back from physics:
       !$acc kernels default(present)
       pt(:,:)=pt(:,:)+dtphys*pdt(:,:)
       pu(:,:)=pu(:,:)+dtphys*pdu(:,:)
       pv(:,:)=pv(:,:)+dtphys*pdv(:,:)
       pq(:,:,:)=pq(:,:,:)+dtphys*pdq(:,:,:)
       
       ! Recomputation of relative geopotential:
       phi_top(:)=0
       !$acc loop private(rho, gdz)
       do l=1,llm
         do ig=1,ngrid
           ! rho=p/RT
           rho=r*pplay(ig,l)/pt(ig,l)
           gdz=(pplev(ig,l)-pplev(ig,l+1))/rho ! layer thickness * g
           pphi(ig,l)=phi_top(ig)+0.5*gdz ! geopotential at mid-layer
           phi_top(ig)=phi_top(ig)+gdz ! geopotential at layer top
         enddo
       enddo
       !$acc end kernels
       
     ENDDO ! of DO istep = 0, day_step-1
   ENDDO ! of DO iday=0,nday-1
   
   !$acc end data
   
   END SUBROUTINE timeloop


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

   SUBROUTINE init_winds_tracers(ucov,vcov,q,nqtot, &
                                 ngrid,pu,pv,pr,pq)

    USE comconst_mod, ONLY: pi
    
    IMPLICIT NONE

    include "dimensions.h"
    include "paramet.h"
    include "comgeom2.h"

     REAL,INTENT(IN) :: ucov(iip1,jjp1,llm) ! covariant meridional velocity
     REAL,INTENT(IN) :: vcov(iip1,jjm,llm) ! covariant zonal wind
     REAL,INTENT(IN) :: q(iip1,jjp1,llm,nqtot) ! tracers
     INTEGER,INTENT(IN) :: nqtot ! # of tracers
     INTEGER,INTENT(IN) :: ngrid ! # of physics columns
     REAL,INTENT(OUT) :: pu(ngrid,llm) ! zonal wind speed on physics grid 
     REAL,INTENT(OUT) :: pv(ngrid,llm) ! meridional wind speed on physics grid
     REAL,INTENT(OUT) :: pr(ngrid,llm) ! vorticity on physics grid
     REAL,INTENT(OUT) :: pq(ngrid,llm,nqtot) ! tracers on physics grid
     
     INTEGER :: i,j,l,iq,ig0
     REAL :: rot(iim,jjm,llm) ! rot, on dynamics grid
     REAL :: z1(iim),zsin(iim),zcos(iim)
     REAL,EXTERNAL :: ssum
     
     ! 1. Compute zonal wind on physics grid
     do l=1,llm
       do j=2,jjm
         ig0=1+(j-2)*iim
         pu(ig0+1,l)=0.5*(ucov(iim,j,l)/cu(iim,j) + ucov(1,j,l)/cu(1,j))
         do i=2,iim
           pu(ig0+i,l)=0.5*(ucov(i-1,j,l)/cu(i-1,j) + ucov(i,j,l)/cu(i,j))
         enddo
       enddo ! of do j=2,jjm
     enddo ! of do l=1,llm
     
     ! North pole: u = 1/pi * Integral[v * cos(long) * d long]
     do l=1,llm
       z1(1)=(rlonu(1)-rlonu(iim)+2.*pi)*vcov(1,1,l)/cv(1,1)
       do i=2,iim
         z1(i)=(rlonu(i)-rlonu(i-1))*vcov(i,1,l)/cv(i,1)
       enddo
       zcos(1:iim)=cos(rlonv(1:iim))*z1(1:iim)
       pu(1,l)=ssum(iim,zcos,1)/pi
     enddo ! of l=1,llm
     
     ! South pole: u = 1/pi * Integral[v * cos(long) * d long]
     do l=1,llm
       z1(1)=(rlonu(1)-rlonu(iim)+2.*pi)*vcov(1,jjm,l)/cv(1,jjm)
       do i=2,iim
         z1(i)=(rlonu(i)-rlonu(i-1))*vcov(i,jjm,l)/cv(i,jjm)
       enddo
       zcos(1:iim)=cos(rlonv(1:iim))*z1(1:iim)
       pu(ngrid,l)=ssum(iim,zcos,1)/pi      
     enddo ! of do l=1,llm
     
     
     ! 2. Compute meridional wind on physics grid
     do l=1,llm
       do j=2,jjm
         ig0=1+(j-2)*iim
         do i=1,iim
           pv(ig0+i,l)=0.5*(vcov(i,j-1,l)/cv(i,j-1) + vcov(i,j,l)/cv(i,j))
         enddo
       enddo ! of do j=2,jjm
     enddo ! of do l=1,llm
     
     ! North pole: v = 1/pi * Integral[v * sin(long) * d long ]
     do l=1,llm
       z1(1)=(rlonu(1)-rlonu(iim)+2.*pi)*vcov(1,1,l)/cv(1,1)
       do i=2,iim
         z1(i)=(rlonu(i)-rlonu(i-1))*vcov(i,1,l)/cv(i,1)
       enddo
       zsin(1:iim)=sin(rlonv(1:iim))*z1(1:iim)
       pv(1,l)=ssum(iim,zsin,1)/pi
     enddo ! of l=1,llm
     
     ! South pole: v = 1/pi * Integral[v * sin(long) * d long ]
     do l=1,llm
       z1(1)=(rlonu(1)-rlonu(iim)+2.*pi)*vcov(1,jjm,l)/cv(1,jjm)
       do i=2,iim
         z1(i)=(rlonu(i)-rlonu(i-1))*vcov(i,jjm,l)/cv(i,jjm)
       enddo
       zsin(1:iim)=sin(rlonv(1:iim))*z1(1:iim)
       pv(ngrid,l)=ssum(iim,zsin,1)/pi
     enddo  ! of do l=1,llm
     
     
     ! 3. Compute vorticity, first on dynamics grid
     do l=1,llm
       do i=1,iim
         do j=1,jjm
           rot(i,j,l)=(vcov(i+1,j,l)-vcov(i,j,l) + ucov(i,j+1,l)-ucov(i,j,l)) / &
                      ((cu(i,j)+cu(i,j+1))*(cv(i+1,j)+cv(i,j))*4)
         enddo
       enddo ! of do i=1,iim
     enddo  ! of do l=1,llm
     
     ! compute vorticity, on physics grid
     do l=1,llm
       ! North Pole:
       pr(1,l)=0.
       do j=2,jjm
         ig0=1+(j-2)*iim
         pr(ig0+1,l)=0.25*(rot(iim,j-1,l)+rot(iim,j,l)+ &
                           rot(1,j-1,l)+rot(1,j,l))
         do i=2,iim
           pr(ig0+i,l)=0.25*(rot(i-1,j-1,l)+rot(i-1,j,l)+ &
                             rot(i,j-1,l)+rot(i,j,l))
         enddo
       enddo ! of do j=2,jjm
       ! South pole
       pr(ngrid,l)=0.
     enddo ! of do l=1,llm
     
     ! 4. Tracers, copy them over to the physics grid
     do iq=1,nqtot
       CALL gr_dyn_fi(llm,iip1,jjp1,ngrid,q(1,1,1,iq),pq(1,1,iq))
     enddo
     
  END SUBROUTINE init_winds_tracers

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

  SUBROUTINE init_temperature_pressure_geopot(ps,masse,teta,phis, &
                                              ngrid,pplev,pplay,pphi,pphis,pt)

    USE comvert_mod, ONLY: ap,bp
    USE exner_hyb_m, ONLY: exner_hyb
    USE comconst_mod, ONLY: cpp

    IMPLICIT NONE

    include "dimensions.h"
    include "paramet.h"

    REAL,INTENT(IN) :: ps(ip1jmp1)   ! surface pressure (Pa) on dyn grid
    REAL,INTENT(IN) :: masse(ip1jmp1,llm) ! air mass on dyn. grid
    REAL,INTENT(IN) :: teta(ip1jmp1,llm) ! potential temperature on dyn grid
    REAL,INTENT(IN) :: phis(ip1jmp1)   ! surface geopotential on dyn grid
    INTEGER,INTENT(IN) :: ngrid ! # of physics columns 
    REAL,INTENT(OUT) :: pplev(ngrid,llm+1) ! pressure at interfaces
    REAL,INTENT(OUT) :: pplay(ngrid, llm) ! pressure at mid-layer
    REAL,INTENT(OUT) :: pphi(ngrid, llm) ! geopotential phi=gz at mid-layer
    REAL,INTENT(OUT) :: pphis(ngrid) ! surface geopotential
    REAL,INTENT(OUT) :: pt(ngrid, llm) ! temperature at mid-layer

    REAL :: p(ip1jmp1,llmp1) ! interlayer pressure on dynamics grid
    REAL :: pks(ip1jmp1) ! Exner at the surface
    REAL :: pk(ip1jmp1,llm) ! Exner at mid-layer
    REAL :: phi(ip1jmp1,llm) ! Geopotential on dynamics grid
    
    REAL :: ppk(ngrid,llm) ! Exner at mid-layer on physics grid
    REAL :: pteta(ngrid,llm) ! potential temperature on physics grid
    
    INTEGER :: ig,l

    ! 1. Compute necessary intermediate variables on dynamics grid
    ! Compute pressure on dynamics grid
    CALL pression(ip1jmp1,ap,bp,ps,p)
    ! Compute Exner on dynamics grid
    CALL exner_hyb(ip1jmp1,ps,p,pks,pk) 
    ! Compute geopotential on dynamics grid
    CALL geopot(ip1jmp1,teta,pk,pks,phis,phi)
    
    ! 2. Output fields, on physics grid
    
    ! Copy over p() to pplev()
    CALL gr_dyn_fi(llm+1,iip1,jjp1,ngrid,p,pplev)
    
    ! Compute pplay() based on ab(),bp() and pplev(:,1) (surface pressure)
    do ig=1,ngrid
      pplay(ig,1:llm)=0.5*(ap(1:llm)+bp(1:llm)*pplev(ig,1)+&
                           ap(2:llm+1)+bp(2:llm+1)*pplev(ig,1))
    enddo
    
    ! Copy surface geopotential phis() to pphis()
    CALL gr_dyn_fi(1,iip1,jjp1,ngrid,phis,pphis)
    
    ! Copy geopotential phi() to pphi()
    CALL gr_dyn_fi(llm,iip1,jjp1,ngrid,phi,pphi)
    ! and make pphi relative to surface
    do l=1,llm
      pphi(1:ngrid,l)=pphi(1:ngrid,l)-pphis(1:ngrid)
    enddo
    
    ! Compute temperature from teta, on physics grid
    CALL gr_dyn_fi(llm,iip1,jjp1,ngrid,pk,ppk)
    CALL gr_dyn_fi(llm,iip1,jjp1,ngrid,teta,pteta)
    do l=1,llm
      pt(1:ngrid,l)=pteta(1:ngrid,l)*ppk(1:ngrid,l)/cpp
    enddo
    

  END SUBROUTINE init_temperature_pressure_geopot

END PROGRAM driver