MODULE diagflux_mod
  USE icosa
  IMPLICIT NONE
  SAVE
  
  TYPE(t_field),POINTER :: &
       f_masst(:), f_qmasst(:), & ! time-integrated mass, tracer mass, 
       f_massfluxt(:), f_qfluxt(:), & ! mass flux and tracer flux
       f_qfluxt_lon(:), f_qfluxt_lat(:) ! scalar flux reconstructed  cell centers
  LOGICAL :: diagflux_on
  !$OMP THREADPRIVATE(diagflux_on)

CONTAINS

  SUBROUTINE init_diagflux
    USE getin_mod
    diagflux_on = .FALSE.
    CALL getin("diagflux", diagflux_on)
    IF(diagflux_on) THEN
       CALL allocate_field(f_masst,  field_t,type_real,llm, name="masst")
       CALL allocate_field(f_qmasst,  field_t,type_real,llm,nqtot, name="qmasst")
       CALL allocate_field(f_massfluxt, field_u,type_real,llm, name="massfluxt")
       CALL allocate_field(f_qfluxt,  field_u,type_real,llm,nqtot, name="qfluxt")
       CALL allocate_field(f_qfluxt_lon,  field_t,type_real,llm,nqtot, name="qfluxt_lon")
       CALL allocate_field(f_qfluxt_lat,  field_t,type_real,llm,nqtot, name="qfluxt_lat")
       CALL zero_qfluxt
    ELSE
       CALL allocate_field(f_masst,  field_t,type_real,0, name="masst")
       CALL allocate_field(f_qmasst,  field_t,type_real,llm,0, name="qmasst")
       CALL allocate_field(f_massfluxt, field_u,type_real,0, name="massfluxt")
       CALL allocate_field(f_qfluxt,  field_u,type_real,llm,0, name="qfluxt")
       CALL allocate_field(f_qfluxt_lon,  field_t,type_real,llm,0, name="qfluxt_lon")
       CALL allocate_field(f_qfluxt_lat,  field_t,type_real,llm,0, name="qfluxt_lat")
    END IF
  END SUBROUTINE init_diagflux

  SUBROUTINE zero_qfluxt
    USE mpipara
    USE omp_para
    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_massfluxt(ind)
       buf2(:,ll_begin:ll_end)=0.
       buf3=f_qmasst(ind)
       buf3(:,ll_begin:ll_end,:)=0.
       buf3=f_qfluxt(ind)
       buf3(:,ll_begin:ll_end,:)=0.
    END DO
  END SUBROUTINE zero_qfluxt

  SUBROUTINE flux_centered_lonlat(scale, f_massflux, f_flux, f_massflux_lon, f_massflux_lat, f_flux_lon, f_flux_lat)
    REAL(rstd), INTENT(IN) :: scale
    TYPE(t_field),POINTER :: f_flux(:), f_flux_lon(:), f_flux_lat(:), &
         f_massflux(:), f_massflux_lon(:), f_massflux_lat(:)
    REAL(rstd), POINTER :: flux(:,:,:), flux_lon(:,:,:), flux_lat(:,:,:), &
         massflux(:,:), massflux_lon(:,:), massflux_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
       massflux=f_massflux(ind)
       massflux_lon=f_massflux_lon(ind)
       massflux_lat=f_massflux_lat(ind)
       CALL compute_flux_centered_lonlat(scale, massflux, massflux_lon, massflux_lat)
    END DO
  END SUBROUTINE flux_centered_lonlat
  
  SUBROUTINE compute_flux_centered_lonlat(scale, flux, flux_lon, flux_lat)
    USE wind_mod
    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_compound(flux_3d, flux_lon, flux_lat)
  END SUBROUTINE compute_flux_centered_lonlat

END MODULE diagflux_mod