MODULE check_conserve_mod
   USE icosa
   USE pression_mod, ONLY: pression
   IMPLICIT NONE
   PRIVATE

   INTEGER, PARAMETER, PUBLIC :: check_basic = 1, check_detailed = 2, &
                                 AAM_dissip = 1, AAM_dyn = 2, AAM_phys = 3
   INTEGER, SAVE :: check_type

   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(:)

   REAL(rstd), SAVE :: AAM_mass, AAM_mass_old, AAM_vel, AAM_vel_old, &
                       AAM_velp, AAM_velp_old, AAM_velm, AAM_velm_old, &
                       AAM_mass_source(3), AAM_vel_source(3) ! read/written only IF is_master
   REAL(rstd), SAVE :: AAM_vel_plus_source(3), AAM_vel_minus_source(3)
   REAL(rstd), SAVE :: mtot0, ztot0, etot0, angtot0, stot0, mqtot0
   !$OMP THREADPRIVATE(check_type, mtot0, ztot0, etot0, angtot0, stot0)

   PUBLIC :: init_check_conserve, check_conserve_detailed, check_conserve
CONTAINS

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

   SUBROUTINE init_check_conserve
      USE getin_mod
      USE omp_para, ONLY: is_master
      USE abort_mod
      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

      IF (check_type /= check_basic) THEN
         CALL abort_acc("check_conservation /= 'basic'")
      END IF
   END SUBROUTINE init_check_conserve

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

   SUBROUTINE check_conserve(f_ps, f_dps, f_ue, f_theta_rhodz, f_phis, f_q, it)
      USE compute_vorticity_mod
      USE caldyn_gcm_mod
      USE exner_mod
      USE omp_para, ONLY: is_master
      USE grid_param
      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(:)
      TYPE(t_field), POINTER :: f_q(:)
      INTEGER, INTENT(IN) :: it

      REAL(rstd), POINTER :: p(:, :), rhodz(:, :)
      INTEGER :: ind
      REAL(rstd) :: mtot, mqtot, angtot, rmsdpdt
      REAL(rstd) :: etot, stot, ang_mass, ang_vel, ang_velp, ang_velm, rmsvtot, ztot

      IF (grid_type == grid_unst) RETURN ! FIXME

      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_qmass_conserve(f_q, f_rhodz, mqtot)
      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_vel, ang_velp, ang_velm, rmsvtot)

      IF (is_master) THEN ! is_master = is_omp_master && is_mpi_master
         ! AAM_mass and AAM_vel must be read/written only IF is_master
         AAM_mass = ang_mass
         AAM_vel = ang_vel
         AAM_velp = ang_velp
         AAM_velm = ang_velm
         angtot = ang_mass + ang_vel
         IF (it == itau0) THEN
            ztot0 = ztot
            mtot0 = mtot
            mqtot0 = mqtot
            etot0 = etot
            angtot0 = angtot
            stot0 = stot
            AAM_mass_old = AAM_mass
            AAM_vel_old = AAM_vel
            AAM_velp_old = AAM_velp
            AAM_velm_old = AAM_velm
            AAM_mass_source = 0.
            AAM_vel_source = 0.
            AAM_vel_plus_source = 0.
            AAM_vel_minus_source = 0.
         END IF
         rmsvtot = SQRT(rmsvtot/mtot)
         ztot = ztot/ztot0 - 1.; mtot = mtot/mtot0 - 1.
         mqtot = (mqtot + 1E-100)/(mqtot0 + 1e-100) - 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, mqtot, rmsdpdt, etot, ztot, stot, rmsvtot, angtot
         WRITE (134, *) mtot, mqtot, rmsdpdt, etot, ztot, stot, rmsvtot, angtot
         WRITE (134, *) "=================================================="
         WRITE (*, 4000) mtot, mqtot, rmsdpdt, etot, ztot, stot, rmsvtot, angtot

