MODULE nudging_mod
   USE icosa, ONLY: rstd
   USE grid_param, ONLY: llm, nqdyn
   USE earth_const, ONLY: radius
   USE omp_para, ONLY: ll_begin, ll_end
   USE domain_mod, ONLY: ndomain, assigned_domain
   USE dimensions, ONLY: swap_dimensions, u_right, u_lup, u_ldown
   USE dimensions, ONLY: iim, jjm, ij_begin_ext, ij_end_ext
   USE geometry, ONLY: swap_geometry
   USE field_mod
   IMPLICIT NONE
   SAVE
   PRIVATE
   ! nudging will be active outside a disc of radius 'radius' centered at 'center_lon', 'center lat'.
   REAL(rstd) :: center_lon, center_lat, nudging_radius, nudging_relaxation_time, &
                 nudging_relaxation_time_in, nudging_relaxation_time_out
   !$OMP THREADPRIVATE(center_lon, center_lat, nudging_radius, nudging_relaxation_time,nudging_relaxation_time_in,nudging_relaxation_time_out)

   TYPE(t_field), POINTER :: f_relax_coef_e(:), f_target_ue(:), &
                             f_relax_coef_i(:), f_target_theta_rhodz(:), f_target_ps(:)
   CHARACTER(LEN=255) :: guided_nudging_field
   INTEGER :: nudging_time

   PUBLIC :: init_guided, guided

CONTAINS

   SUBROUTINE init_guided(f_u, f_theta_rhodz, f_ps)
      USE getin_mod, ONLY: getin
      USE math_const, ONLY: pi
      USE earth_const, ONLY: scale_factor
      TYPE(t_field), POINTER :: f_u(:)! initial condition
      TYPE(t_field), POINTER :: f_theta_rhodz(:)! initial condition
      TYPE(t_field), POINTER :: f_ps(:)! initial condition
      REAL(rstd), POINTER :: ue(:, :), target_ue(:, :), coef_e(:)
      REAL(rstd), POINTER :: theta_rhodz(:, :, :), target_theta_rhodz(:, :, :), coef_i(:)
      REAL(rstd), POINTER :: ps2(:), target_ps(:)
      INTEGER :: ind
      ! read DEF keys describing how to relax
      center_lon = 0.
      CALL getin('nudging_center_lon', center_lon)
      center_lat = 0.
      CALL getin('nudging_center_lat', center_lat)
      nudging_radius = 0.
      CALL getin('nudging_radius', nudging_radius)
      nudging_radius = nudging_radius/scale_factor
      ! we should check that radius>0

      nudging_time = 0
      CALL getin('nudging_time', nudging_time)
      nudging_time = nudging_time/scale_factor
      !nudging_time=0

      nudging_relaxation_time_in = 0.
      nudging_relaxation_time_out = 0.
      CALL getin('nudging_relaxation_time_in', nudging_relaxation_time_in)
      CALL getin('nudging_relaxation_time_out', nudging_relaxation_time_out)
      nudging_relaxation_time_in = nudging_relaxation_time_in/scale_factor
      nudging_relaxation_time_out = nudging_relaxation_time_out/scale_factor
      IF (nudging_relaxation_time_in .EQ. 0. .AND. nudging_relaxation_time_out .EQ. 0.) THEN
         PRINT *, " nudging_relaxation_time_in and nudging_relaxation_time_out values can not be 0. at the same time"
      END IF

      CALL getin("guided_nudging_field", guided_nudging_field)

      SELECT CASE (TRIM(guided_nudging_field))
      CASE ('wind', 'temperature', 'ps', 'wind_temperature', 'wind_temperature_ps')
         CALL allocate_field(f_relax_coef_e, field_u, type_real, name='nudging_coef_e')
         CALL allocate_field(f_target_ue, field_u, type_real, llm, name='nudging_target_e')
         CALL allocate_field(f_relax_coef_i, field_t, type_real, name='nudging_coef_i')
         CALL allocate_field(f_target_theta_rhodz, field_t, type_real, llm, nqdyn, name='nudging_target_theta')
         CALL allocate_field(f_target_ps, field_t, type_real, name='nudging_target_ps')
      CASE DEFAULT
         PRINT *, "Bad selector for varaible init_guided_nudging>", TRIM(guided_nudging_field)
         PRINT *, "options are <wind>, <temperature>,<ps>,<wind_temperature>,<wind_temperature_ps>"
         STOP
      END SELECT
      ! compute relax_coef and target_ue
      center_lon = center_lon*pi/180.
      center_lat = center_lat*pi/180.
      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)
         coef_e = f_relax_coef_e(ind)
         coef_i = f_relax_coef_i(ind)
         target_ue = f_target_ue(ind)
         ue = f_u(ind)
         target_theta_rhodz = f_target_theta_rhodz(ind)
         theta_rhodz = f_theta_rhodz(ind)
         target_ps = f_target_ps(ind)
         ps2 = f_ps(ind)
         CALL compute_relax_coef(coef_e, coef_i)
         CALL compute_target_u(ue, target_ue)
         CALL compute_target_center(theta_rhodz, target_theta_rhodz)
         CALL compute_target_center2(ps2, target_ps)
      END DO
   END SUBROUTINE init_guided

