MODULE caldyn_adv_mod
  USE icosa
  IMPLICIT NONE
  PRIVATE
  PUBLIC :: init_caldyn, caldyn

CONTAINS

  SUBROUTINE init_caldyn
  END SUBROUTINE init_caldyn

  SUBROUTINE check_mass_conservation(f_ps,f_dps)
    USE icosa
    TYPE(t_field),POINTER :: f_ps(:)
    TYPE(t_field),POINTER :: f_dps(:)
    REAL(rstd),POINTER :: ps(:)
    REAL(rstd),POINTER :: dps(:)
    REAL(rstd) :: mass_tot,dmass_tot
    INTEGER :: ind,i,j,ij

    mass_tot=0
    dmass_tot=0

    CALL transfert_request(f_dps,req_i1)
    CALL transfert_request(f_ps,req_i1)

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)

       ps=f_ps(ind)
       dps=f_dps(ind)

       DO j=jj_begin,jj_end
          DO i=ii_begin,ii_end
             ij=(j-1)*iim+i
             IF (domain(ind)%own(i,j)) THEN
                mass_tot=mass_tot+ps(ij)*Ai(ij)/g
                dmass_tot=dmass_tot+dps(ij)*Ai(ij)/g
             ENDIF
          ENDDO
       ENDDO

    ENDDO
    PRINT*, "mass_tot ", mass_tot,"      dmass_tot ",dmass_tot        

  END SUBROUTINE check_mass_conservation


  SUBROUTINE caldyn(write_out,f_phis, f_ps, f_theta_rhodz, f_u, f_q, &
       f_hflux, f_wflux, f_dps, f_dtheta_rhodz, f_du)
    USE icosa
    USE output_field_mod
    USE vorticity_mod
    USE kinetic_mod
    USE theta2theta_rhodz_mod
    IMPLICIT NONE
    LOGICAL,INTENT(IN)    :: write_out
    TYPE(t_field),POINTER :: f_phis(:)
    TYPE(t_field),POINTER :: f_ps(:)
    TYPE(t_field),POINTER :: f_theta_rhodz(:)
    TYPE(t_field),POINTER :: f_u(:)
    TYPE(t_field),POINTER :: f_q(:)
    TYPE(t_field),POINTER :: f_hflux(:), f_wflux(:)
    TYPE(t_field) :: f_dps(:)
    TYPE(t_field) :: f_dtheta_rhodz(:)
    TYPE(t_field) :: f_du(:)

    REAL(rstd),POINTER :: ps(:)
    REAL(rstd),POINTER :: u(:,:)
    REAL(rstd),POINTER :: dps(:)
    REAL(rstd),POINTER :: hflux(:,:), wflux(:,:)
    REAL(rstd),POINTER :: dtheta_rhodz(:,:), du(:,:)  ! set to 0
    INTEGER :: ind

    CALL transfert_request(f_ps,req_i1) 
    CALL transfert_request(f_u,req_e1_vect)

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       ps=f_ps(ind)
       u=f_u(ind)
       dps=f_dps(ind)
       hflux=f_hflux(ind)
       wflux=f_wflux(ind)
       dtheta_rhodz=f_dtheta_rhodz(ind)
       du=f_du(ind)

!       !$OMP PARALLEL DEFAULT(SHARED)
       CALL compute_caldyn(ps,u,hflux, wflux, dps, dtheta_rhodz, du)
