MODULE observable_mod
   USE icosa
   USE diagflux_mod
   USE caldyn_vars_mod
   USE output_field_mod
   IMPLICIT NONE
   PRIVATE

   TYPE(t_field), POINTER, SAVE :: f_buf_i(:), &
                                   f_buf_Fel(:), f_buf_uh(:), & ! horizontal velocity, different from prognostic velocity if NH
                                   f_buf_ulon(:), f_buf_ulat(:)
   TYPE(t_field), POINTER, SAVE :: f_buf_v(:), f_buf_s(:), f_buf_p(:)
   TYPE(t_field), POINTER, SAVE :: f_pmid(:)

! temporary shared variable for caldyn
   TYPE(t_field), POINTER, SAVE :: f_theta(:)

   PUBLIC init_observable, write_output_fields_basic, &
      f_theta, f_buf_i, f_buf_ulon, f_buf_ulat

CONTAINS

   SUBROUTINE init_observable
      USE grid_param
      USE tracer_icosa_mod

      CALL allocate_field(f_buf_i, field_t, type_real, llm, name="buffer_i")
      CALL allocate_field(f_buf_p, field_t, type_real, llm + 1)
      CALL allocate_field(f_buf_ulon, field_t, type_real, llm, name="buf_ulon")
      CALL allocate_field(f_buf_ulat, field_t, type_real, llm, name="buf_ulat")
      CALL allocate_field(f_buf_uh, field_u, type_real, llm, name="buf_uh")
      CALL allocate_field(f_buf_Fel, field_u, type_real, llm + 1, name="buf_F_el")
      CALL allocate_field(f_buf_v, field_z, type_real, llm, name="buf_v")
      CALL allocate_field(f_buf_s, field_t, type_real, name="buf_s")

      CALL allocate_field(f_theta, field_t, type_real, llm, nqdyn, name='theta') ! potential temperature
      CALL allocate_field(f_pmid, field_t, type_real, llm, name='pmid')       ! mid layer pressure
   END SUBROUTINE init_observable

   SUBROUTINE write_output_fields_basic(init, f_phis, f_ps, f_mass, f_geopot, f_theta_rhodz, f_u, f_W, f_q)
      USE xios_mod
      USE disvert_mod
      USE omp_para
      USE time_mod
      USE xios_mod
      USE earth_const
      USE pression_mod, ONLY: pression_mid
      USE compute_pression_mod, ONLY: hydrostatic_pressure
      USE compute_temperature_mod, ONLY: temperature
      USE compute_velocity_mod
      USE compute_vorticity_mod
      USE compute_divergence_mod
      USE vertical_interp_mod
      USE theta2theta_rhodz_mod
      USE compute_omega_mod, ONLY: w_omega
      USE wind_to_lonlat_mod, ONLY: un2ulonlat
      USE kinetic_mod
      USE diagflux_mod
      USE grid_param
      USE tracer_icosa_mod

      LOGICAL, INTENT(IN) :: init
      INTEGER :: l
      REAL :: scalar(1)
      REAL :: mid_ap(llm)
      REAL :: mid_bp(llm)

      TYPE(t_field), POINTER :: f_phis(:), f_ps(:), f_mass(:), f_geopot(:), f_theta_rhodz(:), f_u(:), f_W(:), f_q(:)
