MODULE etat0_mod
  USE icosa
  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
    USE etat0_williamson_mod, ONLY : etat0_williamson_new
    USE etat0_jablonowsky06_mod, ONLY : etat0_jablonowsky06=>etat0
    USE etat0_academic_mod, ONLY : etat0_academic=>etat0  
    USE etat0_dcmip1_mod, ONLY : etat0_dcmip1=>etat0
    USE etat0_dcmip2_mod, ONLY : etat0_dcmip2=>etat0
    USE etat0_dcmip3_mod, ONLY : etat0_dcmip3=>etat0  
    USE etat0_dcmip4_mod, ONLY : etat0_dcmip4=>etat0  
    USE etat0_dcmip5_mod, ONLY : etat0_dcmip5=>etat0  
    USE etat0_heldsz_mod, ONLY : etat0_heldsz=>etat0  
    USE dynetat0_gcm_mod, ONLY : dynetat0_start=>etat0 
    USE dynetat0_hz_mod,  ONLY : dynetat0_hz=>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
    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)
    
    SELECT CASE (TRIM(etat0_type))
    CASE ('isothermal')
       CALL getin_etat0_isothermal
       CALL etat0_collocated(f_phis,f_ps,f_mass,f_theta_rhodz,f_u, f_q)
    CASE ('williamson91.6')
       init_mass=.FALSE.
       CALL etat0_williamson_new(f_phis,f_mass,f_theta_rhodz,f_u, f_q)
    CASE ('jablonowsky06')
!       CALL etat0_jablonowsky06(f_ps,f_phis,f_theta_rhodz,f_u, f_q)
       CALL etat0_collocated(f_phis,f_ps,f_mass,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 ('dcmip1')
       CALL etat0_dcmip1(f_ps,f_phis,f_theta_rhodz,f_u, f_q)
    CASE ('dcmip2_mountain','dcmip2_schaer_noshear','dcmip2_schaer_shear')
       CALL etat0_dcmip2(f_ps,f_phis,f_theta_rhodz,f_u, f_q)
    CASE ('dcmip3')
       CALL etat0_dcmip3(f_ps,f_phis,f_theta_rhodz,f_u, f_q)
     CASE ('dcmip4')
        IF(nqtot<2) THEN
           IF (is_mpi_root)  THEN
              PRINT *, "nqtot must be at least 2 for test case DCMIP4"
           END IF
           STOP
        END IF
       CALL etat0_dcmip4(f_ps,f_phis,f_theta_rhodz,f_u, f_q)
     CASE ('dcmip5')
       CALL etat0_dcmip5(f_ps,f_phis,f_theta_rhodz,f_u, f_q)
     CASE ('readnf_start') 
	  print*,"readnf_start used"	
       CALL dynetat0_start(f_ps,f_phis,f_theta_rhodz,f_u,f_q) 
	CASE ('readnf_hz') 
	  print*,"readnf_hz used"
       CALL dynetat0_hz(f_ps,f_phis,f_theta_rhodz,f_u,f_q) 
   CASE DEFAULT
       PRINT*, 'Bad selector for variable etat0 <',etat0_type, &
            '> options are <jablonowsky06>, <academic>, <dcmip[1-4]> '
       STOP
    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 mpipara
    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(:)
    REAL(rstd),POINTER :: mass(:,:)
    REAL(rstd),POINTER :: phis(:)
    REAL(rstd),POINTER :: theta_rhodz(:,:)
    REAL(rstd),POINTER :: u(:,:)
    REAL(rstd),POINTER :: q(:,:,:)
    INTEGER :: ind

    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)
      u=f_u(ind)
      q=f_q(ind)
      CALL compute_etat0_collocated(ps,mass, phis, theta_rhodz, u, q)

    ENDDO
  END SUBROUTINE etat0_collocated

  SUBROUTINE compute_etat0_collocated(ps,mass, phis, theta_rhodz, u, q)
    USE disvert_mod
    USE theta2theta_rhodz_mod
    USE wind_mod
    USE etat0_jablonowsky06_mod, ONLY : compute_jablonowsky06 => compute_etat0_new
    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) :: theta_rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: u(3*iim*jjm,llm)
    REAL(rstd),INTENT(OUT) :: q(iim*jjm,llm,nqtot)

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

    REAL(rstd) :: lon_e(3*iim*jjm)
    REAL(rstd) :: lat_e(3*iim*jjm)
    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

    DO l=1,llm
       DO j=jj_begin-1,jj_end+1
          DO i=ii_begin-1,ii_end+1
             ij=(j-1)*iim+i
             CALL xyz2lonlat(xyz_i(ij,:)/radius,lon_i(ij),lat_i(ij))
             CALL xyz2lonlat(xyz_e(ij+u_right,:)/radius,lon_e(ij+u_right),lat_e(ij+u_right))
             CALL xyz2lonlat(xyz_e(ij+u_lup,:)/radius,lon_e(ij+u_lup),lat_e(ij+u_lup))
             CALL xyz2lonlat(xyz_e(ij+u_ldown,:)/radius,lon_e(ij+u_ldown),lat_e(ij+u_ldown))
          END DO
       END DO
    END DO

    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('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)
    END SELECT

    CALL compute_temperature2theta_rhodz(ps,temp_i,theta_rhodz,0)       
    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