MODULE diagflux_mod
   USE icosa
   USE omp_para
   USE caldyn_vars_mod
   IMPLICIT NONE
   SAVE
   PRIVATE

   TYPE(t_field), POINTER, PUBLIC :: &
      f_masst(:), f_qmasst(:), & ! time-averaged mass, tracer mass,
      f_massfluxt(:), f_qfluxt(:), & ! time-integrated mass flux and tracer flux
      f_qfluxt_lon(:), f_qfluxt_lat(:), & ! scalar flux reconstructed at cell centers
      f_ulont(:), f_thetat(:), f_epot(:), f_ekin(:), f_enthalpy(:), & ! time-averaged potential E, kinetic E and enthalpy
      f_ulonfluxt(:), f_thetafluxt(:), f_epotfluxt(:), f_ekinfluxt(:), f_enthalpyfluxt(:) ! time averaged 'fluxes' of epot, ekin and enthalpy
   LOGICAL :: diagflux_on
   !$OMP THREADPRIVATE(diagflux_on)

   PUBLIC :: diagflux_on, init_diagflux, zero_qfluxt, qflux_centered_lonlat, flux_centered_lonlat, diagflux_energy

CONTAINS

   SUBROUTINE init_diagflux
      USE getin_mod
      USE abort_mod
      INTEGER :: ll
      diagflux_on = .FALSE.
      CALL getin("diagflux", diagflux_on)
      IF (diagflux_on) THEN
         CALL abort_acc("diagflux /= .FALSE.")
      END IF

      ll = MERGE(llm, 1, diagflux_on)
      CALL allocate_field(f_masst, field_t, type_real, ll, name="masst")
      CALL allocate_field(f_ulont, field_t, type_real, ll, name="ulont")
      CALL allocate_field(f_thetat, field_t, type_real, ll, name="thetat")
      CALL allocate_field(f_epot, field_t, type_real, ll, name="epot")
      CALL allocate_field(f_ekin, field_t, type_real, ll, name="ekin")
      CALL allocate_field(f_enthalpy, field_t, type_real, ll, name="enthalpy")
      CALL allocate_field(f_qmasst, field_t, type_real, ll, nqtot, name="qmasst")
      CALL allocate_field(f_massfluxt, field_u, type_real, ll, name="massfluxt")
      CALL allocate_field(f_ulonfluxt, field_u, type_real, ll, name="ulonfluxt")
      CALL allocate_field(f_thetafluxt, field_u, type_real, ll, name="thetafluxt")
      CALL allocate_field(f_epotfluxt, field_u, type_real, ll, name="epotfluxt")
      CALL allocate_field(f_ekinfluxt, field_u, type_real, ll, name="ekinfluxt")
      CALL allocate_field(f_enthalpyfluxt, field_u, type_real, ll, name="enthalpyfluxt")
      CALL allocate_field(f_qfluxt, field_u, type_real, ll, nqtot, name="qfluxt")
      CALL allocate_field(f_qfluxt_lon, field_t, type_real, ll, nqtot, name="qfluxt_lon")
      CALL allocate_field(f_qfluxt_lat, field_t, type_real, ll, nqtot, name="qfluxt_lat")
      IF (diagflux_on) CALL zero_qfluxt
   END SUBROUTINE init_diagflux

   SUBROUTINE zero_qfluxt
      INTEGER :: ind
      REAL(rstd), POINTER :: buf2(:, :), buf3(:, :, :)
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         buf2 = f_masst(ind); buf2(:, ll_begin:ll_end) = 0.
         buf2 = f_ulont(ind); buf2(:, ll_begin:ll_end) = 0.
         buf2 = f_thetat(ind); buf2(:, ll_begin:ll_end) = 0.
         buf2 = f_epot(ind); buf2(:, ll_begin:ll_end) = 0.
         buf2 = f_ekin(ind); buf2(:, ll_begin:ll_end) = 0.
         buf2 = f_enthalpy(ind); buf2(:, ll_begin:ll_end) = 0.
         buf3 = f_qmasst(ind); buf3(:, ll_begin:ll_end, :) = 0.
         buf2 = f_massfluxt(ind); buf2(:, ll_begin:ll_end) = 0.
         buf2 = f_ulonfluxt(ind); buf2(:, ll_begin:ll_end) = 0.
         buf2 = f_thetafluxt(ind); buf2(:, ll_begin:ll_end) = 0.
         buf2 = f_epotfluxt(ind); buf2(:, ll_begin:ll_end) = 0.
         buf2 = f_ekinfluxt(ind); buf2(:, ll_begin:ll_end) = 0.
         buf2 = f_enthalpyfluxt(ind); buf2(:, ll_begin:ll_end) = 0.
         buf3 = f_qfluxt(ind); buf3(:, ll_begin:ll_end, :) = 0.
      END DO
   END SUBROUTINE zero_qfluxt