!    IF (is_master) PRINT *,'CALL write_output_fields_basic'

      CALL transfert_request(f_ps, req_i1)

      IF (init) THEN
         IF (is_master) PRINT *, 'Creating output files ...'
         scalar(1) = dt
         IF (is_omp_master) CALL xios_send_field("timestep", scalar)
         scalar(1) = preff
         IF (is_omp_master) CALL xios_send_field("preff", scalar)
         IF (is_omp_master) CALL xios_send_field("ap", ap)
         IF (is_omp_master) CALL xios_send_field("bp", bp)
         DO l = 1, llm
            mid_ap(l) = (ap(l) + ap(l + 1))/2
            mid_bp(l) = (bp(l) + bp(l + 1))/2
         END DO
         IF (is_omp_master) CALL xios_send_field("mid_ap", mid_ap)
         IF (is_omp_master) CALL xios_send_field("mid_bp", mid_bp)

         CALL output_field("phis", f_phis)
         CALL output_field("Ai", geom%Ai)
         IF (is_master) PRINT *, '... done creating output files. Writing initial condition ...'
      END IF

      IF (nqdyn > 1) THEN
         CALL divide_by_mass(2, f_mass, f_theta_rhodz, f_buf_i)
         IF (init) THEN
            CALL output_field("dyn_q_init", f_buf_i)
         ELSE
            CALL output_field("dyn_q", f_buf_i)
         END IF
      END IF

      CALL divide_by_mass(1, f_mass, f_theta_rhodz, f_buf_i)
      IF (init) THEN
         CALL output_field("theta_init", f_buf_i)
      ELSE
         CALL output_field("theta", f_buf_i)
      END IF

      CALL hydrostatic_pressure(f_mass, f_theta_rhodz, f_ps, f_pmid)
      CALL temperature(f_pmid, f_q, f_buf_i) ! f_buf_i : IN = theta, out = T

      IF (init) THEN
         CALL output_field("p_init", f_pmid)
         CALL output_field("ps_init", f_ps)
         CALL output_field("mass_init", f_mass)
         CALL output_field("geopot_init", f_geopot)
         CALL output_field("q_init", f_q)

         CALL output_field("temp_init", f_buf_i)
         CALL output_field("vort_init", f_buf_v)
      ELSE
         CALL output_field("p", f_pmid)
         CALL output_field("ps", f_ps)
         CALL output_field("mass", f_mass)
         CALL output_field("geopot", f_geopot)
         CALL output_field("q", f_q)

         CALL output_field("temp", f_buf_i)
         CALL vertical_interp(f_pmid, f_buf_i, f_buf_s, 85000.)
         CALL output_field("t850", f_buf_s)
         CALL vertical_interp(f_pmid, f_buf_i, f_buf_s, 50000.)
         CALL output_field("t500", f_buf_s)
         CALL vertical_interp(f_pmid, f_buf_i, f_buf_s, preff)
         CALL output_field("SST", f_buf_s)
      END IF

      CALL vorticity(f_u, f_buf_v)
      CALL divergence(f_u, f_buf_i)
      IF (init) THEN
         CALL output_field("vort_init", f_buf_v)
         CALL output_field("div_init", f_buf_i)
      ELSE
         CALL output_field("vort", f_buf_v)
         CALL output_field("div", f_buf_i)
      END IF

      CALL velocity(f_geopot, f_ps, f_mass, f_u, f_W, f_buf_Fel, f_buf_uh, f_buf_i)
      CALL transfert_request(f_buf_uh, req_e1_vect)
      CALL un2ulonlat(f_buf_uh, f_buf_ulon, f_buf_ulat)

      IF (init) THEN
         CALL output_field("uz_init", f_buf_i)
         CALL output_field("ulon_init", f_buf_ulon)
         CALL output_field("ulat_init", f_buf_ulat)
         IF (is_master) PRINT *, 'Done writing initial condition ...'

      ELSE
         CALL output_field("uz", f_buf_i)
         CALL output_field("ulon", f_buf_ulon)
         CALL output_field("ulat", f_buf_ulat)

         !       CALL output_field("exner",f_pk)
         !       CALL output_field("pv",f_qv)

         CALL vertical_interp(f_pmid, f_buf_ulon, f_buf_s, 85000.)
         CALL output_field("u850", f_buf_s)
         CALL vertical_interp(f_pmid, f_buf_ulon, f_buf_s, 50000.)
         CALL output_field("u500", f_buf_s)

         CALL vertical_interp(f_pmid, f_buf_ulat, f_buf_s, 85000.)
         CALL output_field("v850", f_buf_s)
         CALL vertical_interp(f_pmid, f_buf_ulat, f_buf_s, 50000.)
         CALL output_field("v500", f_buf_s)

         CALL vertical_interp(f_pmid, f_buf_i, f_buf_s, 85000.)
         CALL output_field("w850", f_buf_s)
         CALL vertical_interp(f_pmid, f_buf_i, f_buf_s, 50000.)
         CALL output_field("w500", f_buf_s)

         CALL vertical_interp_4d(f_pmid, f_q, 1, f_buf_s, 85000.)
         CALL output_field("q1_850", f_buf_s)
         CALL vertical_interp_4d(f_pmid, f_q, 1, f_buf_s, 50000.)
         CALL output_field("q1_500", f_buf_s)

         IF (grid_type == grid_unst) RETURN

         CALL w_omega(f_ps, f_u, f_buf_i)
         CALL output_field("omega", f_buf_i)
         CALL vertical_interp(f_pmid, f_buf_i, f_buf_s, 85000.)
         CALL output_field("omega850", f_buf_s)
         CALL vertical_interp(f_pmid, f_buf_i, f_buf_s, 50000.)
         CALL output_field("omega500", f_buf_s)
      END IF

      IF (grid_type == grid_unst) RETURN

      CALL kinetic(f_u, f_buf_i)
      IF (init) THEN
         CALL output_field("kinetic_trisk_init", f_buf_i)
      ELSE
         CALL output_field("kinetic_trisk", f_buf_i)
      END IF

      CALL kinetic_new(f_u, f_buf_v, f_buf_i)
      IF (init) THEN
         CALL output_field("kinetic_init", f_buf_i)
      ELSE
         CALL output_field("kinetic", f_buf_i)
      END IF

      IF (.NOT. init) THEN
         IF (diagflux_on) THEN
            CALL flux_centered_lonlat(1./(itau_out*dt), f_massfluxt, f_buf_ulon, f_buf_ulat)
            CALL output_field("mass_t", f_masst)
            CALL output_field("massflux_lon", f_buf_ulon)
            CALL output_field("massflux_lat", f_buf_ulat)

            CALL output_energyflux(f_ulont, f_ulonfluxt, "ulon_t", "ulonflux_lon", "ulonflux_lat")
            CALL output_energyflux(f_thetat, f_thetafluxt, "theta_t", "thetaflux_lon", "thetaflux_lat")
            CALL output_energyflux(f_epot, f_epotfluxt, "epot_t", "epotflux_lon", "epotflux_lat")
            CALL output_energyflux(f_ekin, f_ekinfluxt, "ekin_t", "ekinflux_lon", "ekinflux_lat")
            CALL output_energyflux(f_enthalpy, f_enthalpyfluxt, "enthalpy_t", "enthalpyflux_lon", "enthalpyflux_lat")

            CALL qflux_centered_lonlat(1./(itau_out*dt), f_qfluxt, f_qfluxt_lon, f_qfluxt_lat)
            CALL output_field("qmass_t", f_qmasst)
            CALL output_field("qflux_lon", f_qfluxt_lon)
            CALL output_field("qflux_lat", f_qfluxt_lat)
            CALL zero_qfluxt  ! restart accumulating fluxes
         END IF
      END IF
   END SUBROUTINE write_output_fields_basic

   SUBROUTINE output_energyflux(f_energy, f_flux, name_energy, name_fluxlon, name_fluxlat)
      USE diagflux_mod, ONLY: flux_centered_lonlat
      TYPE(t_field), POINTER :: f_energy(:), f_flux(:)
      CHARACTER(*), INTENT(IN) :: name_energy, name_fluxlon, name_fluxlat
      CALL transfert_request(f_flux, req_e1_vect)
      CALL flux_centered_lonlat(1./(itau_out*dt), f_flux, f_buf_ulon, f_buf_ulat)
      CALL output_field(name_energy, f_energy)
      CALL output_field(name_fluxlon, f_buf_ulon)
      CALL output_field(name_fluxlat, f_buf_ulat)
   END SUBROUTINE output_energyflux

   SUBROUTINE divide_by_mass(iq, f_mass, f_theta_rhodz, f_theta)
      INTEGER, INTENT(IN) :: iq
      TYPE(t_field), POINTER :: f_mass(:), f_theta_rhodz(:), f_theta(:)
      REAL(rstd), POINTER :: mass(:, :), theta_rhodz(:, :, :), theta(:, :)
      INTEGER :: ind
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         mass = f_mass(ind)
         theta_rhodz = f_theta_rhodz(ind)
         theta = f_theta(ind)
         theta(:, :) = theta_rhodz(:, :, iq)/mass(:, :)
      END DO
   END SUBROUTINE divide_by_mass

END MODULE observable_mod
