MODULE etat0_heldsz_mod
  USE icosa
  IMPLICIT NONE
  PRIVATE

  TYPE(t_field),POINTER :: f_theta_eq(:)
  TYPE(t_field),POINTER :: f_theta(:)

  REAL(rstd),ALLOCATABLE,SAVE :: knewt_t(:),kfrict(:)
!$OMP THREADPRIVATE(knewt_t,kfrict)
  LOGICAL, SAVE :: done=.FALSE.
!$OMP THREADPRIVATE(done)

  REAL(rstd),SAVE :: teta0,ttp,delt_y,delt_z,eps
!$OMP THREADPRIVATE(teta0,ttp,delt_y,delt_z,eps)

  REAL(rstd),SAVE :: knewt_g, k_f,k_c_a,k_c_s
!$OMP THREADPRIVATE(knewt_g, k_f,k_c_a,k_c_s)

  PUBLIC :: etat0, held_suarez
  
CONTAINS

  SUBROUTINE test_etat0_heldsz
    USE icosa
    USE kinetic_mod
    IMPLICIT NONE
    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(:)
    TYPE(t_field),POINTER :: f_Ki(:)

    REAL(rstd),POINTER :: Ki(:,:)
    INTEGER :: ind

    CALL allocate_field(f_ps,field_t,type_real)
    CALL allocate_field(f_phis,field_t,type_real)
    CALL allocate_field(f_theta_rhodz,field_t,type_real,llm)
    CALL allocate_field(f_u,field_u,type_real,llm)
    CALL allocate_field(f_Ki,field_t,type_real,llm)

    CALL etat0(f_ps,f_phis,f_theta_rhodz,f_u, f_q)
    CALL kinetic(f_u,f_Ki)

    CALL writefield('ps',f_ps)
    CALL writefield('theta',f_theta_rhodz)
  END SUBROUTINE test_etat0_heldsz

  SUBROUTINE etat0(f_ps,f_phis,f_theta_rhodz,f_u, f_q)
    USE icosa
    USE theta2theta_rhodz_mod
    IMPLICIT NONE
    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 :: ps(:)
    REAL(rstd),POINTER :: phis(:)
    REAL(rstd),POINTER :: theta_rhodz(:,:)
    REAL(rstd),POINTER :: u(:,:)
    REAL(rstd),POINTER :: q(:,:,:)
    REAL(rstd),POINTER :: theta_eq(:,:) 
    REAL(rstd),POINTER :: theta(:,:) 

    INTEGER :: ind

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

       theta_eq=f_theta_eq(ind) 
       CALL compute_Teq(lat_i,theta_eq) ! FIXME : already done by Init_Teq 

       ps=f_ps(ind)
       phis=f_phis(ind)
       u=f_u(ind)
       ps(:)=1e5
       phis(:)=0.
       u(:,:)=0

       theta_rhodz=f_theta_rhodz(ind)
       theta=f_theta(ind)
       CALL compute_etat0_heldsz(theta_eq,theta)
       CALL compute_theta2theta_rhodz(ps,theta,theta_rhodz,1)
       IF(nqtot>0) THEN
          q=f_q(ind)
          q(:,:,1)=1e-2
          IF(nqtot>1) q(:,:,2)=0
          IF(nqtot>2) q(:,:,3:)=1e-2
       END IF
    ENDDO
  END SUBROUTINE etat0

  SUBROUTINE init_Teq
    USE icosa
    USE disvert_mod, ONLY : ap,bp
    IMPLICIT NONE
    REAL(rstd),POINTER :: clat(:) 
    REAL(rstd),POINTER :: theta_eq(:,:) 
    REAL(rstd) :: zsig
    INTEGER :: ind, l

    IF(.NOT.done) THEN
       done = .TRUE.
       
       CALL allocate_field(f_theta,field_t,type_real,llm)
       CALL allocate_field(f_theta_eq,field_t,type_real,llm)
       ALLOCATE(knewt_t(llm)); ALLOCATE( kfrict(llm)) 

       k_f=1.                !friction 
       CALL getin('k_j',k_f)
       k_f=1./(daysec*k_f)
       k_c_s=4.  !cooling surface
       CALL getin('k_c_s',k_c_s)
       k_c_s=1./(daysec*k_c_s)
       k_c_a=40. !cooling free atm
       CALL getin('k_c_a',k_c_a)
       k_c_a=1./(daysec*k_c_a)
       ! Constants for Teta equilibrium profile
       teta0=315.     ! mean Teta (S.H. 315K)
       CALL getin('teta0',teta0)
       ttp=200.       ! Tropopause temperature (S.H. 200K)
       CALL getin('ttp',ttp)
       eps=0.         ! Deviation to N-S symmetry(~0-20K)
       CALL getin('eps',eps)
       delt_y=60.     ! Merid Temp. Gradient (S.H. 60K)
       CALL getin('delt_y',delt_y)
       delt_z=10.     ! Vertical Gradient (S.H. 10K)
       CALL getin('delt_z',delt_z)

       !-----------------------------------------------------------
       knewt_g=k_c_a 
       DO l=1,llm
          zsig=ap(l)/preff+bp(l)
          knewt_t(l)=(k_c_s-k_c_a)*MAX(0.,(zsig-0.7)/0.3) 
          kfrict(l)=k_f*MAX(0.,(zsig-0.7)/0.3) 
       ENDDO

       DO ind=1,ndomain
          IF (.NOT. assigned_domain(ind)) CYCLE
          CALL swap_dimensions(ind)
          CALL swap_geometry(ind)
          theta_eq=f_theta_eq(ind)
          CALL compute_Teq(lat_i,theta_eq)
       ENDDO

    ELSE
       PRINT *, 'Init_Teq called twice'
       CALL ABORT
    END IF

  END SUBROUTINE init_Teq

  SUBROUTINE compute_Teq(lat,theta_eq)
    USE icosa
    USE disvert_mod
    IMPLICIT NONE  
    REAL(rstd),INTENT(IN) :: lat(iim*jjm)
    REAL(rstd),INTENT(OUT) :: theta_eq(iim*jjm,llm) 

    REAL(rstd) :: r, zsig, ddsin, tetastrat, tetajl
    INTEGER :: i,j,l,ij

    DO l=1,llm
       zsig=ap(l)/preff+bp(l)
       tetastrat=ttp*zsig**(-kappa)
       DO j=jj_begin-1,jj_end+1
          DO i=ii_begin-1,ii_end+1
             ij=(j-1)*iim+i
	     ddsin=SIN(lat(ij)) 
	     tetajl=teta0-delt_y*ddsin*ddsin+eps*ddsin &
                  -delt_z*(1.-ddsin*ddsin)*log(zsig)
	     theta_eq(ij,l)=MAX(tetajl,tetastrat) 
          ENDDO
       ENDDO
    ENDDO
  END SUBROUTINE compute_Teq

  SUBROUTINE compute_etat0_heldsz(theta_eq, theta)
    USE icosa
    USE disvert_mod
    USE pression_mod
    USE exner_mod
    USE geopotential_mod
    USE theta2theta_rhodz_mod
    IMPLICIT NONE  
    REAL(rstd),INTENT(IN) :: theta_eq(iim*jjm,llm) 
    REAL(rstd),INTENT(OUT) :: theta(iim*jjm,llm) 

    REAL(rstd) :: r  ! random number
    INTEGER :: i,j,l,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 RANDOM_NUMBER(r); r = 0.0 
             theta(ij,l)=theta_eq(ij,l)*(1.+0.0005*r)
          ENDDO
       ENDDO
    ENDDO

  END SUBROUTINE compute_etat0_heldsz


  SUBROUTINE held_suarez(f_ps,f_theta_rhodz,f_u) 
    USE icosa
    IMPLICIT NONE
    TYPE(t_field),POINTER :: f_theta_rhodz(:)
    TYPE(t_field),POINTER :: f_u(:)
    TYPE(t_field),POINTER :: f_ps(:)
    REAL(rstd),POINTER :: theta_rhodz(:,:)
    REAL(rstd),POINTER :: u(:,:)
    REAL(rstd),POINTER :: ps(:)
    REAL(rstd),POINTER :: theta_eq(:,:)
    REAL(rstd),POINTER :: theta(:,:)
    REAL(rstd),POINTER :: clat(:)
    INTEGER::ind

    DO ind=1,ndomain
       IF (.NOT. assigned_domain(ind)) CYCLE
       CALL swap_dimensions(ind)
       CALL swap_geometry(ind)
       theta_rhodz=f_theta_rhodz(ind)
       u=f_u(ind)
       ps=f_ps(ind) 
       theta_eq=f_theta_eq(ind) 
       theta=f_theta(ind) 
       CALL compute_heldsz(ps,theta_eq,lat_i, theta_rhodz,u, theta) 
    ENDDO
  END SUBROUTINE held_suarez

  SUBROUTINE compute_heldsz(ps,theta_eq,lat, theta_rhodz,u, theta) 
    USE icosa
    USE theta2theta_rhodz_mod
    IMPLICIT NONE
    REAL(rstd),INTENT(IN)    :: ps(iim*jjm) 
    REAL(rstd),INTENT(IN)    :: theta_eq(iim*jjm,llm) 
    REAL(rstd),INTENT(IN)    :: lat(iim*jjm) 
    REAL(rstd),INTENT(INOUT) :: theta_rhodz(iim*jjm,llm)
    REAL(rstd),INTENT(INOUT) :: u(3*iim*jjm,llm)
    REAL(rstd),INTENT(OUT)   :: theta(iim*jjm,llm) 

    INTEGER :: i,j,l,ij

    CALL compute_theta_rhodz2theta(ps,theta_rhodz,theta,1)
    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
             theta(ij,l)=theta(ij,l) - dt*(theta(ij,l)-theta_eq(ij,l))* &
                  (knewt_g+knewt_t(l)*COS(lat(ij))**4 )
          ENDDO
       ENDDO
    ENDDO
    CALL compute_theta2theta_rhodz(ps,theta,theta_rhodz,1)

    Do l=1,llm
       u(:,l)=u(:,l)*(1.-dt*kfrict(l))
    END DO

  END SUBROUTINE compute_heldsz

END MODULE etat0_heldsz_mod