!----------------------------- Compute relaxation coefficients ------------------------------

   SUBROUTINE compute_relax_coef(coef_e, coef_i)
      USE geometry, ONLY: lon_e, lat_e, lon_i, lat_i
      REAL(rstd), INTENT(OUT) :: coef_e(iim*3*jjm), coef_i(iim*jjm)
      INTEGER :: l, ij
      DO ij = ij_begin_ext, ij_end_ext
         coef_e(ij + u_right) = relax_coef(lon_e(ij + u_right), lat_e(ij + u_right))
         coef_e(ij + u_lup) = relax_coef(lon_e(ij + u_lup), lat_e(ij + u_lup))
         coef_e(ij + u_ldown) = relax_coef(lon_e(ij + u_ldown), lat_e(ij + u_ldown))
         coef_i(ij) = relax_coef(lon_i(ij), lat_i(ij))
      END DO
   END SUBROUTINE compute_relax_coef

   FUNCTION relax_coef(lon, lat)
      USE spherical_geom_mod, ONLY: dist_lonlat
      USE time_mod!, ONLY : dt ! time step for a full RK step
      REAL(rstd), INTENT(IN) :: lon, lat
      REAL(rstd) :: relax_coef, dist, c
      ! NB : dist is computed on unit sphere
      IF (nudging_radius > 0.) THEN
         CALL dist_lonlat(lon, lat, center_lon, center_lat, dist)
         c = tanh((1.-radius*dist/nudging_radius)*20.) ! 1 inside circle, -1 outside
         IF (c > 0.99) c = 1
         IF (c < -0.99) c = -1
         c = .5*(1.+c) ! rescale to [0,1] range ; ! 1 inside circle, 0 outside
      ELSE
         c = 0.
      END IF

      nudging_relaxation_time = c*nudging_relaxation_time_in + (1.-c)*nudging_relaxation_time_out
      relax_coef = 1.-dt/(dt + nudging_relaxation_time)
   END FUNCTION relax_coef

