MODULE etat0_mod
  USE icosa
  IMPLICIT NONE 	
  PRIVATE

    CHARACTER(len=255),SAVE :: etat0_type
!$OMP THREADPRIVATE(etat0_type)

    REAL(rstd) :: etat0_temp

    PUBLIC :: etat0, etat0_type

CONTAINS
  
  SUBROUTINE etat0(f_ps,f_mass,f_phis,f_theta_rhodz,f_u, f_q)
    USE mpipara, ONLY : is_mpi_root
    USE disvert_mod
    ! Generic interface
    USE etat0_dcmip1_mod, ONLY : getin_etat0_dcmip1=>getin_etat0
    USE etat0_dcmip2_mod, ONLY : getin_etat0_dcmip2=>getin_etat0
    USE etat0_dcmip4_mod, ONLY : getin_etat0_dcmip4=>getin_etat0
    USE etat0_dcmip5_mod, ONLY : getin_etat0_dcmip5=>getin_etat0
    USE etat0_williamson_mod, ONLY : getin_etat0_williamson=>getin_etat0
    USE etat0_temperature_mod, ONLY: getin_etat0_temperature=>getin_etat0
    ! Ad hoc interfaces
    USE etat0_academic_mod, ONLY : etat0_academic=>etat0  
    USE etat0_heldsz_mod, ONLY : etat0_heldsz=>etat0  
    USE etat0_venus_mod,  ONLY : etat0_venus=>etat0  
    USE etat0_start_file_mod, ONLY : etat0_start_file=>etat0  

    IMPLICIT NONE
    TYPE(t_field),POINTER :: f_ps(:)
    TYPE(t_field),POINTER :: f_mass(:)
    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 :: ps(:), mass(:,:)
    LOGICAL :: init_mass, collocated
    INTEGER :: ind,i,j,ij,l

    ! most etat0 routines set ps and not mass
    ! in that case and if caldyn_eta == eta_lag
    ! the initial distribution of mass is taken to be the same
    ! as what the mass coordinate would dictate
    ! however if etat0_XXX defines mass then the flag init_mass must be set to .FALSE.
    ! otherwise mass will be overwritten 
    init_mass = (caldyn_eta == eta_lag)

    etat0_type='jablonowsky06'
    CALL getin("etat0",etat0_type)
    
    !------------------- Generic interface ---------------------
    collocated=.TRUE.
    SELECT CASE (TRIM(etat0_type))
    CASE ('isothermal')
       CALL getin_etat0_isothermal
    CASE ('temperature_profile')
       CALL getin_etat0_temperature
    CASE ('jablonowsky06')
    CASE ('dcmip1')
        CALL getin_etat0_dcmip1
    CASE ('dcmip2_mountain','dcmip2_schaer_noshear','dcmip2_schaer_shear')
       CALL getin_etat0_dcmip2
    CASE ('dcmip3')
    CASE ('dcmip4')
        CALL getin_etat0_dcmip4
    CASE ('dcmip5')
        CALL getin_etat0_dcmip5
    CASE ('williamson91.6')
       init_mass=.FALSE.
       CALL getin_etat0_williamson
    CASE DEFAULT
       collocated=.FALSE.
    END SELECT

    !------------------- Ad hoc interfaces --------------------
    SELECT CASE (TRIM(etat0_type))
    CASE ('start_file')
       CALL etat0_start_file(f_ps,f_phis,f_theta_rhodz,f_u, f_q)
    CASE ('academic')
       CALL etat0_academic(f_ps,f_phis,f_theta_rhodz,f_u, f_q)
    CASE ('held_suarez')
       PRINT *,"Held & Suarez (1994) test case"
       CALL etat0_heldsz(f_ps,f_phis,f_theta_rhodz,f_u, f_q)
    CASE ('venus')
       CALL etat0_venus(f_ps, f_phis, f_theta_rhodz, f_u, f_q)
       PRINT *, "Venus (Lebonnois et al., 2012) test case"
   CASE DEFAULT
      IF(collocated) THEN
         CALL etat0_collocated(f_phis,f_ps,f_mass,f_theta_rhodz,f_u, f_q)
      ELSE
         PRINT*, 'Bad selector for variable etat0 <',etat0_type, &
              '> options are <jablonowsky06>, <academic>, <dcmip[1-4]> '
         STOP
      END IF
    END SELECT

    IF(init_mass) THEN ! initialize mass distribution using ps
