MODULE check_conserve_mod
  USE icosa 
  IMPLICIT NONE 

  PRIVATE 
    TYPE(t_field),POINTER,SAVE :: f_pk(:), f_pks(:), f_p(:)
    TYPE(t_field),POINTER,SAVE :: f_vort(:)
    TYPE(t_field),POINTER,SAVE :: f_rhodz(:) 

    INTEGER, PARAMETER :: check_basic=1, check_detailed=2
    INTEGER, SAVE :: check_type
  PUBLIC init_check_conserve, check_conserve, check_conserve_detailed
    REAL(rstd),SAVE:: mtot0,ztot0,etot0,angtot0,stot0,rmsvtot0
!$OMP THREADPRIVATE(check_type, mtot0,ztot0,etot0,angtot0,stot0,rmsvtot0)
     
CONTAINS 

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

  SUBROUTINE init_check_conserve
    USE icosa
    USE getin_mod
    USE omp_para, ONLY : is_master
    IMPLICIT NONE
    CHARACTER(LEN=255) :: check_type_str
    CALL allocate_field(f_pk,field_t,type_real,llm)
    CALL allocate_field(f_p,field_t,type_real,llm+1)
    CALL allocate_field(f_pks,field_t,type_real)
    CALL allocate_field(f_rhodz,field_t,type_real,llm)
    CALL allocate_field(f_vort,field_z,type_real,llm)
    
    check_type_str='basic'
    CALL getin("check_conservation",check_type_str)
    SELECT CASE(TRIM(check_type_str))
    CASE('basic')
       check_type=check_basic
    CASE('detailed')
       check_type=check_detailed
    CASE DEFAULT
       IF(is_master) PRINT*,'Bad selector <', TRIM(check_type_str), &
            '> for variable check_conservation : options are <basic>, <detailed>'
       STOP
    END SELECT
  END SUBROUTINE init_check_conserve

!--------------------------------- Basic check --------------------------------

  SUBROUTINE check_conserve(f_ps,f_dps,f_ue,f_theta_rhodz,f_phis,it) 
  USE icosa 
  USE pression_mod
  USE vorticity_mod
  USE caldyn_gcm_mod
  USE exner_mod 
  USE mpipara, ONLY : is_mpi_master, comm_icosa
  USE omp_para
  IMPLICIT NONE
    TYPE(t_field),POINTER :: f_ps(:)
    TYPE(t_field),POINTER :: f_dps(:)
    TYPE(t_field),POINTER :: f_ue(:)
    TYPE(t_field),POINTER :: f_theta_rhodz(:)
    TYPE(t_field),POINTER :: f_phis(:)
    INTEGER, INTENT(IN) :: it

    REAL(rstd),POINTER :: p(:,:),rhodz(:,:) 
    INTEGER::ind,ierr
    REAL(rstd) :: mtot, angtot, rmsdpdt
    REAL(rstd) :: etot, stot, ang_mass, ang_vit, rmsvtot, ztot

    CALL transfert_request(f_ue,req_e1_vect)
    CALL pression(f_ps,f_p)

    DO ind=1,ndomain 
      IF (.NOT. assigned_domain(ind)) CYCLE
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      p=f_p(ind) 
      rhodz=f_rhodz(ind)
      CALL compute_rhodz(p,rhodz) 
    END DO

    CALL vorticity(f_ue,f_vort)
    CALL check_mass_conserve(f_ps,f_dps,mtot,rmsdpdt)
    CALL check_PV(ztot) 
    CALL exner(f_ps,f_p,f_pks,f_pk)
    CALL check_energy(f_ue,f_theta_rhodz,f_phis, etot, stot, ang_mass, ang_vit, rmsvtot) 
    angtot  = ang_mass + ang_vit 

    IF (is_master) THEN 
       IF ( it == itau0  ) THEN 
          ztot0 = ztot
          mtot0 = mtot
          etot0 = etot 
          angtot0  = angtot
          stot0 = stot 
       END IF

       rmsvtot=SQRT(rmsvtot/mtot) 
       ztot=ztot/ztot0-1. ; mtot=mtot/mtot0-1. 
       etot=etot/etot0-1. ; angtot=angtot/angtot0-1. ; stot=stot/stot0-1. 
       rmsdpdt= daysec*1.e-2*sqrt(rmsdpdt/ncell_glo)  
       
       OPEN(134,file="checkconsicosa.txt",position='append')
       WRITE(134,4000)mtot,rmsdpdt,etot,ztot,stot,rmsvtot,angtot
       WRITE(134,*)mtot,rmsdpdt,etot,ztot,stot,rmsvtot,angtot 
       WRITE(134,*)"=================================================="
       WRITE(*,4000)mtot,rmsdpdt,etot,ztot,stot,rmsvtot,angtot
       