!       !$OMP END PARALLEL
    ENDDO

    IF (write_out) THEN
       CALL output_field("wflux",f_wflux)
       CALL output_field("ps",f_ps)
    ENDIF
    !    CALL check_mass_conservation(f_ps,f_dps)

  END SUBROUTINE caldyn

  SUBROUTINE compute_caldyn(ps,u, hflux,wflux,dps, dtheta_rhodz,du)
    USE icosa
    USE disvert_mod
    IMPLICIT NONE
    REAL(rstd),INTENT(IN)  :: ps(iim*jjm)
    REAL(rstd),INTENT(IN)  :: u(iim*3*jjm,llm)
    REAL(rstd),INTENT(OUT) :: du(iim*3*jjm,llm), hflux(iim*3*jjm,llm) ! hflux in kg/s
    REAL(rstd),INTENT(OUT) :: dtheta_rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: dps(iim*jjm)
    REAL(rstd),INTENT(OUT) :: wflux(iim*jjm,llm+1) ! vertical mass flux (kg/m2/s)

    REAL(rstd),ALLOCATABLE :: rhodz(:,:)
    REAL(rstd),ALLOCATABLE :: divm(:,:)  ! mass flux divergence

    INTEGER :: i,j,ij,l    
    LOGICAL,SAVE :: first=.TRUE.

    ALLOCATE(rhodz(iim*jjm,llm))
    ALLOCATE(divm(iim*jjm,llm))    ! mass flux divergence

    dtheta_rhodz(:,:)=0.
    du(:,:)=0.

!!! Compute mass
    DO l = 1, llm
       DO j=jj_begin-1,jj_end+1
          DO i=ii_begin-1,ii_end+1
             ij=(j-1)*iim+i
             rhodz(ij,l) = (ap(l)-ap(l+1) + ps(ij)*(bp(l)-bp(l+1)) )/g
          ENDDO
       ENDDO
    ENDDO

    DO l = 1, llm
!!!  Mass fluxes
       DO j=jj_begin-1,jj_end+1
          DO i=ii_begin-1,ii_end+1
             ij=(j-1)*iim+i
             hflux(ij+u_right,l)=0.5*(rhodz(ij,l)+rhodz(ij+t_right,l))*u(ij+u_right,l)*le(ij+u_right)
             hflux(ij+u_lup,l)=0.5*(rhodz(ij,l)+rhodz(ij+t_lup,l))*u(ij+u_lup,l)*le(ij+u_lup)
             hflux(ij+u_ldown,l)=0.5*(rhodz(ij,l)+rhodz(ij+t_ldown,l))*u(ij+u_ldown,l)*le(ij+u_ldown)
          ENDDO
       ENDDO
!!! Horizontal divergence of fluxes
       DO j=jj_begin,jj_end
          DO i=ii_begin,ii_end
             ij=(j-1)*iim+i  
             ! divm = +div(mass flux), sign convention as in Ringler et al. 2012, eq. 21
             divm(ij,l)= 1./Ai(ij)*(ne(ij,right)*hflux(ij+u_right,l)   +  &
                  ne(ij,rup)*hflux(ij+u_rup,l)       +  &  
                  ne(ij,lup)*hflux(ij+u_lup,l)       +  &  
                  ne(ij,left)*hflux(ij+u_left,l)     +  &  
                  ne(ij,ldown)*hflux(ij+u_ldown,l)   +  &  
                  ne(ij,rdown)*hflux(ij+u_rdown,l))
          ENDDO
       ENDDO
    ENDDO

!!! cumulate mass flux divergence from top to bottom
    DO  l = llm-1, 1, -1
       !$OMP DO
       DO j=jj_begin,jj_end
          DO i=ii_begin,ii_end
             ij=(j-1)*iim+i
             divm(ij,l) = divm(ij,l) + divm(ij,l+1)
          ENDDO
       ENDDO
    ENDDO

!!! Compute vertical mass flux
    DO l = 1,llm-1
       DO j=jj_begin,jj_end
          DO i=ii_begin,ii_end
             ij=(j-1)*iim+i
             ! w = int(z,ztop,div(flux)dz) + B(eta)dps/dt
             ! => w>0 for upward transport
             wflux( ij, l+1 ) = divm( ij, l+1 ) - bp(l+1) * divm( ij, 1 )
          ENDDO
       ENDDO
    ENDDO

    ! compute dps, set vertical mass flux at the surface to 0
    DO j=jj_begin,jj_end
       DO i=ii_begin,ii_end
          ij=(j-1)*iim+i
          wflux(ij,1)  = 0.
          ! dps/dt = -int(div flux)dz
          dps(ij)=-divm(ij,1) * g 
       ENDDO
    ENDDO

    DEALLOCATE(rhodz)
    DEALLOCATE(divm)
    
  END SUBROUTINE compute_caldyn

END MODULE caldyn_adv_mod