4000     FORMAT(10x, 'masse', 5x, 'advec mass', 5x, 'rmsdpdt', 6x, 'energie', 3x, 'enstrophie' &
                , 5x, 'entropie', 5x, 'rmsv', 5x, 'mt.ang', /, 'GLOB  ' &
                , 2e10.3, e13.5, 5e13.3/)
         CLOSE (134)

         IF (check_type == check_detailed) THEN
            WRITE (*, '(A,6E12.4)') 'AAM_mass : time,old,new,dissip,dyn,phys', dt*it, AAM_mass_old, AAM_mass, &
               AAM_mass_source(AAM_dissip), AAM_mass_source(AAM_dyn), AAM_mass_source(AAM_phys)
            WRITE (*, '(A,6E12.4)') 'AAM_vel  : time,old,new,dissip,dyn,phys', dt*it, AAM_vel_old, AAM_vel, &
               AAM_vel_source(AAM_dissip), AAM_vel_source(AAM_dyn), AAM_vel_source(AAM_phys)
            WRITE (*, '(A,6E12.4)') 'AAM_vel+ : time,old,new,dissip,dyn,phys', dt*it, AAM_velp_old, AAM_velp, &
               AAM_vel_plus_source(AAM_dissip), AAM_vel_plus_source(AAM_dyn), AAM_vel_plus_source(AAM_phys)
            WRITE (*, '(A,6E12.4)') 'AAM_vel- : time,old,new,dissip,dyn,phys', dt*it, AAM_velm_old, AAM_velm, &
               AAM_vel_minus_source(AAM_dissip), AAM_vel_minus_source(AAM_dyn), AAM_vel_minus_source(AAM_phys)
            WRITE (*, '(A,6E12.4)') 'AAM_dyn budget mass + vel:', AAM_mass_source(AAM_dyn) + AAM_vel_source(AAM_dyn)
            AAM_mass_old = AAM_mass
            AAM_vel_old = AAM_vel
            AAM_velp_old = AAM_velp
            AAM_velm_old = AAM_velm
            AAM_mass_source = 0.
            AAM_vel_source = 0.
            AAM_vel_minus_source = 0.
            AAM_vel_plus_source = 0.
         END IF
      END IF

   END SUBROUTINE check_conserve

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

   SUBROUTINE check_conserve_detailed(it, tag, f_ps, f_dps, f_ue, f_theta_rhodz, f_phis)
      USE abort_mod, ONLY: abort_acc
      USE pression_mod
      USE vorticity_mod
      USE caldyn_gcm_mod
      USE exner_mod
      USE omp_para, ONLY: is_master
      INTEGER, 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
      REAL(rstd) :: etot, stot, rmsv, ang_mass, ang_vel, ang_velp, ang_velm

      IF (check_type == check_detailed) THEN
         CALL abort_acc("!check_detailed")
         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, ang_mass, ang_vel, ang_velp, ang_velm, rmsv)

         IF (is_master) THEN
            ! XXX=mass,vel
            ! AAM_XXX = value of total AAM_XXX at last check
            ! ang_XXX = value of total AAM_XXX now
            ! AAM_XXX_source(tag) = source of AAA_XXX due to tag=dissip,dyn,phys
            !WRITE(*,'(A,3E12.2)') 'AAM_mass', tag, ang_mass, AAM_mass, ang_mass-AAM_mass ! FIXME
            !WRITE(*,'(A,3E12.2)') 'AAM_vel ', tag, ang_vel, AAM_vel, ang_vel-AAM_vel ! FIXME
            AAM_mass_source(tag) = AAM_mass_source(tag) + ang_mass - AAM_mass
            AAM_vel_source(tag) = AAM_vel_source(tag) + ang_vel - AAM_vel
            AAM_vel_plus_source(tag) = AAM_vel_plus_source(tag) + ang_velp - AAM_velp
            AAM_vel_minus_source(tag) = AAM_vel_minus_source(tag) + ang_velm - AAM_velm
            !if ((ang_vel - AAM_vel) .gt. 0.) then
            !  AAM_vel_plus_source(tag) = AAM_vel_plus_source(tag) + ang_vel - AAM_vel
            !else
            !  AAM_vel_minus_source(tag) = AAM_vel_minus_source(tag) + ang_vel - AAM_vel
            !endif
            AAM_mass = ang_mass
            AAM_vel = ang_vel
            AAM_velp = ang_velp
            AAM_velm = ang_velm
         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
      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
      END IF
   END SUBROUTINE global_sum

   SUBROUTINE check_mass_conserve(f_ps, f_dps, mtot, rmsdpdt)
      USE mpi_mod
      USE mpipara
      USE transfert_omp_mod
      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

      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)
               END IF
            END DO
         END DO
      END DO

      CALL global_sum(mloc, mtot)
      CALL global_sum(rmsloc, rmsdpdt)
   END SUBROUTINE check_mass_conserve

   SUBROUTINE check_qmass_conserve(f_q, f_rhodz, mqtot)
      USE mpi_mod
      USE mpipara
      USE transfert_omp_mod
      USE omp_para
      IMPLICIT NONE
      TYPE(t_field), POINTER :: f_q(:)
      TYPE(t_field), POINTER :: f_rhodz(:)
      REAL(rstd), POINTER :: q(:, :, :)
      REAL(rstd), POINTER :: rhodz(:, :)
      REAL(rstd), INTENT(OUT) :: mqtot

      INTEGER :: ind, i, j, ij, l
      REAL :: mloc

      mloc = 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)
         q = f_q(ind)
         rhodz = f_rhodz(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
                  DO l = 1, llm
                     mloc = mloc + sum(q(ij, l, :)*Ai(ij)*rhodz(ij, l))
                  END DO
               END IF
            END DO
         END DO
      END DO

      CALL global_sum(mloc, mqtot)
   END SUBROUTINE check_qmass_conserve

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

   SUBROUTINE check_energy(f_ue, f_theta_rhodz, f_phis, etot, &
                           stot, AAM_mass_tot, AAM_vel_tot, AAM_velp_tot, AAM_velm_tot, rmsvtot)
      USE pression_mod
      USE vorticity_mod
      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, AAM_velp_tot, AAM_velm_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, AAM_velp, AAM_velm, rmsv
      INTEGER :: ind

      e = 0; s = 0; AAM_mass = 0; AAM_vel = 0; rmsv = 0
      AAM_velp = 0; AAM_velm = 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(:, :, 1), pk, phis, &
                             e, s, AAM_mass, AAM_vel, AAM_velp, AAM_velm, 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)
      CALL global_sum(AAM_velp, AAM_velp_tot)
      CALL global_sum(AAM_velm, AAM_velm_tot)
   END SUBROUTINE check_energy

   SUBROUTINE compute_energy(ind, u, p, rhodz, theta_rhodz, pk, phis, e, s, AAM_mass, AAM_vel, AAM_velp, AAM_velm, rmsv)
      USE disvert_mod
      USE wind_to_lonlat_mod
      USE omp_para
      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, AAM_velp, AAM_velm, rmsv

      REAL(rstd) :: ucenter(iim*jjm, llm, 3)
      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(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)
                  if (ulon(ij, l) > 0.) then
                     AAM_velp = AAM_velp + rad*cos(lat)*mass_ij*ulon(ij, l)
                  else
                     AAM_velm = AAM_velm + rad*cos(lat)*mass_ij*ulon(ij, l)
                  end if
               END IF
            END DO
         END DO
      END DO

   END SUBROUTINE compute_energy

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

   SUBROUTINE check_PV(ztot)
      REAL(rstd), INTENT(OUT) :: ztot
      REAL(rstd), POINTER :: vort(:, :)
      REAL(rstd), POINTER :: rhodz(:, :)
      INTEGER :: ind
      REAL(rstd) :: z

      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)
      END DO
      CALL global_sum(z, ztot)
   END SUBROUTINE check_PV

   SUBROUTINE compute_PV(vort, rhodz, z)
      USE disvert_mod
      USE omp_para
      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

      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)

            END DO
         END DO

      END DO

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

   SUBROUTINE compute_rhodz(p, rhodz)
      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)
            END DO
         END DO
      END DO

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

END MODULE check_conserve_mod