4000   FORMAT(10x,'masse',5x,'rmsdpdt',5x,'energie',5x,'enstrophie'  &
            ,5x,'entropie',5x,'rmsv',5x,'mt.ang',/,'GLOB  '                &
            ,e10.3,e13.6,5e13.3/)     
       close(134)
      
    END IF
  END SUBROUTINE check_conserve

!------------------------ Detailed check (Lauritzen et al., 2014) -----------------------------

  SUBROUTINE check_conserve_detailed(tag, f_ps,f_dps,f_ue,f_theta_rhodz,f_phis,it) 
    USE icosa 
    USE pression_mod
    USE vorticity_mod
    USE caldyn_gcm_mod
    USE exner_mod 
    USE mpipara, ONLY : is_mpi_root, comm_icosa
    IMPLICIT NONE
    CHARACTER(LEN=*), INTENT(IN) :: tag
    TYPE(t_field),POINTER :: f_ps(:)
    TYPE(t_field),POINTER :: f_dps(:)
    TYPE(t_field),POINTER :: f_ue(:)
    TYPE(t_field),POINTER :: f_theta_rhodz(:)
    TYPE(t_field),POINTER :: f_phis(:)
    INTEGER, INTENT(IN) ::it

    REAL(rstd),POINTER :: p(:,:),rhodz(:,:) 
    INTEGER::ind,ierr
    REAL(rstd) :: mtot, ztot, rmsdpdt, etot,stot,rmsv, AAM_mass, AAM_vel
    
    IF(check_type == check_detailed) THEN
       etot=0.0; stot=0.0; rmsv=0.0 
       AAM_vel=0. ; AAM_vel=0.
       ztot = 0.0 

       CALL transfert_request(f_ue,req_e1_vect)
       CALL pression(f_ps,f_p)
       DO ind=1,ndomain 
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          p=f_p(ind) 
          rhodz=f_rhodz(ind)
          CALL compute_rhodz(p,rhodz) 
       END DO
       CALL exner(f_ps,f_p,f_pks,f_pk)
       CALL check_energy(f_ue,f_theta_rhodz,f_phis, etot, stot, AAM_mass, AAM_vel, rmsv) 
       
       IF (is_mpi_root) THEN 
          !$OMP MASTER       
          IF ( it == itau0  ) Then 
             ztot0 = ztot
             mtot0 = mtot
             etot0 = etot 
             angtot0  = AAM_mass + AAM_vel
             stot0 = stot 
          END IF
          
          PRINT *, TRIM(tag), AAM_mass, AAM_vel
          !$OMP END MASTER 
       END IF
    END IF
  END SUBROUTINE check_conserve_detailed
 
!---------------------------------------------------------------------

  SUBROUTINE global_sum(loc_sum, glob_sum)
    USE mpi_mod
    USE mpipara
    USE transfert_omp_mod
    IMPLICIT NONE
    REAL(rstd), INTENT(IN) :: loc_sum
    REAL(rstd), INTENT(OUT) :: glob_sum
    REAL(rstd) :: sum

    CALL reduce_sum_omp(loc_sum, sum)
    IF (using_mpi) THEN
!$OMP MASTER    
      CALL MPI_REDUCE(sum, glob_sum, 1, MPI_REAL8, MPI_SUM, 0, comm_icosa, ierr)