!----------------------------- Copy initial condition as target ------------------------------

   SUBROUTINE compute_target_u(ue, target_ue)
      REAL(rstd), INTENT(OUT)  :: target_ue(iim*3*jjm, llm)
      REAL(rstd), INTENT(IN) :: ue(iim*3*jjm, llm)
      INTEGER :: l, ij
      DO l = ll_begin, ll_end
         DO ij = ij_begin_ext, ij_end_ext
            target_ue(ij + u_right, l) = ue(ij + u_right, l)
            target_ue(ij + u_lup, l) = ue(ij + u_lup, l)
            target_ue(ij + u_ldown, l) = ue(ij + u_ldown, l)
         END DO
      END DO
   END SUBROUTINE compute_target_u

   SUBROUTINE compute_target_center(theta_rhodz, target_theta_rhodz)
      REAL(rstd), INTENT(OUT)  :: target_theta_rhodz(iim*jjm, llm, nqdyn)
      REAL(rstd), INTENT(IN) :: theta_rhodz(iim*jjm, llm, nqdyn)
      INTEGER :: l, ij, iq
      DO iq = 1, nqdyn
         DO l = ll_begin, ll_end
            DO ij = ij_begin_ext, ij_end_ext
               target_theta_rhodz(ij, l, iq) = theta_rhodz(ij, l, iq)
            END DO
         END DO
      END DO
   END SUBROUTINE compute_target_center

   SUBROUTINE compute_target_center2(theta_rhodz, target_theta_rhodz)
      REAL(rstd), INTENT(OUT)  :: target_theta_rhodz(iim*jjm, nqdyn)
      REAL(rstd), INTENT(IN) :: theta_rhodz(iim*jjm, nqdyn)
      INTEGER :: ij, iq
      DO iq = 1, nqdyn
         DO ij = ij_begin_ext, ij_end_ext
            target_theta_rhodz(ij, iq) = theta_rhodz(ij, iq)
         END DO
      END DO
   END SUBROUTINE compute_target_center2