!------------------------------------ Reconstruct fluxes at cell centers ---------------------------------------

   SUBROUTINE qflux_centered_lonlat(scale, f_flux, f_flux_lon, f_flux_lat)
      REAL(rstd), INTENT(IN) :: scale
      TYPE(t_field), POINTER :: f_flux(:), f_flux_lon(:), f_flux_lat(:)
      REAL(rstd), POINTER :: flux(:, :, :), flux_lon(:, :, :), flux_lat(:, :, :)
      INTEGER :: ind, itrac
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         flux = f_flux(ind)
         flux_lon = f_flux_lon(ind)
         flux_lat = f_flux_lat(ind)
         DO itrac = 1, nqtot
            CALL compute_flux_centered_lonlat(scale, flux(:, :, itrac), flux_lon(:, :, itrac), flux_lat(:, :, itrac))
         END DO
      END DO
   END SUBROUTINE qflux_centered_lonlat

   SUBROUTINE flux_centered_lonlat(scale, f_flux, f_flux_lon, f_flux_lat)
      REAL(rstd), INTENT(IN) :: scale
      TYPE(t_field), POINTER :: f_flux(:), f_flux_lon(:), f_flux_lat(:)
      REAL(rstd), POINTER :: flux(:, :), flux_lon(:, :), flux_lat(:, :)
      INTEGER :: ind
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         flux = f_flux(ind)
         flux_lon = f_flux_lon(ind)
         flux_lat = f_flux_lat(ind)
         CALL compute_flux_centered_lonlat(scale, flux, flux_lon, flux_lat)
      END DO
   END SUBROUTINE flux_centered_lonlat

   SUBROUTINE compute_flux_centered_lonlat(scale, flux, flux_lon, flux_lat)
      USE wind_to_lonlat_mod, ONLY: compute_flux_centered, compute_wind_centered_lonlat
      REAL(rstd), INTENT(IN) :: scale
      REAL(rstd), INTENT(IN) :: flux(3*iim*jjm, llm)
      REAL(rstd), INTENT(OUT) :: flux_lon(iim*jjm, llm), flux_lat(iim*jjm, llm)
      REAL(rstd) :: flux_3d(iim*jjm, llm, 3)
      CALL compute_flux_centered(scale, flux, flux_3d)
      CALL compute_wind_centered_lonlat(flux_3d, flux_lon, flux_lat)
   END SUBROUTINE compute_flux_centered_lonlat

