MODULE nudging_mod
  USE icosa, ONLY : rstd
  USE grid_param, ONLY : llm, nqdyn
  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, time
  !$OMP THREADPRIVATE(center_lon, center_lat, nudging_radius, time)

  TYPE(t_field),POINTER :: f_relax_coef_e(:), f_target_ue(:), &
       f_relax_coef_i(:), f_target_theta_rhodz(:)

  PUBLIC :: init_guided, guided
  
CONTAINS

  SUBROUTINE init_guided(f_u,f_theta_rhodz)
    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
    REAL(rstd), POINTER :: ue(:,:), target_ue(:,:), coef_e(:) 
    REAL(rstd), POINTER :: theta_rhodz(:,:,:), target_theta_rhodz(:,:,:), coef_i(:) 
    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
!    time=0.
!    CALL getin('nudging_time', time)

    ! we should check that radius>0

    !SELECT CASE(TRIM(nudg_name))
    !    CASE ('f_u')
        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')
    !    CASE ('f_theta_rhodz')
        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')
    !    CASE DEFAULT
    !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)
            CALL compute_relax_coef(coef_e, coef_i)
            CALL compute_target_u(ue, target_ue)
            CALL compute_target_center(theta_rhodz, target_theta_rhodz)
    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
    REAL(rstd), INTENT(IN) :: lon,lat
    REAL(rstd) :: relax_coef, dist, c
    ! NB : dist is computed on unit sphere
    CALL dist_lonlat(lon, lat, center_lon, center_lat, dist) 
    c = tanh((1.-radius*dist/nudging_radius)*20.) ! 1 inside circle, -1 outside
    relax_coef = .5*(1.+c) ! rescale to [0,1] range ; ! 1 inside circle, 0 outside
  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

!----------------------------- Relax towards target ------------------------------
  
  SUBROUTINE guided(tt, f_ps, f_theta_rhodz, f_u, f_q)
    REAL(rstd), INTENT(IN):: tt
    TYPE(t_field),POINTER :: f_ps(:)
    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(:)  
    REAL(rstd), POINTER :: target_ue(:,:), ue(:,:), coef_e(:)
    REAL(rstd), POINTER :: target_theta_rhodz(:,:,:), theta_rhodz(:,:,:), coef_i(:) 
    INTEGER :: ind

    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)
       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)
    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
  
END MODULE nudging_mod