!----------------------------- Relax towards target ------------------------------

   SUBROUTINE guided(tt, f_ps, f_theta_rhodz, f_u, f_q)
      USE xios_mod
      USE theta2theta_rhodz_mod
      USE wind_from_lonlat_mod
      USE transfert_mod
      USE time_mod, ONLY: dt ! time step for a full RK step
      USE vertical_remap_mod
      USE pression_mod, ONLY: pression_mid
      REAL(rstd), INTENT(IN):: tt
      TYPE(t_field), POINTER :: f_ps(:)
      TYPE(t_field), POINTER, SAVE :: f_pmid_target(:)
      TYPE(t_field), POINTER :: f_phis(:)
      TYPE(t_field), POINTER :: f_theta_rhodz(:)
      TYPE(t_field), POINTER :: f_u(:)
      TYPE(t_field), POINTER :: f_q(:)
      TYPE(t_field), POINTER, SAVE :: f_T_guided(:), f_ulon_guided(:), f_ulat_guided(:), f_T_guided_interp(:), &
                                      f_ulon_guided_interp(:), f_ulat_guided_interp(:), f_ps_guided(:)
      REAL(rstd), POINTER :: target_ue(:, :), ue(:, :), coef_e(:)
      REAL(rstd), POINTER :: target_theta_rhodz(:, :, :), theta_rhodz(:, :, :), coef_i(:)
      REAL(rstd), POINTER :: target_ps(:), ps2(:)
      INTEGER :: ind

      IF (abs(MOD(tt, REAL(nudging_time)) - dt) < 1.0D-2) THEN
         CALL allocate_field(f_ps_guided, field_t, type_real, name='nudging_ps')
         CALL xios_read_field("ps_guided_read", f_ps_guided)
         CALL transfert_request(f_ps_guided, req_i0)
         CALL transfert_request(f_ps_guided, req_i1)
         CALL allocate_field(f_T_guided, field_t, type_real, llm, name='nudging_T')
         CALL allocate_field(f_T_guided_interp, field_t, type_real, llm, name='nudging_T')
         CALL allocate_field(f_ulon_guided_interp, field_t, type_real, llm, name='nudging_T')
         CALL allocate_field(f_ulat_guided_interp, field_t, type_real, llm, name='nudging_T')
         CALL allocate_field(f_ulon_guided, field_t, type_real, llm, name='nudging_ulon')
         CALL allocate_field(f_ulat_guided, field_t, type_real, llm, name='nudging_ulat')
         CALL allocate_field(f_pmid_target, field_t, type_real, llm, name='target_pressure')
         CALL xios_read_field("T_guided_read", f_T_guided)
         CALL pression_mid(f_ps, f_pmid_target)
         CALL vertical_remap_extdata(f_T_guided, f_pmid_target, f_T_guided_interp)
         CALL transfert_request(f_T_guided, req_i0)
         CALL transfert_request(f_T_guided_interp, req_i0)
         CALL transfert_request(f_T_guided, req_i1)
         CALL transfert_request(f_T_guided_interp, req_i1)
         CALL temperature2theta_rhodz(f_ps, f_T_guided_interp, f_target_theta_rhodz)
         CALL xios_read_field("ulon_guided_read", f_ulon_guided)
         CALL xios_read_field("ulat_guided_read", f_ulat_guided)
         CALL vertical_remap_extdata(f_ulon_guided, f_pmid_target, f_ulon_guided_interp)
         CALL vertical_remap_extdata(f_ulat_guided, f_pmid_target, f_ulat_guided_interp)
         CALL transfert_request(f_ulon_guided, req_i0)
         CALL transfert_request(f_ulon_guided_interp, req_i0)
         CALL transfert_request(f_ulat_guided, req_i0)
         CALL transfert_request(f_ulat_guided_interp, req_i0)
         CALL transfert_request(f_ulon_guided, req_i1)
         CALL transfert_request(f_ulon_guided_interp, req_i1)
         CALL transfert_request(f_ulat_guided, req_i1)
         CALL transfert_request(f_ulat_guided_interp, req_i1)
         !CALL xios_write_field("ulat_guided_out",f_ulat_guided)
         CALL xios_write_field("ulat_guided_out", f_T_guided_interp) ! FIX
         CALL xios_write_field("ulon_guided_out", f_ulon_guided)
         CALL xios_write_field("T_guided_out", f_T_guided)
         CALL xios_write_field("PS_guided_out", f_ps_guided)
         CALL ulonlat2un(f_ulon_guided_interp, f_ulat_guided, f_target_ue)
         CALL deallocate_field(f_T_guided)
         CALL deallocate_field(f_T_guided_interp)
         CALL deallocate_field(f_ulon_guided)
         CALL deallocate_field(f_ulon_guided_interp)
         CALL deallocate_field(f_ulat_guided)
         CALL deallocate_field(f_ulat_guided_interp)
         CALL deallocate_field(f_pmid_target)
      END IF

      DO ind = 1, ndomain
         IF (.NOT. assigned_domain(ind)) CYCLE
         CALL swap_dimensions(ind)
         CALL swap_geometry(ind)

         SELECT CASE (TRIM(guided_nudging_field))
         CASE ('wind')
            coef_e = f_relax_coef_e(ind)
            target_ue = f_target_ue(ind)
            ue = f_u(ind)
            CALL compute_guided_u(coef_e, target_ue, ue)
         CASE ('temperature')
            coef_i = f_relax_coef_i(ind)
            target_theta_rhodz = f_target_theta_rhodz(ind)
            theta_rhodz = f_theta_rhodz(ind)
            CALL compute_guided_center(coef_i, target_theta_rhodz, theta_rhodz)
         CASE ('ps')
            coef_i = f_relax_coef_i(ind)
            target_ps = f_ps_guided(ind)
            ps2 = f_ps(ind)
            CALL compute_guided_center2(coef_i, target_ps, ps2)
         CASE ('wind_temperature')
            coef_e = f_relax_coef_e(ind)
            target_ue = f_target_ue(ind)
            ue = f_u(ind)
            CALL compute_guided_u(coef_e, target_ue, ue)
            coef_i = f_relax_coef_i(ind)
            target_theta_rhodz = f_target_theta_rhodz(ind)
            theta_rhodz = f_theta_rhodz(ind)
            CALL compute_guided_center(coef_i, target_theta_rhodz, theta_rhodz)
         CASE ('wind_temperature_ps')
            coef_e = f_relax_coef_e(ind)
            target_ue = f_target_ue(ind)
            ue = f_u(ind)
            CALL compute_guided_u(coef_e, target_ue, ue)
            coef_i = f_relax_coef_i(ind)
            target_theta_rhodz = f_target_theta_rhodz(ind)
            theta_rhodz = f_theta_rhodz(ind)
            CALL compute_guided_center(coef_i, target_theta_rhodz, theta_rhodz)
            target_ps = f_ps_guided(ind)
            ps2 = f_ps(ind)
            CALL compute_guided_center2(coef_i, target_ps, ps2)
         CASE DEFAULT
            PRINT *, "Bad selector for varaible <guided_nudging_field>", TRIM(guided_nudging_field)
            PRINT *, "options are <wind>, <temperature>,<ps>,<wind_temperature>,<wind_temperature_ps>"
            STOP
         END SELECT
      END DO

   END SUBROUTINE guided

   SUBROUTINE compute_guided_u(coef_e, target_ue, ue)
      REAL(rstd), INTENT(IN) :: coef_e(iim*3*jjm)
      REAL(rstd), INTENT(IN) :: target_ue(iim*3*jjm, llm)
      REAL(rstd), INTENT(INOUT) :: ue(iim*3*jjm, llm)
      INTEGER :: l, ij
      DO l = ll_begin, ll_end
         DO ij = ij_begin_ext, ij_end_ext
            ue(ij + u_right, l) = ue(ij + u_right, l)*coef_e(ij + u_right) + &
                                  target_ue(ij + u_right, l)*(1.-coef_e(ij + u_right))
            ue(ij + u_lup, l) = ue(ij + u_lup, l)*coef_e(ij + u_lup) + &
                                target_ue(ij + u_lup, l)*(1.-coef_e(ij + u_lup))
            ue(ij + u_ldown, l) = ue(ij + u_ldown, l)*coef_e(ij + u_ldown) + &
                                  target_ue(ij + u_ldown, l)*(1.-coef_e(ij + u_ldown))
         END DO
      END DO
   END SUBROUTINE compute_guided_u

   SUBROUTINE compute_guided_center(coef_i, target_theta_rhodz, theta_rhodz)
      REAL(rstd), INTENT(IN) :: coef_i(iim*jjm)
      REAL(rstd), INTENT(IN) :: target_theta_rhodz(iim*jjm, llm, nqdyn)
      REAL(rstd), INTENT(INOUT) :: theta_rhodz(iim*jjm, llm, nqdyn)
      INTEGER :: l, ij, iq
      DO iq = 1, nqdyn
         DO l = ll_begin, ll_end
            DO ij = ij_begin_ext, ij_end_ext
               theta_rhodz(ij, l, iq) = theta_rhodz(ij, l, iq)*coef_i(ij) + &
                                        target_theta_rhodz(ij, l, iq)*(1.-coef_i(ij))
            END DO
         END DO
      END DO
   END SUBROUTINE compute_guided_center

   SUBROUTINE compute_guided_center2(coef_i, target_theta_rhodz, theta_rhodz)
      REAL(rstd), INTENT(IN) :: coef_i(iim*jjm)
      REAL(rstd), INTENT(IN) :: target_theta_rhodz(iim*jjm, nqdyn)
      REAL(rstd), INTENT(INOUT) :: theta_rhodz(iim*jjm, nqdyn)
      INTEGER :: ij, iq
      DO iq = 1, nqdyn
         DO ij = ij_begin_ext, ij_end_ext
            theta_rhodz(ij, iq) = theta_rhodz(ij, iq)*coef_i(ij) + &
                                  target_theta_rhodz(ij, iq)*(1.-coef_i(ij))
         END DO
      END DO
   END SUBROUTINE compute_guided_center2

END MODULE nudging_mod
