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

    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_mod
    USE transfert_mod
    USE time_mod, ONLY : dt ! time step for a full RK step
    USE vertical_remap_mod
    USE compute_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)
    ENDIF


    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