!------------------------------------ Compute energy fluxes ---------------------------------------

   SUBROUTINE diagflux_energy(frac, f_phis, f_rhodz, f_theta_rhodz, f_u, f_geopot, f_theta, f_pk, f_hfluxt)
      REAL(rstd), INTENT(IN) :: frac
      TYPE(t_field), POINTER :: f_phis(:), f_rhodz(:), f_theta_rhodz(:), f_u(:), f_geopot(:), f_theta(:), f_pk(:), f_hfluxt(:)
      REAL(rstd), POINTER :: phis(:), rhodz(:, :), theta_rhodz(:, :, :), u(:, :), &
                             geopot(:, :), theta(:, :, :), pk(:, :), hfluxt(:, :), &
                             ulont(:, :), thetat(:, :), epot(:, :), ekin(:, :), enthalpy(:, :), &
                             thetaflux(:, :), ulonflux(:, :), epotflux(:, :), ekinflux(:, :), enthalpyflux(:, :)
      INTEGER :: ind
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         hfluxt = f_hfluxt(ind)
         phis = f_phis(ind)
         rhodz = f_rhodz(ind)
         theta_rhodz = f_theta_rhodz(ind)
         u = f_u(ind)
         geopot = f_geopot(ind)
         theta = f_theta(ind) ! buffer
         pk = f_pk(ind) ! buffer
         ulont = f_ulont(ind)
         thetat = f_thetat(ind)
         epot = f_epot(ind)
         ekin = f_ekin(ind)
         enthalpy = f_enthalpy(ind)
         ulonflux = f_ulonfluxt(ind)
         thetaflux = f_thetafluxt(ind)
         epotflux = f_epotfluxt(ind)
         ekinflux = f_ekinfluxt(ind)
         enthalpyflux = f_enthalpyfluxt(ind)
         CALL compute_diagflux_energy(frac, hfluxt, phis, rhodz, theta_rhodz, u, geopot, theta, pk, &
                                      ulont, thetat, epot, ekin, enthalpy, &
                                      ulonflux, thetaflux, epotflux, ekinflux, enthalpyflux)
      END DO
   END SUBROUTINE diagflux_energy

   SUBROUTINE compute_diagflux_energy(frac, massflux, phis, rhodz, theta_rhodz, ue, geopot, theta, pk, &
                                      ulon, thetat, epot, ekin, enthalpy, &
                                      ulon_flux, thetat_flux, epot_flux, ekin_flux, enthalpy_flux)
      USE disvert_mod, ONLY: ptop
      REAL(rstd), INTENT(IN) :: frac
      REAL(rstd), INTENT(IN) :: massflux(3*iim*jjm, llm), ue(3*iim*jjm, llm), &
                                phis(iim*jjm), rhodz(iim*jjm, llm), theta_rhodz(iim*jjm, llm, nqtot)
      REAL(rstd), INTENT(INOUT) :: geopot(iim*jjm, llm + 1), theta(iim*jjm, llm), pk(iim*jjm, llm) ! theta,pk = buffers
      REAL(rstd), INTENT(INOUT), DIMENSION(iim*jjm, llm)   ::  ulon, thetat, epot, ekin, enthalpy
      REAL(rstd), INTENT(INOUT), DIMENSION(3*iim*jjm, llm) ::  ulon_flux, thetat_flux, epot_flux, ekin_flux, enthalpy_flux
      REAL(rstd) :: energy, p_ik, theta_ik, temp_ik, gv, Rd, cx, cy, cz, ux, uy, uz, ue_le, ulon_i
      INTEGER :: ij, l, ij_omp_begin_ext, ij_omp_end_ext
      Rd = kappa*cpp
      ! even if loops are of the _ext variant, we still need halo exchanges before reconstructing fluxes at cell centers
      ! => loop over interior region
      CALL distrib_level(ij_begin_ext, ij_end_ext, ij_omp_begin_ext, ij_omp_end_ext)

      !--------------------------------------------------------------------------
      !---------------------------- energy_fluxes ----------------------------------
      ! First diagnose geopotential and temperature, column-wise
      !$OMP BARRIER
      !DIR$ SIMD
      DO ij = ij_omp_begin_ext, ij_omp_end_ext
         pk(ij, llm) = ptop + .5*g*rhodz(ij, llm)
      END DO
      DO l = llm - 1, 1, -1
         !DIR$ SIMD
         DO ij = ij_omp_begin_ext, ij_omp_end_ext
            pk(ij, l) = pk(ij, l + 1) + (.5*g)*(rhodz(ij, l) + rhodz(ij, l + 1))
         END DO
      END DO
      ! NB : at this point pressure is stored in array pk
      ! pk then serves as buffer to store temperature
      SELECT CASE (caldyn_thermo)
      CASE (thermo_theta)
         DO l = 1, llm
            !DIR$ SIMD
            DO ij = ij_omp_begin_ext, ij_omp_end_ext
               p_ik = pk(ij, l)
               theta_ik = theta_rhodz(ij, l, 1)/rhodz(ij, l)
               theta(ij, l) = theta_ik
               temp_ik = theta_ik*(p_ik/preff)**kappa
               gv = (g*Rd)*temp_ik/p_ik
               pk(ij, l) = temp_ik
               geopot(ij, l + 1) = geopot(ij, l) + gv*rhodz(ij, l)
            END DO
         END DO
      CASE (thermo_entropy)
         DO l = 1, llm
            !DIR$ SIMD
            DO ij = ij_omp_begin_ext, ij_omp_end_ext
               p_ik = pk(ij, l)
               theta_ik = theta_rhodz(ij, l, 1)/rhodz(ij, l)
               temp_ik = Treff*exp((theta_ik + Rd*log(p_ik/preff))/cpp)
               theta(ij, l) = Treff*exp(theta_ik/cpp)
               gv = (g*Rd)*temp_ik/p_ik ! specific volume v = Rd*T/p
               pk(ij, l) = temp_ik
               geopot(ij, l + 1) = geopot(ij, l) + gv*rhodz(ij, l)
            END DO
         END DO
      END SELECT
      !$OMP BARRIER
      ! Now accumulate energies and energy fluxes
      ! NB : at this point temperature is stored in array pk
      ! pk then serves as buffer to store energy
      ! enthalpy
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            energy = cpp*pk(ij, l)
            enthalpy(ij, l) = enthalpy(ij, l) + frac*rhodz(ij, l)*energy
            pk(ij, l) = energy
         END DO
      END DO
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            enthalpy_flux(ij + u_right, l) = enthalpy_flux(ij + u_right, l) + &
                                             .5*massflux(ij + u_right, l)*(pk(ij, l) + pk(ij + t_right, l))
            enthalpy_flux(ij + u_lup, l) = enthalpy_flux(ij + u_lup, l) + &
                                           .5*massflux(ij + u_lup, l)*(pk(ij, l) + pk(ij + t_lup, l))
            enthalpy_flux(ij + u_ldown, l) = enthalpy_flux(ij + u_ldown, l) + &
                                             .5*massflux(ij + u_ldown, l)*(pk(ij, l) + pk(ij + t_ldown, l))
         END DO
      END DO
      ! potential energy
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            energy = .5*(geopot(ij, l + 1) + geopot(ij, l))
            epot(ij, l) = epot(ij, l) + frac*rhodz(ij, l)*energy
            pk(ij, l) = energy
         END DO
      END DO
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            epot_flux(ij + u_right, l) = epot_flux(ij + u_right, l) + .5*massflux(ij + u_right, l)*(pk(ij, l) + pk(ij + t_right, l))
            epot_flux(ij + u_lup, l) = epot_flux(ij + u_lup, l) + .5*massflux(ij + u_lup, l)*(pk(ij, l) + pk(ij + t_lup, l))
            epot_flux(ij + u_ldown, l) = epot_flux(ij + u_ldown, l) + .5*massflux(ij + u_ldown, l)*(pk(ij, l) + pk(ij + t_ldown, l))
         END DO
      END DO
      ! theta
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            energy = theta(ij, l)
            thetat(ij, l) = thetat(ij, l) + frac*rhodz(ij, l)*energy
            pk(ij, l) = energy
         END DO
      END DO
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            thetat_flux(ij + u_right, l) = thetat_flux(ij + u_right, l) + &
                                           .5*massflux(ij + u_right, l)*(pk(ij, l) + pk(ij + t_right, l))
            thetat_flux(ij + u_lup, l) = thetat_flux(ij + u_lup, l) + &
                                         .5*massflux(ij + u_lup, l)*(pk(ij, l) + pk(ij + t_lup, l))
            thetat_flux(ij + u_ldown, l) = thetat_flux(ij + u_ldown, l) + &
                                           .5*massflux(ij + u_ldown, l)*(pk(ij, l) + pk(ij + t_ldown, l))
         END DO
      END DO
      ! kinetic energy
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            energy = 0.d0
            energy = energy + le(ij + u_rup)*de(ij + u_rup)*ue(ij + u_rup, l)**2
            energy = energy + le(ij + u_lup)*de(ij + u_lup)*ue(ij + u_lup, l)**2
            energy = energy + le(ij + u_left)*de(ij + u_left)*ue(ij + u_left, l)**2
            energy = energy + le(ij + u_ldown)*de(ij + u_ldown)*ue(ij + u_ldown, l)**2
            energy = energy + le(ij + u_rdown)*de(ij + u_rdown)*ue(ij + u_rdown, l)**2
            energy = energy + le(ij + u_right)*de(ij + u_right)*ue(ij + u_right, l)**2
            energy = energy*(.25/Ai(ij))
            ekin(ij, l) = ekin(ij, l) + frac*rhodz(ij, l)*energy
            pk(ij, l) = energy
         END DO
      END DO
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            ekin_flux(ij + u_right, l) = ekin_flux(ij + u_right, l) + .5*massflux(ij + u_right, l)*(pk(ij, l) + pk(ij + t_right, l))
            ekin_flux(ij + u_lup, l) = ekin_flux(ij + u_lup, l) + .5*massflux(ij + u_lup, l)*(pk(ij, l) + pk(ij + t_lup, l))
            ekin_flux(ij + u_ldown, l) = ekin_flux(ij + u_ldown, l) + .5*massflux(ij + u_ldown, l)*(pk(ij, l) + pk(ij + t_ldown, l))
         END DO
      END DO
      ! ulon
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            cx = centroid(ij, 1)
            cy = centroid(ij, 2)
            cz = centroid(ij, 3)
            ux = 0.; uy = 0.; uz = 0.
            ue_le = ne_rup*ue(ij + u_rup, l)*le(ij + u_rup)
            ux = ux + ue_le*(.5*(xyz_v(ij + z_rup, 1) + xyz_v(ij + z_up, 1)) - cx)
            uy = uy + ue_le*(.5*(xyz_v(ij + z_rup, 2) + xyz_v(ij + z_up, 2)) - cy)
            uz = uz + ue_le*(.5*(xyz_v(ij + z_rup, 3) + xyz_v(ij + z_up, 3)) - cz)
            ue_le = ne_lup*ue(ij + u_lup, l)*le(ij + u_lup)
            ux = ux + ue_le*(.5*(xyz_v(ij + z_lup, 1) + xyz_v(ij + z_up, 1)) - cx)
            uy = uy + ue_le*(.5*(xyz_v(ij + z_lup, 2) + xyz_v(ij + z_up, 2)) - cy)
            uz = uz + ue_le*(.5*(xyz_v(ij + z_lup, 3) + xyz_v(ij + z_up, 3)) - cz)
            ue_le = ne_left*ue(ij + u_left, l)*le(ij + u_left)
            ux = ux + ue_le*(.5*(xyz_v(ij + z_lup, 1) + xyz_v(ij + z_ldown, 1)) - cx)
            uy = uy + ue_le*(.5*(xyz_v(ij + z_lup, 2) + xyz_v(ij + z_ldown, 2)) - cy)
            uz = uz + ue_le*(.5*(xyz_v(ij + z_lup, 3) + xyz_v(ij + z_ldown, 3)) - cz)
            ue_le = ne_ldown*ue(ij + u_ldown, l)*le(ij + u_ldown)
            ux = ux + ue_le*(.5*(xyz_v(ij + z_ldown, 1) + xyz_v(ij + z_down, 1)) - cx)
            uy = uy + ue_le*(.5*(xyz_v(ij + z_ldown, 2) + xyz_v(ij + z_down, 2)) - cy)
            uz = uz + ue_le*(.5*(xyz_v(ij + z_ldown, 3) + xyz_v(ij + z_down, 3)) - cz)
            ue_le = ne_rdown*ue(ij + u_rdown, l)*le(ij + u_rdown)
            ux = ux + ue_le*(.5*(xyz_v(ij + z_rdown, 1) + xyz_v(ij + z_down, 1)) - cx)
            uy = uy + ue_le*(.5*(xyz_v(ij + z_rdown, 2) + xyz_v(ij + z_down, 2)) - cy)
            uz = uz + ue_le*(.5*(xyz_v(ij + z_rdown, 3) + xyz_v(ij + z_down, 3)) - cz)
            ue_le = ne_right*ue(ij + u_right, l)*le(ij + u_right)
            ux = ux + ue_le*(.5*(xyz_v(ij + z_rup, 1) + xyz_v(ij + z_rdown, 1)) - cx)
            uy = uy + ue_le*(.5*(xyz_v(ij + z_rup, 2) + xyz_v(ij + z_rdown, 2)) - cy)
            uz = uz + ue_le*(.5*(xyz_v(ij + z_rup, 3) + xyz_v(ij + z_rdown, 3)) - cz)
            ulon_i = ux*elon_i(ij, 1) + uy*elon_i(ij, 2) + uz*elon_i(ij, 3)
            energy = ulon_i*(1./Ai(ij))
            ulon(ij, l) = ulon(ij, l) + frac*rhodz(ij, l)*energy
            pk(ij, l) = energy
         END DO
      END DO
      DO l = ll_begin, ll_end
         !DIR$ SIMD
         DO ij = ij_begin_ext, ij_end_ext
            ulon_flux(ij + u_right, l) = ulon_flux(ij + u_right, l) + .5*massflux(ij + u_right, l)*(pk(ij, l) + pk(ij + t_right, l))
            ulon_flux(ij + u_lup, l) = ulon_flux(ij + u_lup, l) + .5*massflux(ij + u_lup, l)*(pk(ij, l) + pk(ij + t_lup, l))
            ulon_flux(ij + u_ldown, l) = ulon_flux(ij + u_ldown, l) + .5*massflux(ij + u_ldown, l)*(pk(ij, l) + pk(ij + t_ldown, l))
         END DO
      END DO
      !---------------------------- energy_fluxes ----------------------------------
      !--------------------------------------------------------------------------
   END SUBROUTINE compute_diagflux_energy

END MODULE diagflux_mod
