MODULE diagflux_mod
  USE icosa
  IMPLICIT NONE
  SAVE
  
  TYPE(t_field),POINTER :: f_qfluxt(:), f_qfluxt_lon(:), f_qfluxt_lat(:) ! time-integrated flux of scalars and its reconstruction at 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_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
       STOP ! FIXME
       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 :: qfluxt(:,:,:)
    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       qfluxt=f_qfluxt(ind)
       qfluxt(:,ll_begin:ll_end,:)=0.
    END DO
  END SUBROUTINE zero_qfluxt

  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, 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 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