!$OMP END MASTER
    ELSE
      glob_sum=loc_sum
    ENDIF 
  END SUBROUTINE global_sum

  SUBROUTINE check_mass_conserve(f_ps,f_dps,mtot,rmsdpdt)
  USE mpi_mod
  USE mpipara
  USE transfert_omp_mod
  USE icosa
  USE omp_para
  IMPLICIT NONE
    TYPE(t_field),POINTER :: f_ps(:)
    TYPE(t_field),POINTER :: f_dps(:)
    REAL(rstd),POINTER :: ps(:),dps(:) 
    REAL(rstd), INTENT(OUT) :: mtot, rmsdpdt

    INTEGER :: ind,i,j,ij  
    REAL :: mloc, rmsloc
    REAL :: mloc_mpi, rmsloc_mpi

    mloc=0.0; rmsloc=0.0
    DO ind=1,ndomain
      IF (.NOT. assigned_domain(ind) .OR. .NOT. is_omp_level_master) CYCLE
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      dps=f_dps(ind) 
      ps=f_ps(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
             mloc=mloc+ps(ij)*Ai(ij)
             rmsloc=rmsloc+dps(ij)*dps(ij)
	  ENDIF
        ENDDO
      ENDDO
    ENDDO
    
    CALL global_sum(mloc, mtot)
    CALL global_sum(rmsloc, rmsdpdt)
  END SUBROUTINE check_mass_conserve

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

  SUBROUTINE check_energy(f_ue,f_theta_rhodz,f_phis, etot, &
                      stot, AAM_mass_tot, AAM_vel_tot, rmsvtot)  
  USE icosa
  USE pression_mod
  USE vorticity_mod
  IMPLICIT NONE
    TYPE(t_field), POINTER :: f_ue(:)
    TYPE(t_field), POINTER :: f_theta_rhodz(:)
    TYPE(t_field), POINTER :: f_phis(:)
    REAL(rstd),INTENT(OUT) :: etot, stot, AAM_mass_tot, AAM_vel_tot, rmsvtot
    REAL(rstd), POINTER :: ue(:,:)
    REAL(rstd), POINTER :: p(:,:) 
    REAL(rstd), POINTER :: theta_rhodz(:,:) 
    REAL(rstd), POINTER :: pk(:,:) 
    REAL(rstd), POINTER :: phis(:) 
    REAL(rstd), POINTER :: rhodz(:,:) 
    REAL(rstd) :: e, s, AAM_mass, AAM_vel, rmsv
    INTEGER :: ind

    e = 0 ; s = 0 ; AAM_mass = 0 ; AAM_vel=0 ; rmsv = 0

    DO ind=1,ndomain 
      IF (.NOT. assigned_domain(ind)) CYCLE
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      ue=f_ue(ind) 
      p=f_p(ind)
      rhodz=f_rhodz(ind) 
      theta_rhodz=f_theta_rhodz(ind) 
      pk=f_pk(ind) 
      phis=f_phis(ind) 
      CALL compute_energy(ind,ue,p,rhodz,theta_rhodz,pk,phis, &
           e, s, AAM_mass, AAM_vel, rmsv)
    END DO 

    CALL global_sum(e, etot)
    CALL global_sum(s, stot)
    CALL global_sum(rmsv, rmsvtot)
    CALL global_sum(AAM_mass, AAM_mass_tot)
    CALL global_sum(AAM_vel, AAM_vel_tot)
  END SUBROUTINE check_energy

  SUBROUTINE compute_energy(ind,u,p,rhodz,theta_rhodz,pk,phis, e, s, AAM_mass, AAM_vel, rmsv)
    USE icosa
    USE disvert_mod
    USE wind_mod
    USE omp_para
    IMPLICIT NONE
    INTEGER,INTENT(IN)::ind 
    REAL(rstd),INTENT(IN) :: u(3*iim*jjm,llm)
    REAL(rstd),INTENT(IN) :: p(iim*jjm,llm+1)
    REAL(rstd),INTENT(IN) :: rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(IN) :: theta_rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(IN) :: pk(iim*jjm,llm)
    REAL(rstd),INTENT(IN) :: phis(iim*jjm)
    REAL(rstd),INTENT(INOUT) :: e, s, AAM_mass, AAM_vel, rmsv
    
    REAL(rstd) :: ucenter(iim*jjm,3,llm)
    REAL(rstd) :: ulon(iim*jjm,llm)
    REAL(rstd) :: ulat(iim*jjm,llm)
    
    REAL(rstd) :: mass_ij, theta_ij, KE_ij, rad,radomeg,lat
    INTEGER :: i,j,ij,l
    
    CALL compute_wind_centered(u,ucenter)
    CALL compute_wind_centered_lonlat_compound(ucenter,ulon,ulat)
    rad = radius 
    radomeg = rad*omega
    
    DO l = ll_begin,ll_end
       DO j=jj_begin,jj_end
          DO i=ii_begin,ii_end
             ij=(j-1)*iim+i
             IF (domain(ind)%own(i,j)) THEN
                lat = lat_i(ij)  
                mass_ij = rhodz(ij,l)*Ai(ij) 
                theta_ij = theta_rhodz(ij,l)/rhodz(ij,l)
                KE_ij = 1/(4*Ai(ij))*(le(ij+u_right)*de(ij+u_right)*u(ij+u_right,l)**2 +   &
                     le(ij+u_rup)*de(ij+u_rup)*u(ij+u_rup,l)**2 +           &
                     le(ij+u_lup)*de(ij+u_lup)*u(ij+u_lup,l)**2 +           &
                     le(ij+u_left)*de(ij+u_left)*u(ij+u_left,l)**2 +        &
                     le(ij+u_ldown)*de(ij+u_ldown)*u(ij+u_ldown,l)**2 +     &
                     le(ij+u_rdown)*de(ij+u_rdown)*u(ij+u_rdown,l)**2 )  
                rmsv = rmsv + 2*mass_ij*KE_ij
                s    = s + Ai(ij)*theta_rhodz(ij,l)
                e    = e + mass_ij*(phis(ij)+theta_ij*pk(ij,l)+KE_ij)
                AAM_mass = AAM_mass + rad*cos(lat)*mass_ij*radomeg*cos(lat)
                AAM_vel  = AAM_vel  + rad*cos(lat)*mass_ij*ulon(ij,l)
             END IF
          END DO
       END DO
    END DO
    
END SUBROUTINE compute_energy

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

  SUBROUTINE check_PV(ztot)
  USE icosa
  USE mpi_mod
  USE mpipara
  USE transfert_omp_mod
  IMPLICIT NONE
     REAL(rstd),INTENT(OUT) :: ztot
     REAL(rstd), POINTER :: vort(:,:)
     REAL(rstd), POINTER :: rhodz(:,:) 
     INTEGER :: ind
     REAL(rstd) :: z,z_mpi

     z=0
     DO ind=1,ndomain 
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       vort=f_vort(ind)
       rhodz=f_rhodz(ind) 
       CALL compute_PV(vort,rhodz,z) 
     ENDDO

    IF (using_mpi) THEN
      CALL reduce_sum_omp(z, z_mpi)
!$OMP MASTER    
      CALL MPI_REDUCE(z_mpi, ztot, 1, MPI_REAL8, MPI_SUM, 0, comm_icosa, ierr)
!$OMP END MASTER
    ELSE
      ztot=z
    ENDIF 
 
  END SUBROUTINE check_PV 

  SUBROUTINE compute_PV(vort,rhodz,z) 
  USE icosa
  USE disvert_mod
  USE omp_para
  IMPLICIT NONE 
    REAL(rstd),INTENT(IN) :: vort(2*iim*jjm,llm)
    REAL(rstd),INTENT(IN) :: rhodz(iim*jjm,llm) 
    REAL(rstd),INTENT(INOUT) :: z
    REAL(rstd)::qv1,qv2 
    REAL(rstd)::hv1,hv2  
    INTEGER :: i,j,ij,l,ij2

    hv1 = 0.0 ; hv2 = 0.0 
   
    DO l = ll_begin,ll_end
      DO j=jj_begin+1,jj_end-1
        DO i=ii_begin+1,ii_end-1
          ij=(j-1)*iim+i
          hv1 =  Riv2(ij,vup)        * rhodz(ij,l)           &
            + Riv2(ij+t_rup,vldown) * rhodz(ij+t_rup,l)     &
            + Riv2(ij+t_lup,vrdown) * rhodz(ij+t_lup,l)
    
          qv1 = ( vort(ij+z_up,l)+fv(ij+z_up) )/hv1
          
          hv2  =  Riv2(ij,vdown)     * rhodz(ij,l)          &
            + Riv2(ij+t_ldown,vrup) * rhodz(ij+t_ldown,l)  &
            + Riv2(ij+t_rdown,vlup) * rhodz(ij+t_rdown,l)
                     
          qv2 =( vort(ij+z_down,l)+fv(ij+z_down) )/hv2

          z = z + (qv1*qv1*hv1 + qv2*qv2*hv2) 

        ENDDO
      ENDDO
      
    ENDDO

  END SUBROUTINE compute_PV
!---------------------------------------------------------------------

   SUBROUTINE compute_rhodz(p,rhodz) 
   USE icosa 
   IMPLICIT NONE 
     REAL(rstd),INTENT(IN) :: p(iim*jjm,llm+1)
     REAL(rstd),INTENT(OUT):: rhodz(iim*jjm,llm)
     REAL(rstd) ::mass(iim*jjm,llm)
     INTEGER :: i,j,ij,l 
	
     DO l = 1, llm
      DO j=jj_begin,jj_end
       DO i=ii_begin,ii_end
         ij=(j-1)*iim+i
         mass(ij,l) = ( p(ij,l) - p(ij,l+1) ) * Ai(ij)/g
         rhodz(ij,l) = mass(ij,l) / Ai(ij)
       ENDDO
      ENDDO
     ENDDO

   END SUBROUTINE compute_rhodz
!====================================================================


 END MODULE check_conserve_mod