!       !$OMP BARRIER
       DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          mass=f_mass(ind); ps=f_ps(ind)
          CALL compute_rhodz(.TRUE., ps, mass)
       END DO
    END IF

  END SUBROUTINE etat0

  SUBROUTINE etat0_collocated(f_phis,f_ps,f_mass,f_theta_rhodz,f_u, f_q)
    USE theta2theta_rhodz_mod
    IMPLICIT NONE
    TYPE(t_field),POINTER :: f_ps(:)
    TYPE(t_field),POINTER :: f_mass(:)
    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_temp(:)
    REAL(rstd),POINTER :: ps(:)
    REAL(rstd),POINTER :: mass(:,:)
    REAL(rstd),POINTER :: phis(:)
    REAL(rstd),POINTER :: theta_rhodz(:,:)
    REAL(rstd),POINTER :: temp(:,:)
    REAL(rstd),POINTER :: u(:,:)
    REAL(rstd),POINTER :: q(:,:,:)
    INTEGER :: ind

    CALL allocate_field(f_temp,field_t,type_real,llm,name='temp')

    DO ind=1,ndomain
      IF (.NOT. assigned_domain(ind)) CYCLE
      CALL swap_dimensions(ind)
      CALL swap_geometry(ind)
      ps=f_ps(ind)
      mass=f_mass(ind)
      phis=f_phis(ind)
      theta_rhodz=f_theta_rhodz(ind)
      temp=f_temp(ind)
      u=f_u(ind)
      q=f_q(ind)

      IF( TRIM(etat0_type)=='williamson91.6' ) THEN 
       CALL compute_etat0_collocated(ps,mass, phis, theta_rhodz, u, q)
      ELSE
       CALL compute_etat0_collocated(ps,mass, phis, temp, u, q)
      ENDIF
    ENDDO
    
    IF( TRIM(etat0_type)/='williamson91.6' ) CALL temperature2theta_rhodz(f_ps,f_temp,f_theta_rhodz)
    
    CALL deallocate_field(f_temp)
    
  END SUBROUTINE etat0_collocated

  SUBROUTINE compute_etat0_collocated(ps,mass, phis, temp_i, u, q)
    USE wind_mod
    USE etat0_jablonowsky06_mod, ONLY : compute_jablonowsky06 => compute_etat0
    USE etat0_dcmip1_mod, ONLY : compute_dcmip1 => compute_etat0
    USE etat0_dcmip2_mod, ONLY : compute_dcmip2 => compute_etat0
    USE etat0_dcmip3_mod, ONLY : compute_dcmip3 => compute_etat0
    USE etat0_dcmip4_mod, ONLY : compute_dcmip4 => compute_etat0
    USE etat0_dcmip5_mod, ONLY : compute_dcmip5 => compute_etat0
    USE etat0_williamson_mod, ONLY : compute_w91_6 => compute_etat0
    USE etat0_temperature_mod, ONLY: compute_etat0_temperature => compute_etat0
    IMPLICIT NONE
    REAL(rstd),INTENT(INOUT) :: ps(iim*jjm)
    REAL(rstd),INTENT(INOUT) :: mass(iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: phis(iim*jjm)
    REAL(rstd),INTENT(OUT) :: temp_i(iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: u(3*iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: q(iim*jjm,llm,nqtot)

    REAL(rstd) :: ulon_i(iim*jjm,llm)
    REAL(rstd) :: ulat_i(iim*jjm,llm)

    REAL(rstd) :: ps_e(3*iim*jjm)
    REAL(rstd) :: mass_e(3*iim*jjm,llm)
    REAL(rstd) :: phis_e(3*iim*jjm)
    REAL(rstd) :: temp_e(3*iim*jjm,llm)
    REAL(rstd) :: ulon_e(3*iim*jjm,llm)
    REAL(rstd) :: ulat_e(3*iim*jjm,llm)
    REAL(rstd) :: q_e(3*iim*jjm,llm,nqtot)

    INTEGER :: l,i,j,ij

    SELECT CASE (TRIM(etat0_type))
    CASE ('isothermal')
       CALL compute_etat0_isothermal(iim*jjm, phis, ps, temp_i, ulon_i, ulat_i, q)
       CALL compute_etat0_isothermal(3*iim*jjm, phis_e, ps_e, temp_e, ulon_e, ulat_e, q_e)
    CASE ('temperature_profile')
       CALL compute_etat0_temperature(iim*jjm, phis, ps, temp_i, ulon_i, ulat_i, q)
       CALL compute_etat0_temperature(3*iim*jjm, phis_e, ps_e, temp_e, ulon_e, ulat_e, q_e)
    CASE('jablonowsky06')
       CALL compute_jablonowsky06(iim*jjm,lon_i,lat_i, phis, ps, temp_i, ulon_i, ulat_i)
       CALL compute_jablonowsky06(3*iim*jjm,lon_e,lat_e, phis_e, ps_e, temp_e, ulon_e, ulat_e)
    CASE('dcmip1')
       CALL compute_dcmip1(iim*jjm,lon_i,lat_i, phis, ps, temp_i, ulon_i, ulat_i, q)
       CALL compute_dcmip1(3*iim*jjm,lon_e,lat_e, phis_e, ps_e, temp_e, ulon_e, ulat_e, q_e)
    CASE ('dcmip2_mountain','dcmip2_schaer_noshear','dcmip2_schaer_shear')
       CALL compute_dcmip2(iim*jjm,lon_i,lat_i, phis, ps, temp_i, ulon_i, ulat_i)
       CALL compute_dcmip2(3*iim*jjm,lon_e,lat_e, phis_e, ps_e, temp_e, ulon_e, ulat_e)      
    CASE('dcmip3')
       CALL compute_dcmip3(iim*jjm,lon_i,lat_i, phis, ps, temp_i, ulon_i, ulat_i, q)
       CALL compute_dcmip3(3*iim*jjm,lon_e,lat_e, phis_e, ps_e, temp_e, ulon_e, ulat_e, q_e)
    CASE('dcmip4')
       CALL compute_dcmip4(iim*jjm,lon_i,lat_i, phis, ps, temp_i, ulon_i, ulat_i, q)
       CALL compute_dcmip4(3*iim*jjm,lon_e,lat_e, phis_e, ps_e, temp_e, ulon_e, ulat_e, q_e)
    CASE('dcmip5')
       CALL compute_dcmip5(iim*jjm,lon_i,lat_i, phis, ps, temp_i, ulon_i, ulat_i, q)
       CALL compute_dcmip5(3*iim*jjm,lon_e,lat_e, phis_e, ps_e, temp_e, ulon_e, ulat_e, q_e)
    CASE('williamson91.6')
       CALL compute_w91_6(iim*jjm,lon_i,lat_i, phis, mass(:,1), temp_i(:,1), ulon_i(:,1), ulat_i(:,1))
       CALL compute_w91_6(3*iim*jjm,lon_e,lat_e, phis_e, mass_e(:,1), temp_e(:,1), ulon_e(:,1), ulat_e(:,1))
    END SELECT

    CALL compute_wind_perp_from_lonlat_compound(ulon_e, ulat_e, u)

  END SUBROUTINE compute_etat0_collocated

!----------------------------- Resting isothermal state --------------------------------

  SUBROUTINE getin_etat0_isothermal
    etat0_temp=300
    CALL getin("etat0_isothermal_temp",etat0_temp)
  END SUBROUTINE getin_etat0_isothermal

  SUBROUTINE compute_etat0_isothermal(ngrid, phis, ps, temp, ulon, ulat, q)
    IMPLICIT NONE  
    INTEGER, INTENT(IN)    :: ngrid
    REAL(rstd),INTENT(OUT) :: phis(ngrid)
    REAL(rstd),INTENT(OUT) :: ps(ngrid)
    REAL(rstd),INTENT(OUT) :: temp(ngrid,llm)
    REAL(rstd),INTENT(OUT) :: ulon(ngrid,llm)
    REAL(rstd),INTENT(OUT) :: ulat(ngrid,llm)
    REAL(rstd),INTENT(OUT) :: q(ngrid,llm,nqtot)
    phis(:)=0
    ps(:)=preff
    temp(:,:)=etat0_temp
    ulon(:,:)=0
    ulat(:,:)=0
    q(:,:,:)=0
  END SUBROUTINE compute_etat0_isothermal

END MODULE etat0_mod