module matrix_math
	implicit none

	private
	public						::	    get_levi_civita,                        &
                                        crossP,                                 &
                                        zheevr_wrapper,                         &
                                        zheevd_wrapper,                         &
                                        zheevx_wrapper,                         &
                                        is_equal_vect,                          &
                                        is_equal_mat,                           &
                                        is_herm_mat,                            &
                                        is_skew_herm_mat,                       &
                                        convert_tens_to_vect,                   &
                                        blas_matmul,                            &
                                        matrix_comm,                            &
                                        get_linspace

    interface crossP
        module procedure   real_crossP
        module procedure   cplx_crossP
    end interface crossP

    interface  is_equal_vect
        module procedure    real_is_equal_vect
        module procedure    cplx_is_equal_vect
    end interface is_equal_vect

    interface is_equal_mat
        module procedure d_is_equal_mat
        module procedure z_is_equal_mat
    end interface is_equal_mat

    interface convert_tens_to_vect
        module procedure    real_tens_to_vect
        module procedure    cplx_tens_to_vect
    end interface convert_tens_to_vect

    interface blas_matmul
        module procedure    d_blas_matmul
        module procedure    z_blas_matmul
    end interface blas_matmul

    interface is_herm_mat
        module procedure    z_is_herm_mat
    end interface is_herm_mat

    interface is_skew_herm_mat
        module procedure    z_is_skew_herm_mat
    end interface is_skew_herm_mat

    interface matrix_comm
        module procedure    d_matrix_comm
        module procedure    z_matrix_comm
    end interface matrix_comm

    save

    integer,        parameter   ::  dp              = kind(0.d0)

	contains

!public:
    pure subroutine get_linspace(min, max, n_mesh, linspace)
        real(dp),                      intent(in)      ::   min, max
        integer,                       intent(in)      ::   n_mesh
        real(dp),   allocatable,       intent(out)     ::   linspace(:)
        real(dp)                                       ::   delta
        integer                                        ::   idx
        !
        !
        if( n_mesh > 1) then
            !
            !   RETURN LIST
            allocate(   linspace(n_mesh)  )
            delta   =   (   max -   min)    /   real(n_mesh-1,dp)
            !
            do idx  =   1, n_mesh
                linspace(idx)   =       min     +  real(idx-1,dp)    * delta
            end do
        else
            !
            !   RETURN SCALAR
            allocate(   linspace(1) )
            linspace    =   min
        end if
        !
        !
        return
    end subroutine

    integer pure function my_Levi_Civita(i,j,k)
        !Hard coded Levi Civita tensor
        integer,        intent(in)      :: i,j,k
        logical                         :: even, odd
        !
        !
        even    = (i==1 .and. j==2 .and. k==3) .or. (i==2 .and. j==3 .and. k==1) .or. (i==3 .and. j==1 .and. k==2)
        odd     = (i==3 .and. j==2 .and. k==1) .or. (i==1 .and. j==3 .and. k==2) .or. (i==2 .and. j==1 .and. k==3)
        !
        if(even)        then
                                my_Levi_Civita  =  1
        else if(odd)    then
                                my_Levi_Civita  = -1
        else
                                my_Levi_Civita  =  0
        end if
        !
        return
    end function

    pure subroutine get_levi_civita(lc_tens)
        integer,    intent(out)   ::  lc_tens(3,3,3)
        !
        lc_tens         =    0
        !
        lc_tens(1,2,3)  =   +1
        lc_tens(2,3,1)  =   +1
        lc_tens(3,1,2)  =   +1
        !
        lc_tens(1,3,2) =   -1
        lc_tens(3,2,1) =   -1
        lc_tens(2,1,3) =   -1
        !
        return
    end subroutine

	subroutine zheevr_wrapper(a, w ,z, m)
    	!https://software.intel.com/en-us/mkl-developer-reference-fortran-heevr#6ADF761A-127A-4C9B-9A2A-1A8AA4602CE1
    	!with ability to solve only selected eigenvalues (& vectors)
    	complex(dp),	intent(inout)			:: 	a(:,:)				!Hermitian matrix to be solved
    	real(dp),		intent(out)				:: 	w(:)				!selected eigenvalues in ascending order
    	complex(dp),	intent(out)				:: 	z(:,:)				!each column represents an eigenvector (	columns 1,2,...,m	)
    	integer,		intent(out)				:: 	m					!number of eigenvalues found
    	character(len=1)	 					:: 	jobz, range, uplo
    	integer									:: 	n, lda, il, iu, ldz, lwork, lrwork, liwork,  info
    	integer,		allocatable				:: 	isuppz(:) , iwork(:)
    	real(dp)								:: 	vl, vu, abstol
    	real(dp),		allocatable				:: 	rwork(:)
    	complex(dp),	allocatable				:: 	work(:)
    	!
    	jobz	= 'V'			!also compute eigenvectors
    	range	= 'I'
    	uplo	= 'U'
    	n 		= size(a,1)
    	lda		= n
    	vl		= 0.0_dp
    	vu  	= 0.0_dp
    	il		= 1
    	iu		= size(z,2)		!number of eigenvalues to compute
    	abstol	= 1e-15_dp
    	ldz		= n
    	lwork	= 2*n
    	lrwork	= 24*n
    	liwork	= 10*n
    	if( size(z,1)/= ldz ) write(*,*)"[eigSolver2]: z array has wrong size"
    	if( size(w)/= n) write(*,*)"[eigSolver2]; w array has wrong size"
    	!
    	allocate( isuppz(2*iu)	)
    	allocate(  work(lwork)	)
    	allocate( rwork(lrwork)	)
    	allocate( iwork(liwork)	)
    	!
    	call zheevr(jobz, range, uplo, n, a, lda, vl, vu, il, iu, abstol, &
    					m, w, z, ldz, isuppz, work, lwork, rwork, lrwork, iwork, liwork, info)
    	!
    	call errCheck(n,info,jobz)
        if( info /= 0)  stop 'zheevr'
    	return
    end subroutine

    subroutine zheevd_wrapper(A, w)
		!solves standard (non general) eigenvalue problem with following mkl routine
		!return eigVectors stored in A, eigValues in w
		!https://software.intel.com/en-us/node/469182
		complex(dp) , intent(inout)   :: A(:,:)
		real(dp)    , intent(out)     :: w(:)
            !
            !set up workspace arrays (assuming jobz='V' and using zheevd)
            complex(dp), allocatable, dimension(:)   :: work
            real(dp)   , allocatable, dimension(:)   :: rwork
            integer    , allocatable, dimension(:)   :: iwork
            character(len=1)                         :: jobz,uplo
            integer                                  :: n, info,lwork,lrwork,liwork
            n		= size(A,1)
            if(n /= size(A,2)) then
            	write(*,*)"[eigSolver]: WARNING the matrix to solve is not a square matrix"
            end if
            lwork  	=   n*n + 2*n
            lrwork 	= 2*n*n + 5*n + 1
            liwork 	=         5*n + 3
            allocate(  work( lwork) )
            allocate( rwork(lrwork) )
            allocate( iwork(liwork) )
            !
            jobz='V'
            uplo='U' !test this behaviour
            !
            !solve eigenvalue problem
            call zheevd(jobz, uplo, n, A, n , w, work, lwork, rwork, lrwork, iwork, liwork, info)
            !
            !check if system was solved correctly
            call errCheck(n,info,jobz)
            if( info /= 0)  stop 'zheevd'
            return
	end subroutine

    subroutine zheevx_wrapper(A, w)
        complex(dp),        intent(inout)       ::  A(:,:)
        real(dp),           intent(in)          ::  w(:)
        real(dp)                                ::  vl, vu
        integer                                 ::  il, iu, lwork, lrwork, liwork, info, num_wann
        complex(dp),        allocatable         ::  z(:,:), work(:)
        real(dp),           allocatable         ::  rwork(:)
        integer,            allocatable         ::  iwork(:), ifail(:)
        !
        if (size(A,1)== size(A,2)) then
            num_wann    =   size(A,1)
        else
            stop "zheevx_wrapper was given non symmetric matrix!"
        end if
        !
        lwork=12*num_wann
        lrwork=17*num_wann
        liwork=15*num_wann
        allocate(   ifail(          num_wann            ))
        allocate(   z(             num_wann, num_wann   ))
        allocate( rwork(lrwork),    work(lwork),    iwork(liwork) )
        !
        !
        call zheevx('V', 'A', 'U', size(A,1), A(:,:), size(A,1), vl, vu, il, iu, 0.0_dp, size(w,1),&
                            w(:), z(:,:), size(A,1), work, lwork, rwork, iwork, ifail, info)
        !
        !   return eigenvectors
        A(:,:)  = z(:,:)
        !
        !   debug
        call errCheck(num_wann,info,'V')
        if( info /= 0)  stop 'zheevx'
        return
    end subroutine

!    subroutine uni_gauge_trafo(U_mat, M_mat)
!        !   does a gauge trafo of the gauge-covariant matrix M_mat
!        !   by applying the unitary matrix U_mat as shown in
!        !   Eq.(21), PRB 74, 195118, (2006)
!        !
!        !   since U is defined differently in the paper and in the lapack
!        !   consider the following
!        !
!        !   want:   M^(hbar)
!        !
!        !    paper:      U^*    H^(W)   U   = H^(H)
!        !           and M^(hbar) = U^* M^(W) U
!        !
!        !   lapack:
!        !            U  H^(W)   U^*     = H^(H)
!        !   therefore here M^(hbar) = U M^(W) U^*
!
!        complex(dp),        intent(in)      ::  U_mat(:,:)
!        complex(dp),        intent(inout)   ::  M_mat(:,:)
!        complex(dp),        allocatable     ::  U_dag(:,:)
!        !
!        allocate(U_dag(  size(U_mat,1),size(U_mat,2)  ))
!        U_dag   = conjg(    transpose( U_mat )  )
!        !
!        !       WORKING:
!       ! M_mat   =   blas_matmul(    M_mat,  U_mat   )
!       ! M_mat   =   blas_matmul(    U_dag,  M_mat   )
!        !
!        M_mat   =   blas_matmul(    blas_matmul(U_dag, M_mat),  U_mat)
!        !       NOT WORKING:
!        !M_mat   =   blas_matmul(    M_mat,  U_dag   )
!        !M_mat   =   blas_matmul(    U_mat,  M_mat   )
!        !
!        return
!    end subroutine

!----------------------------------------------------------------------------------------------------------------------
!----------------------------------------------------------------------------------------------------------------------
!------------INTERFACES------------------------------------------------------------------------------------------------
!
!----------------------------------------------------------------------------------------------------------------------
!
!
    pure function real_crossP(a,b)
        !cross product of two real 3dim vectors a,b
        real(dp), dimension(3)              :: real_crossP
        real(dp), dimension(3), intent(in)  :: a, b
        !
        real_crossP(1)   =   a(2) * b(3)     -   a(3) * b(2)
        real_crossP(2)   =   a(3) * b(1)     -   a(1) * b(3)
        real_crossP(3)   =   a(1) * b(2)     -   a(2) * b(1)
        !
        return
    end function

    pure function cplx_crossP(a,b)
        !cross product of two complex 3dim vectors a,b
        complex(dp)                 :: cplx_crossP(3)
        complex(dp),    intent(in)  :: a(3), b(3)
        !
        cplx_crossP(1)   =   a(2) * b(3)     -   a(3) * b(2)
        cplx_crossP(2)   =   a(3) * b(1)     -   a(1) * b(3)
        cplx_crossP(3)   =   a(1) * b(2)     -   a(2) * b(1)
        !
        return
    end function
!
!----------------------------------------------------------------------------------------------------------------------
!
!
    logical pure function real_is_equal_vect(acc, a,b)
        real(dp),     intent(in)     ::   acc
        real(dp),     intent(in)     ::   a(:), b(:)
        integer                      ::   idx
        !
        real_is_equal_vect  =   ( size(a,1) == size(b,1) )
        !
        if(real_is_equal_vect) then
            loop: do idx = 1, size(a,1)
                real_is_equal_vect =    real_is_equal_vect  .and.      abs(   a(idx) - b(idx)  )  < acc
                if(.not. real_is_equal_vect) exit loop
            end do loop
        end if
        !
        return
    end function

    logical pure function cplx_is_equal_vect(acc, a,b)
        real(dp),           intent(in)  ::   acc
        complex(dp),        intent(in)  ::   a(:), b(:)
        integer                         ::   idx
        !
        cplx_is_equal_vect = ( size(a,1) == size(b,1) )
        !
        if(cplx_is_equal_vect) then
            loop: do idx = 1, size(a,1)
                cplx_is_equal_vect =    cplx_is_equal_vect  .and.      abs(   a(idx) - b(idx)  )  < acc
                 if(.not. cplx_is_equal_vect) exit loop
            end do loop
        end if
        !
        return
    end function
!
!----------------------------------------------------------------------------------------------------------------------
!
!
    logical pure function  d_is_equal_mat(acc, A, B)
        real(dp),   intent(in)          ::  acc
        real(dp),   intent(in)          ::  A(:,:), B(:,:)
        real(dp)                        ::  delta
        integer                         ::  n, m
        !
        d_is_equal_mat  =   (   size(A,1) == size(B,1) .and. size(A,2) == size(B,2) )
        !
        if( d_is_equal_mat  ) then
            loop_out: do n = 1, size(A,2)
                do m = 1, size(A,1)
                    delta   =   abs(    A(m,n) - B(m,n) )
                    !
                    d_is_equal_mat  = d_is_equal_mat .and. (delta < acc)
                    !
                    if(.not. d_is_equal_mat) exit loop_out
                end do
            end do loop_out
        end if
        !
        return
    end function

    logical pure function z_is_equal_mat(acc, A, B)
        real(8),       intent(in)      ::  acc
        complex(8),    intent(in)      ::  A(:,:), B(:,:)
        real(8)                        ::  delta
        integer                         ::  n, m
        !
        z_is_equal_mat  =   (   size(A,1) == size(B,1) .and. size(A,2) == size(B,2) )
        !
        if( z_is_equal_mat  ) then
            loop_out: do n = 1, size(A,2)
                do m = 1, size(A,1)
                    delta   =   abs(    A(m,n) - B(m,n) )
                    !
                    z_is_equal_mat  = z_is_equal_mat .and. (delta < acc)
                    !
                    if(.not. z_is_equal_mat)    exit loop_out
                end do
            end do loop_out
        end if
        !
        return
    end function
!----------------------------------------------------------------------------------------------------------------------
!
!
    logical function z_is_herm_mat(H, max_err)
        complex(dp),    intent(in)      ::      H(:,:)
        real(dp),       intent(out)     ::      max_err
        real(dp)                        ::      dH_nm
        integer                         ::      n, m
        !
        max_err         =       -1.0_dp
        z_is_herm_mat   =       ( size(H,1) == size(H,2)  )
        !
        if(z_is_herm_mat)   then
            columns: do m = 1, size(H,2)
                do n = 1, size(H,1)
                    !
                    dH_nm   =  abs( H(n,m) - conjg(H(m,n)) )
                    !
                    if(  n/=m       .and.  dH_nm     >   1e-6_dp   ) then
                        z_is_herm_mat   = .false.
                        !write(*,*)  '[z_is_herm_mat]: herm. criterion violated by ', abs( H(n,m) - conjg(H(m,n)) )
                        !exit columns
                        if(dH_nm > max_err)     max_err =   dH_nm
                    end if
                    !
                    !
                end do
            end do columns
        else
            write(*,*)  '[z_is_herm_mat]: WARNING  matrix is non symmetric...'
        end if
        !
        return
    end function
!----------------------------------------------------------------------------------------------------------------------
!
!
    logical function z_is_skew_herm_mat(H, max_err)
        complex(dp),    intent(in)      ::      H(:,:)
        real(dp),       intent(out)     ::      max_err
        real(dp)                        ::      dH_nm
        integer                         ::      n, m
        !
        max_err              =      -1.0_dp
        z_is_skew_herm_mat   =      ( size(H,1) == size(H,2)  )
        !
        if(z_is_skew_herm_mat)   then
            columns: do m = 1, size(H,2)
                do n = 1, size(H,1)
                    !
                    dH_nm   =   abs( H(n,m) + conjg(H(m,n)) )
                    !
                    if(  n/=m       .and.   dH_nm     >   1e-6_dp   ) then
                        z_is_skew_herm_mat   = .false.
                        !exit columns
                        if(dH_nm > max_err)     max_err =   dH_nm
                    end if
                    !
                    !
                end do
            end do columns
        end if
        !
        return
    end function
!----------------------------------------------------------------------------------------------------------------------
!
!
    pure subroutine real_tens_to_vect(tens, vect)
        !   converts om_tens to a vector by applying Levi Cevita Tensor
        !   see PRB 74, 195118 (2006)   Eq.(5)
        real(dp),           intent(in)      ::  tens(3,3)
        real(dp),           intent(out)     ::  vect(3)
        integer                             ::  a, b, c
        !
        vect =  0.0_dp
        do c = 1, 3
            !
            do b = 1, 3
                do a = 1,3
                    vect(c) = vect(c) + real(my_Levi_Civita(a,b,c),dp) * tens(a,b)
                end do
            end do
            !
        end do
        return
    end subroutine

    pure subroutine cplx_tens_to_vect(tens, vect)
        !   converts om_tens to a vector by applying Levi Cevita Tensor
        !   see PRB 74, 195118 (2006)   Eq.(5)
        complex(dp),        intent(in)      ::  tens(3,3)
        complex(dp),        intent(out)     ::  vect(3)
        integer                             ::  a, b, c
        !
        vect =   cmplx(0.0_dp,0.0_dp, dp)
        do c = 1, 3
            !
            do b = 1, 3
                do a = 1,3
                    vect(c) = vect(c) + real(my_Levi_Civita(a,b,c),dp) * tens(a,b)
                end do
            end do
            !
        end do
        return
    end subroutine
!----------------------------------------------------------------------------------------------------------------------
!
!
    function d_matrix_comm(A, B) result(C)
        !   calculates the commutator of matrix A with matrix B
        !       C   = [A,B] =   AB - BA
        real(dp),    intent(in)      ::  A(:,:),     B(:,:)
        real(dp),    allocatable     ::  C(:,:),     tmp(:,:)
        !
        if(size(A,2) /= size(B,1)   .or. size(B,2) /= size(A,1) ) stop "real_matrix_comm: mat sizes not matching"
        !
        !
        allocate(   C(  size(A,1),size(B,2))     )
        allocate(   tmp(size(A,1),size(B,2))     )

        !
        !call blas_matmul(A, B,      C   )
        C   =   blas_matmul(A,B)
        !call blas_matmul(B, A,      tmp )
        tmp =   blas_matmul(B,A)
        C   = C - tmp
        !
        return
    end function

    function z_matrix_comm(A,  B) result(C)
        !   calculates the commutator of matrix A with matrix B
        !       C   = [A,B] =   AB - BA
        complex(dp),    intent(in)      ::  A(:,:),     B(:,:)
        complex(dp),    allocatable     ::  C(:,:),     tmp(:,:)
        !
        if(size(A,2) /= size(B,1)   .or. size(B,2) /= size(A,1) ) stop "real_matrix_comm: mat sizes not matching"
        !
        !
        allocate(   C(size(A,1),size(B,2))     )
        allocate(   tmp(size(A,1),size(B,2))   )
        !
        !call blas_matmul(A, B,      C   )
        C   =   blas_matmul(A,B)
        !call blas_matmul(B, A,      tmp )
        tmp   =   blas_matmul(B,A)
        C   = C - tmp
        !
        return
    end function
!----------------------------------------------------------------------------------------------------------------------

!private helpers
    subroutine errCheck(n,info, jobz)
        !small subroutine used in subroutine eigSolver
        implicit none
        integer    , intent(in)        :: n,info
        character(len=1), intent(in)   :: jobz
        !

        if(info > 0) then
              write(*,*) '[solver/errCheck]: WARNING, Problem solving the eigenvalue problem: '
              if(jobz .EQ. 'N') write(*,*) '[solver/errCheck]: the algorithm failed to converge; ', info ,&
                            ' off-diagonal elements of an intermediate tridiagonal form did not converge to zero;'
              if(jobz .EQ. 'V') write(*,*) '[solver/errCheck]: the algorithm failed to compute an',&
            ' eigenvalue while working on the submatrix lying in rows and columns ', info/(n+1), ' through ', mod(info, n+1)
        elseif(info < 0) then
             write(*,*) '[solver/errCheck]: the ',info,'-th parameter had an illegal value.'
        endif
        !
        return
    end subroutine

    function d_blas_matmul(A, B)   result(C)
        !   C := alpha*op(A)*op(B) + beta*C
        real(dp),        intent(in)         ::  A(:,:), B(:,:)
        real(dp),   allocatable             ::  C(:,:)
        integer                             ::  m, n, k
        !
        m       =   size(A,1)
        n       =   size(B,2)
        k       =   size(A,2)
        if( k /= size(B,1))   stop  '[d_blas_matmul]:   warning matrix size do not match correctly (k<->size(B,1))'
        if( m /= size(C,1))   stop  '[d_blas_matmul]:   warning matrix size do not match correctly (m<->size(C,1))'
        if( n /= size(C,2))   stop  '[d_blas_matmul]:   warning matrix size do not match correctly (n<->size(C,2))'
        !
        allocate(   C(m,n)     )
        C       =   0.0_dp
        !
        !
        !        CALL DGEMM('N','N',M,N,K,ALPHA,A,M,B,K,BETA,C,M)
        call dgemm('N', 'N', m, n, k, 1.0_dp, A(:,:), m, B(:,:), k, 0.0_dp, C(:,:), m)
        !
        !
        return
    end function

    function z_blas_matmul(A, B)     result(C)
        !   C := alpha*op(A)*op(B) + beta*C
        complex(dp),        intent(in)      ::  A(:,:), B(:,:)
        complex(dp),        allocatable     ::  C(:,:)
        integer                             ::  m, n, k
        !
        m       =   size(A,1)
        k       =   size(A,2)
        n       =   size(B,2)
        if( k /= size(B,1))   stop  '[z_blas_matmul]:   warning matrix size do not match correctly (k<->size(B,1))'
        !
        !
        allocate(   C(m,n)     )
        C      =    cmplx(0.0_dp,0.0_dp,dp)
        !
        call zgemm('N', 'N', m, n, k, 1.0_dp, A(:,:), m,  B(:,:), k, 0.0_dp,    C(:,:), m)
        !
        !
        return
    end function

end module

!    subroutine normalCheck(n,eigVec)
!       !check if eigVec are orthonormal this should always be the case
!       integer     , intent(in)      :: n
!       complex(dp) , intent(in)      :: eigVec(:,:)
!       integer                       :: i1,i2
!       real(dp)                      :: tmp
!       do i2=1,n
!             tmp = 0.0_dp
!             do i1=1,n
!                   tmp = tmp + dconjg(eigVec(i1,i2))*eigVec(i1,i2)
!             end do
!             if(abs(tmp)-1.0_dp>1e-12_dp) then
!                   write(*,*)"[solver/normalCheck]: WARNING eigVec not normal: ",dsqrt(tmp)
!             end if
!       end do
!       !
!       return
!   end subroutine


      SUBROUTINE elmhes(a,n,np)
      INTEGER n,np
      REAL a(np,np)
      INTEGER i,j,m
      REAL x,y
      do 17 m=2,n-1
        x=0.
        i=m
        do 11 j=m,n
          if(abs(a(j,m-1)).gt.abs(x))then
            x=a(j,m-1)
            i=j
          endif
11      continue
        if(i.ne.m)then
          do 12 j=m-1,n
            y=a(i,j)
            a(i,j)=a(m,j)
            a(m,j)=y
12        continue
          do 13 j=1,n
            y=a(j,i)
            a(j,i)=a(j,m)
            a(j,m)=y
13        continue
        endif
        if(x.ne.0.)then
          do 16 i=m+1,n
            y=a(i,m-1)
            if(y.ne.0.)then
              y=y/x
              a(i,m-1)=y
              do 14 j=m,n
                a(i,j)=a(i,j)-y*a(m,j)
14            continue
              do 15 j=1,n
                a(j,m)=a(j,m)+y*a(j,i)
15            continue
            endif
16        continue
        endif
17    continue
      return
      END

NR-Functions/Numerical Recipes- Example & Functions/Functions/elmhes.for


! { dg-do run }

program test
  implicit none

  interface check
    procedure check_r4
    procedure check_r8
  end interface check

  real(kind=4) :: x4
  real(kind=8) :: x8

  x8 = 1.9_8 ; x4 = 1.9_4

  call check(erfc_scaled(x8), erfc_scaled(1.9_8))
  call check(erfc_scaled(x4), erfc_scaled(1.9_4))

contains
  subroutine check_r4 (a, b)
    real(kind=4), intent(in) :: a, b
    if (abs(a - b) > 1.e-5 * abs(b)) call abort
  end subroutine
  subroutine check_r8 (a, b)
    real(kind=8), intent(in) :: a, b
    if (abs(a - b) > 1.e-7 * abs(b)) call abort
  end subroutine
end program test


subroutine young_frame(u,v)
implicit none

integer*8 u,v
integer*8 i,j,k,l
integer*8 ii,jj,m,kk
integer*8 a(u),b(v)
integer*8 hka(u),hkb(v),f_lam
integer*16 h1,h2
integer*8 a1(u),b1(v),a11(u),b11(v)
integer*8 flag,flag2
integer*8 pp
integer*8 alpha
integer*8, dimension(:), allocatable  :: comb
integer*8, dimension(:,:), allocatable  :: rem
integer*8, dimension(:,:), allocatable  :: remb
integer*8, dimension(:,:), allocatable :: young
integer AllocateStatus
integer*8 v1,num
integer*16  tabl(15)

! a = first row of young frame
! b = second row of young frame
! u = size (a)
! v = size (b)
! hka = hooklength of first row
! hkb = hooklength of second row
! f_lam = no. of standard young tabulae

print*,"u and v", u,v
do i = 1,u
  a(i) = i
end do

do j=1,v
  b(j) = u+j
end do

! Computing number of standard yound tabuleaux using
! hook length formula f(lam) = N! /(h1*h2* ... hn)

! hook lengths are stored in hka and hkb

if (v .ne. 0) then

  do i=1,v
    hka(i) = 2 + (u-i)
  end do

  do i=v+1,u
   hka(i) = 1 + (u-i)
  end do

 else

  do i=1,u
    hka(i) = 1 + (u-i)
  end do

end if

do j=1,v
 hkb(j) = 1 + (v-j)
end do

!write(22,*) "hook lengths"

!write(22,*) hka
!write(22,*) hkb

! computing number of standard young tabulaex

h1=1
do i=1,u
  h1=h1*hka(i)
end do

h2=1
do j=1,v
  h2=h2*hkb(j)
end do

if (u+v.le.20) then
f_lam = Factorial(u+v)/(h1*h2)

print*,"Fatcorial of ", u+v,"=", Factorial (u+v)
print*, "No. of standard yound tabulae = ", f_lam
write(22,*), "No. of standard yound tabulae = ", f_lam

end if

do i=1,13
read(51,*) num,tabl(i)
end do

if (u+v.gt.20) then
f_lam = tabl(u+v-20)/(h1*h2)
print*,"Fatcorial of ", u+v,"=", tabl(u+v-20)
print*, "No. of standard yound tabulae = ", f_lam
write(22,*), "No. of standard yound tabulae = ", f_lam
end if

!allocate (young(f_lam,u+v))

allocate ( young(f_lam,u+v), STAT = AllocateStatus)
   IF (AllocateStatus /= 0) STOP "*** Not enough memory ***"

! generating all the standard young tabulae

write(22,*) "Standard Young Frame = 1"

write(22,*) a
write(22,*) b

write(22,*) " "

young(1,1:size(a)) = a(:)
young(1,size(a)+1:size(a)+size(b))=b(:)

m=2

!print*, young(1,:),"above"

a1 = a
b1 = b

! no of ways to pull down numbers from above row to below row

v1 = v
if (u.eq.v) then
 v1 =v-1
end if

pp=0
do i=1,v1

  k = Factorial (u-1)/(Factorial(i)*Factorial (u-1-i))
  l = Factorial (v)/(Factorial(i)*Factorial (v-i))

!  print*, k,l,k*l,"k and l"
! pp= pp+k*l

  alpha = 1
!  print*, u-1,i,"comb"
  allocate ( comb(i))
  Call gen (alpha,u-1,i,comb,1)
  deallocate (comb)

  CLOSE ( 23, STATUS='KEEP')
  allocate (rem(k,i))

  do j=1,k
  read(23,*) rem(j,:)
  end do

  do j=1,k
!  write(*,*) rem(j,:)
  end do

!----------------------------------------

  alpha = 1
  allocate ( comb(i))
  Call gen (alpha,v,i,comb,2)
  deallocate (comb)

  CLOSE ( 24, STATUS='KEEP')
  allocate (remb(l,i))

  do j=1,l
  read(24,*) remb(j,:)
  end do

  do j=1,l
!  write(*,*) remb(j,:)
  end do

!----------------------------------------

 do kk = 1,k
  do jj = 1,l
  do j=1,i

!   print*, rem(kk,j)+1, remb(jj,j),jj,"inside"
!   print*, a1(rem(kk,j)+1), b1(remb(jj,j)),"inside"
   Call Swap(a1(rem(kk,j)+1),b1(remb(jj,j)))
  end do
!  stop
  Call Sort (a1,u)
  Call Sort (b1,v)

!-------------------------------
  flag = 1
 do j=1,v
  if (a1(j).gt.b1(j)) then
   flag=flag*0
  end if
 end do

 if (flag.eq.1) then

 young(m,1:size(a1)) = a1(:)
 young(m,size(a1)+1:size(a1)+size(b1))=b1(:)
!print*,young(m,:)
 write(22,*) "Standard Young Frame = ",m

 write(22,*) a1
 write(22,*) b1

 write(22,*) " "
 m=m+1
 end if

 a1 = a
 b1 = b

!-----------------------------

 end do
 end do
  deallocate (rem)
  deallocate (remb)
  CLOSE ( 23, STATUS='DELETE' )
  CLOSE ( 24, STATUS='DELETE' )
!print*, young(1,:) ,"loop",i
end do

stop
print*, "Satndard Young Tabulae constructed, checking degenracy"
!print*, young(1,:)

!stop
flag2 = 0
do j=1,f_lam
 do k=j+1,f_lam

flag = 0

  do i=1,u+v
   if (young(j,i).eq.young(k,i)) then
   flag = flag + 1
   end if
  end do

  if (flag.eq.u+v) then
   flag2 =flag2+1
   print*, "degeneracy number = ", j,k
   print*, young(j,1:u)
   print*, young(j,u+1:u+v)
  end if

 end do
end do

print*,"No. of degeneracy = ", flag2
deallocate (young)

!-------------------------------------------------------------------

CONTAINS

! http://www.cs.mtu.edu/~shene/COURSES/cs201/NOTES/chap08/sorting.f90

! --------------------------------------------------------------------
! INTEGER*8 FUNCTION  FindMinimum():
!    This function returns the location of the minimum in the section
! between Start and End.
! --------------------------------------------------------------------

   INTEGER*8 FUNCTION  FindMinimum(x, Start, End)
      IMPLICIT  NONE
      INTEGER*8, DIMENSION(1:), INTENT(IN) :: x
      INTEGER*8, INTENT(IN)                :: Start, End
      INTEGER*8                            :: Minimum
      INTEGER*8                            :: Location
      INTEGER*8                            :: i

      Minimum  = x(Start)               ! assume the first is the min
      Location = Start                  ! record its position
      DO i = Start+1, End               ! start with next elements
         IF (x(i) < Minimum) THEN       !   if x(i) less than the min?
            Minimum  = x(i)             !      Yes, a new minimum found
            Location = i                !      record its position
         END IF
      END DO
      FindMinimum = Location            ! return the position
   END FUNCTION  FindMinimum

! --------------------------------------------------------------------
! SUBROUTINE  Swap():
!    This subroutine swaps the values of its two formal arguments.
! --------------------------------------------------------------------

   SUBROUTINE  Swap(a, b)
      IMPLICIT  NONE
      INTEGER*8, INTENT(INOUT) :: a, b
      INTEGER*8                :: Temp

      Temp = a
      a    = b
      b    = Temp
   END SUBROUTINE  Swap

! --------------------------------------------------------------------
! SUBROUTINE  Sort():
!    This subroutine receives an array x() and sorts it into ascending
! order.
! --------------------------------------------------------------------

   SUBROUTINE  Sort(x, Size)
      IMPLICIT  NONE
      INTEGER*8, DIMENSION(1:), INTENT(INOUT) :: x
      INTEGER*8, INTENT(IN)                   :: Size
      INTEGER*8                               :: i
      INTEGER*8                               :: Location

      DO i = 1, Size-1                      ! except for the last
         Location = FindMinimum(x, i, Size) ! find min from this to last
         CALL  Swap(x(i), x(Location)) ! swap this and the minimum
      END DO
   END SUBROUTINE  Sort

! Stanford - https://web.stanford.edu/class/me200c/tutorial_90/08_subprograms.html
RECURSIVE FUNCTION Factorial(n)  RESULT(Fact)

IMPLICIT NONE
INTEGER*8 :: Fact
INTEGER*8, INTENT(IN) :: n

IF (n == 0) THEN
   Fact = 1
ELSE
   Fact = n * Factorial(n-1)
END IF

END FUNCTION Factorial

end subroutine

recursive subroutine gen (m,n_max,m_max,comb,l)

    implicit none
    integer*8 m
    integer*8 n, n_max,m_max
    integer*8 comb(m_max)
    integer l

    if (m > m_max) then
     if (l.eq.1) then
      write (23,*) comb
      write(223,*) comb
     end if

     if (l.eq.2) then
      write(24,*) comb
      write(224,*) comb
     end if
    else
      do n = 1, n_max
        if ((m == 1) .or. (n > comb (m - 1))) then
          comb (m) = n
          call gen (m + 1,n_max,m_max,comb,l)
        end if
      end do
    end if

end subroutine gen


! { dg-do run }
! PR 20436: This used to give a runtime error.
program nested_reshape
  implicit none
  real :: k(8,2)
  real :: o(8,2)

  k = reshape((/1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0, &
            9.0,10.0,11.0,12.0,13.0,14.0,15.0,16.0/), (/8,2/))

  o = reshape(reshape(k, (/2,8/), order=(/2,1/)), (/8,2/))
end program


      subroutine fraction(irho)
c--
c  This calculates the fraction passing and the fraction trapped.
c--
C-----------------------------------------------
C   M o d u l e s
C-----------------------------------------------
      use parambs
      use vmec0
      implicit none
C-----------------------------------------------
C   D u m m y   A r g u m e n t s
C-----------------------------------------------
      integer irho
C-----------------------------------------------
C   L o c a l   P a r a m e t e r s
C-----------------------------------------------
      real(rprec), parameter :: one = 1, zero = 0
      integer, parameter :: n_lambda = 41
C-----------------------------------------------
C   L o c a l   V a r i a b l e s
C-----------------------------------------------
      integer :: l
      real(rprec), dimension(n_lambda) ::
     1   alambda, avgvpov, antgrand, answer

C-----------------------------------------------

      dlambda = one/(n_lambda - 1)

c  Fill lambda mesh, evaluate V||/V

      do l = 1, n_lambda
         alambda(l) = (l - 1)*dlambda
         avgvpov(l) = sum(sqrt(abs(one - alambda(l)*bfield))*gsqrt_b)
      end do

c  Form integrand

      where (avgvpov .gt. zero) antgrand = alambda/avgvpov

c  Do integral for the passing fraction

      call simpun (alambda, antgrand, n_lambda, answer)

      fpassing(irho) = .75_dp*avgbobm2*answer(n_lambda)*sum_gsqrt_b
      ftrapped(irho) = one - fpassing(irho)

      end subroutine fraction


!> \file test_mocsy.f90
!! \BRIEF Fortran 90 program to test mocsy.f90
PROGRAM test_derivauto

   USE msingledouble
   USE mderivauto
   USE mvars

   IMPLICIT NONE

!  For vars routine (called below)
!  "vars" Output variables:
   REAL(kind=rx), DIMENSION(1) :: ph, pco2, fco2, co2, hco3, co3, OmegaA, OmegaC, BetaD, rhoSW, p, tempis
!  "vars" Input variables
   INTEGER :: N
   REAL(kind=rx), DIMENSION(1) :: temp, sal, alk, dic, sil, phos, Patm, depth, lat
   REAL(kind=rx), DIMENSION(6,1) :: ph_deriv, pco2_deriv, fco2_deriv, co2_deriv, &
                     hco3_deriv, co3_deriv, omegaa_deriv, omegac_deriv
!  "vars" Input options
   CHARACTER(10) :: optCON, optT, optP, optB, optKf, optK1K2

!  Local variables:
   INTEGER :: i
   REAL(kind=r8) :: H
   REAL(kind=r8), DIMENSION(6) :: H_deriv
   CHARACTER*4 :: invar(6)

   invar(1) = 'Alk '
   invar(2) = 'DIC '
   invar(3) = 'Phos'
   invar(4) = 'Sil '
   invar(5) = 'T   '
   invar(6) = 'S   '

!> Typical options for observations
   optCON  = 'mol/kg'  ! input concentrations are in MOL/KG
   optT    = 'Tinsitu' ! input temperature, variable 'temp' is actually IN SITU temp [°C]
   optP    = 'm'      ! input variable 'depth' is in 'DECIBARS'
   optB    = 'l10'
   optK1K2 = 'l'
   optKf   = 'dg'
!> Simple input data (with CONCENTRATION units typical for DATA)
!> (based on observed average surface concentrations from S. Ocean (south of 60°S)--GLODAP and WOA2009)
   DO i = 1,1
     temp(i)   = 18.0            !Can be "Potential temperature" or "In situ temperature" (see optT below)
     sal(i)    = 35.0           !Salinity (practical scale)
     alk(i)    = 2300.*1.e-6      ! Convert obs. S. Ocean ave surf ALK (umol/kg) to mocsy data units (mol/kg)
     dic(i)    = 2000.*1.e-6      ! Convert obs. S. Ocean ave surf DIC (umol/kg) to mocsy data units (mol/kg)
     sil(i)    = 60.*1.e-6
     phos(i)   = 2.*1.e-6
     depth(i)  = 0.
     Patm(i)   = 1.0            !Atmospheric pressure (atm)
     lat(i)    = 0.
     N = i
   END DO

   call vars(ph, pco2, fco2, co2, hco3, co3, OmegaA, OmegaC, BetaD, rhoSW, p, tempis,         &  ! OUTPUT
             temp, sal, alk, dic, sil, phos, Patm, depth, lat, 1,                             &  ! INPUT
             optCON, optT, optP, optB=optB, optK1K2=optK1K2, optKf=optKf    )

   call derivauto(ph_deriv, pco2_deriv, fco2_deriv, co2_deriv, hco3_deriv, co3_deriv,   &
                omegaa_deriv, omegac_deriv,                                             &
                temp, sal, alk, dic, sil, phos, Patm, depth, lat, 1,                    &
                optCON, optT, optP, optB=optB, optK1K2=optK1K2, optKf=optKf )

    ! [H+] concentration
    H = 10.0**(-ph(1))
    ! derivative of [H+] deduced from that of pH
    H_deriv = - H * ph_deriv(:,1) * log(10.0)

    write (*,*) "derivatives :"
    write (*,*) "===========  "

    write (*,*) "      dH/dx         dpCO2/dx       dfCO2/dx", &
          "       d[CO2*]/dx      d[HCO3-]/dx    d[CO3--]/dx     dOmegaA/dx     dOmegaC/dx"

    DO i = 1,6
       write (*,"(A3, 8ES15.6)") invar(i), H_deriv(i), pco2_deriv(i,1), fco2_deriv(i,1), co2_deriv(i,1), hco3_deriv(i,1), &
       co3_deriv(i,1), omegaa_deriv(i,1), omegac_deriv(i,1)
    END DO

  STOP
END PROGRAM test_derivauto


! NCLFORTSTART
SUBROUTINE DCOMPUTEPI(pi, pressure, nx, ny, nz)
    USE wrf_constants, ONLY : P1000MB, RD, CP

    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: pi

    INTEGER, INTENT(IN) :: nx, ny, nz
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(OUT) :: pi
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: pressure
! NCLEND

    INTEGER i, j, k

    !REAL(KIND=8), PARAMETER :: P1000MB=100000.D0, R_D=287.D0, CP=7.D0*R_D/2.D0

    !$OMP PARALLEL DO COLLAPSE(3) SCHEDULE(runtime)
    DO k = 1,nz
        DO j = 1,ny
            DO i = 1,nx
                pi(i,j,k) = (pressure(i,j,k)/P1000MB)**(RD/CP)
            END DO
        END DO
    END DO
    !$OMP END PARALLEL DO

END SUBROUTINE DCOMPUTEPI

! Temperature from potential temperature in kelvin.
!NCLFORTSTART
SUBROUTINE DCOMPUTETK(tk, pressure, theta, nx)
    USE wrf_constants, ONLY : P1000MB, RD, CP

    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: tk

    INTEGER, INTENT(IN) :: nx
    !REAL(KIND=8) :: pi
    REAL(KIND=8), DIMENSION(nx), INTENT(IN) :: pressure
    REAL(KIND=8), DIMENSION(nx), INTENT(IN) :: theta
    REAL(KIND=8), DIMENSION(nx), INTENT(OUT) :: tk

! NCLEND

    INTEGER :: i

    !REAL(KIND=8), PARAMETER :: P1000MB=100000.D0, RD=287.D0, CP=7.D0*RD/2.D0

    !$OMP PARALLEL DO SCHEDULE(runtime)
    DO i = 1,nx
        !pi = (pressure(i)/P1000MB)**(RD/CP)
        !tk(i) = pi * theta(i)
        tk(i) = (pressure(i)/P1000MB)**(RD/CP) * theta(i)
    END DO
    !$OMP END PARALLEL DO

    RETURN

END SUBROUTINE DCOMPUTETK

! NCLFORTSTART
SUBROUTINE DINTERP3DZ(data3d, out2d, zdata, desiredloc, nx, ny, nz, missingval)
    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: out2d

    INTEGER, INTENT(IN) :: nx, ny, nz
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) ::  data3d
    REAL(KIND=8), DIMENSION(nx,ny), INTENT(OUT) :: out2d
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: zdata
    REAL(KIND=8), INTENT(IN) :: desiredloc
    REAL(KIND=8), INTENT(IN) :: missingval

! NCLEND

    INTEGER :: i,j,kp,ip,im
    LOGICAL :: dointerp
    REAL(KIND=8) :: w1,w2

    ! does vertical coordinate increase or decrease with increasing k?
    ! set offset appropriately

    ip = 0
    im = 1
    IF (zdata(1,1,1) .GT. zdata(1,1,nz)) THEN
        ip = 1
        im = 0
    END IF

    !$OMP PARALLEL DO COLLAPSE(2) PRIVATE(i,j,kp,dointerp,w1,w2) &
    !$OMP FIRSTPRIVATE(ip,im) SCHEDULE(runtime)
    DO i = 1,nx
        DO j = 1,ny
            ! Initialize to missing.  Was initially hard-coded to -999999.
            out2d(i,j) = missingval
            dointerp = .FALSE.
            kp = nz

            DO WHILE ((.NOT. dointerp) .AND. (kp >= 2))
                IF (((zdata(i,j,kp-im) < desiredloc) .AND. (zdata(i,j,kp-ip) > desiredloc))) THEN
                    w2 = (desiredloc - zdata(i,j,kp-im))/(zdata(i,j,kp-ip) - zdata(i,j,kp-im))
                    w1 = 1.D0 - w2
                    out2d(i,j) = w1*data3d(i,j,kp-im) + w2*data3d(i,j,kp-ip)
                    dointerp = .TRUE.
                END IF
                kp = kp - 1
            END DO
        END DO
    END DO
    !$OMP END PARALLEL DO

    RETURN

END SUBROUTINE DINTERP3DZ

! PORT DZSTAG HERE

! NCLFORTSTART
SUBROUTINE DZSTAG(znew, nx, ny, nz, z, nxz, nyz ,nzz, terrain)
    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: znew

    INTEGER, INTENT(IN) :: nx, ny, nz, nxz, nyz, nzz
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(OUT) :: znew
    REAL(KIND=8), DIMENSION(nxz,nyz,nzz), INTENT(IN) :: z
    REAL(KIND=8), DIMENSION(nxz,nyz), INTENT(IN) :: terrain
! NCLEND

    INTEGER :: i,j,k,ii,im1,jj,jm1

    ! check for u, v, or w (x,y,or z) staggering
    ! for x and y stag, avg z to x, y, point
    IF (nx .GT. nxz) THEN
        DO k = 1,nz
            DO j = 1,ny
                DO i = 1,nx
                    ii = MIN(i,nxz)
                    im1 = MAX(i-1,1)
                    znew(i,j,k) = 0.5D0*(z(ii,j,k) + z(im1,j,k))
                END DO
            END DO
        END DO

    ELSE IF (ny .GT. nyz) THEN
        DO k = 1,nz
            DO j = 1,NY
                jj = MIN(j,nyz)
                jm1 = MAX(j-1,1)
                DO i = 1,nx
                    znew(i,j,k) = 0.5D0*(z(i,jj,k) + z(i,jm1,k))
                END DO
            END DO
        END DO

    ! w (z) staggering
    ELSE IF (nz .GT. nzz) THEN
        DO j = 1,ny
            DO i = 1,nx
                znew(i,j,1) = terrain(i,j)
            END DO
        END DO

        DO k = 2,nz
            DO j = 1,ny
                DO i = 1,nx
                    znew(i,j,k) = znew(i,j,k-1) + 2.D0*(z(i,j,k-1) - znew(i,j,k-1))
                END DO
            END DO
        END DO

    END IF

    RETURN

    END SUBROUTINE DZSTAG

! NCLFORTSTART
SUBROUTINE DINTERP2DXY(v3d, v2d, xy, nx, ny, nz, nxy)

    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: v2d

    INTEGER, INTENT(IN) :: nx, ny, nz, nxy
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: v3d
    REAL(KIND=8), DIMENSION(nxy,nz), INTENT(OUT) :: v2d
    REAL(KIND=8), DIMENSION(2,nxy), INTENT(IN) :: xy

! NCLEND

    INTEGER :: i, j, k, ij
    REAL(KIND=8) :: w11, w12, w21, w22, wx, wy

    !$OMP PARALLEL DO PRIVATE(i,j,k,ij,w11,w12,w21,w22,wx,wy) &
    !$OMP SCHEDULE(runtime)
    DO ij = 1,nxy
        i = MAX(1,MIN(nx-1,INT(xy(1,ij)+1)))
        j = MAX(1,MIN(ny-1,INT(xy(2,ij)+1)))
        wx = DBLE(i+1) - (xy(1,ij)+1)
        wy = DBLE(j+1) - (xy(2,ij)+1)
        w11 = wx*wy
        w21 = (1.D0-wx)*wy
        w12 = wx*(1.D0-wy)
        w22 = (1.D0-wx)* (1.D0-wy)
        DO k = 1,nz
            v2d(ij,k) = w11*v3d(i,j,k) + w21*v3d(i+1,j,k) + &
            w12*v3d(i,j+1,k) + w22*v3d(i+1,j+1,k)
        END DO
    END DO
    !$OMP END PARALLEL DO

    RETURN

END SUBROUTINE DINTERP2DXY

! NCLFORTSTART
SUBROUTINE DINTERP1D(v_in, v_out, z_in, z_out, vmsg, nz_in, nz_out)

    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: v_out

    INTEGER, INTENT(IN) :: nz_in, nz_out
    REAL(KIND=8), DIMENSION(nz_in), INTENT(IN) :: v_in, z_in
    REAL(KIND=8), DIMENSION(nz_out), INTENT(IN) :: z_out
    REAL(KIND=8), DIMENSION(nz_out), INTENT(OUT) :: v_out
    REAL(KIND=8), INTENT(IN) :: vmsg

! NCLEND

    INTEGER :: kp,k,im,ip
    LOGICAL :: interp
    REAL(KIND=8) :: height,w1,w2

    ! does vertical coordinate increase of decrease with increasing k?
    ! set offset appropriately

    ip = 0
    im = 1
    IF (z_in(1) .GT. z_in(nz_in)) THEN
        ip = 1
        im = 0
    END IF

    !$OMP PARALLEL DO PRIVATE(kp, k, interp, height, w1, w2) &
    !$OMP FIRSTPRIVATE(ip, im) SCHEDULE(runtime)
    DO k = 1,nz_out
        v_out(k) = vmsg

        interp = .FALSE.
        kp = nz_in
        height = z_out(k)

        DO WHILE ((.NOT. interp) .AND. (kp .GE. 2))
            IF (((z_in(kp-im) .LE. height) .AND. (z_in(kp-ip) .GT. height))) THEN
                w2 = (height - z_in(kp-im))/(z_in(kp-ip) - z_in(kp-im))
                w1 = 1.D0 - w2
                v_out(k) = w1*v_in(kp-im) + w2*v_in(kp-ip)
                interp = .TRUE.
            END IF
            kp = kp - 1
        END DO
    END DO
    !$OMP END PARALLEL DO

    RETURN

END SUBROUTINE DINTERP1D

! This routine assumes
!    index order is (i,j,k)
!    wrf staggering
!
!    units: pressure (Pa), temperature(K), height (m), mixing ratio
!     (kg kg{-1}) availability of 3d p, t, and qv; 2d terrain; 1d
! half-level zeta string
!    output units of SLP are Pa, but you should divide that by 100 for the
!          weather weenies.
!    virtual effects are included
!

! NCLFORTSTART
SUBROUTINE DCOMPUTESEAPRS(nx, ny, nz, z, t, p, q, sea_level_pressure, &
                          t_sea_level, t_surf, level, errstat, errmsg)
    USE wrf_constants, ONLY : ALGERR, RD, G, USSALR

    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: sea_level_pressure

    !     Estimate sea level pressure.
    INTEGER, INTENT(IN) :: nx, ny, nz
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: z
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: t, p, q
    !     The output is the 2d sea level pressure.
    REAL(KIND=8), DIMENSION(nx,ny), INTENT(OUT) :: sea_level_pressure
    INTEGER, DIMENSION(nx,ny), INTENT(INOUT) ::  level
    REAL(KIND=8), DIMENSION(nx,ny), INTENT(INOUT) :: t_surf, t_sea_level
    INTEGER, INTENT(INOUT) :: errstat
    CHARACTER(LEN=*), INTENT(INOUT) :: errmsg

! NCLFORTEND

    !     Some required physical constants:

    !REAL(KIND=8), PARAMETER :: R=287.04D0, G=9.81D0, GAMMA=0.0065D0

    !     Specific constants for assumptions made in this routine:
    REAL(KIND=8), PARAMETER :: TC=273.16D0+17.5D0, PCONST=10000

    LOGICAL, PARAMETER :: ridiculous_mm5_test=.TRUE.
    !      PARAMETER (ridiculous_mm5_test = .FALSE.)

    !     Local variables:

    INTEGER :: i, j, k
    INTEGER :: klo, khi
    INTEGER :: errcnt, bad_i, bad_j

    REAL(KIND=8) :: bad_sfp
    REAL(KIND=8) :: plo, phi, tlo, thi, zlo, zhi
    REAL(KIND=8) :: p_at_pconst, t_at_pconst, z_at_pconst
    !REAL(KIND=8) :: z_half_lowest

    LOGICAL :: l1, l2, l3, found

    !  Find least zeta level that is PCONST Pa above the surface.  We
    !  later use this level to extrapolate a surface pressure and
    !  temperature, which is supposed to reduce the effect of the diurnal
    !  heating cycle in the pressure field.

    errstat = 0
    errcnt = 0
    bad_i = -1
    bad_j = -1
    bad_sfp = -1

    !$OMP PARALLEL DO COLLAPSE(2) PRIVATE(i,j,k,found) REDUCTION(+:errcnt) &
    !$OMP SCHEDULE(runtime)
    DO j = 1,ny
        DO i = 1,nx
            level(i,j) = -1

            k = 1
            found = .FALSE.
            DO WHILE ((.NOT. found) .AND. (k <= nz))
                IF (p(i,j,k) < p(i,j,1)-PCONST) THEN
                    level(i,j) = k
                    found = .TRUE.
                END IF
                k = k + 1
            END DO

            IF (level(i,j) == -1) THEN
                errcnt = errcnt + 1
                !$OMP CRITICAL
                ! Only do this the first time
                IF (bad_i .EQ. -1) THEN
                    bad_i = i
                    bad_j = j
                    bad_sfp = p(i,j,1) / 100.
                END IF
                !$OMP END CRITICAL
            END IF
        END DO
    END DO
    !$OMP END PARALLEL DO

    IF (errcnt > 0) THEN
        errstat = ALGERR
        WRITE(errmsg,*) "Error in finding 100 hPa up.  i=", bad_i, "j=", bad_j, "sfc_p=", bad_sfp
        RETURN
    END IF

    !     Get temperature PCONST Pa above surface.  Use this to extrapolate
    !     the temperature at the surface and down to sea level.
    !$OMP PARALLEL DO COLLAPSE(2) PRIVATE(i,j,klo,khi,plo, &
    !$OMP phi,tlo,thi,zlo,zhi,p_at_pconst,t_at_pconst,z_at_pconst) &
    !$OMP REDUCTION(+:errcnt) SCHEDULE(runtime)
    DO j = 1,ny
        DO i = 1,nx

            klo = MAX(level(i,j)-1, 1)
            khi = MIN(klo+1, nz-1)

            IF (klo == khi) THEN
                errcnt = errcnt + 1
                !$OMP CRITICAL
                IF (bad_i .EQ. -1) THEN
                    bad_i = i
                    bad_j = j
                END IF
                !$OMP END CRITICAL
            END IF

            plo = p(i,j,klo)
            phi = p(i,j,khi)
            tlo = t(i,j,klo)*(1.D0 + 0.608D0*q(i,j,klo))
            thi = t(i,j,khi)*(1.D0 + 0.608D0*q(i,j,khi))
            zlo = z(i,j,klo)
            zhi = z(i,j,khi)
            p_at_pconst = p(i,j,1) - PCONST
            t_at_pconst = thi - (thi-tlo)*LOG(p_at_pconst/phi)*LOG(plo/phi)
            z_at_pconst = zhi - (zhi-zlo)*LOG(p_at_pconst/phi)*LOG(plo/phi)
            !
            t_surf(i,j) = t_at_pconst * (p(i,j,1)/p_at_pconst)**(USSALR*RD/G)
            t_sea_level(i,j) = t_at_pconst + USSALR*z_at_pconst
        END DO
    END DO
    !$OMP END PARALLEL DO

    IF (errcnt > 0) THEN
        errstat = ALGERR
        WRITE(errmsg,*) "Error trapping levels at i=", bad_i, "j=", bad_j
        RETURN
    END IF

    ! If we follow a traditional computation, there is a correction to the
    ! sea level temperature if both the surface and sea level
    ! temperatures are *too* hot.

    IF (ridiculous_mm5_test) THEN
        !$OMP PARALLEL DO COLLAPSE(2) PRIVATE(l1,l2,l3) SCHEDULE(runtime)
        DO j = 1,ny
            DO i = 1,nx
                l1 = t_sea_level(i,j) < TC
                l2 = t_surf(i,j) <= TC
                l3 = .NOT. l1
                IF (l2 .AND. l3) THEN
                    t_sea_level(i,j) = TC
                ELSE
                    t_sea_level(i,j) = TC - 0.005D0*(t_surf(i,j)-TC)**2
                END IF
            END DO
        END DO
        !$OMP END PARALLEL DO
    END IF

    !     The grand finale: ta da!
    !$OMP PARALLEL DO COLLAPSE(2) SCHEDULE(runtime)
    DO j = 1,ny
        DO i = 1,nx
            !z_half_lowest = z(i,j,1)

    ! Convert to hPa in this step, by multiplying by 0.01. The original
    ! Fortran routine didn't do this, but the NCL script that called it
    ! did, so we moved it here.
            sea_level_pressure(i,j) = 0.01*(p(i,j,1)*EXP((2.D0*G*z(i,j,1))/&
                (RD*(t_sea_level(i,j) + t_surf(i,j)))))
        END DO
    END DO
    !$OMP END PARALLEL DO

    !     PRINT *,'sea pres input at weird location i=20,j=1,k=1'
    !     PRINT *,'t=',t(20,1,1),t(20,2,1),t(20,3,1)
    !     PRINT *,'z=',z(20,1,1),z(20,2,1),z(20,3,1)
    !     PRINT *,'p=',p(20,1,1),p(20,2,1),p(20,3,1)
    !     PRINT *,'slp=',sea_level_pressure(20,1),
    !    *         sea_level_pressure(20,2),sea_level_pressure(20,3)

    RETURN

END SUBROUTINE DCOMPUTESEAPRS

! Double precision version. If you make a change here, you
! must make the same change below to filter2d.

! NCLFORTSTART
SUBROUTINE DFILTER2D(a, b, nx, ny, it, missing, cenweight)

    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: a

    INTEGER, INTENT(IN) :: nx, ny, it
    REAL(KIND=8), DIMENSION(nx, ny), INTENT(INOUT) :: a
    REAL(KIND=8), INTENT(IN) :: missing, cenweight
    REAL(KIND=8), DIMENSION(nx, ny), INTENT(INOUT) :: b

! NCLEND

    INTEGER :: i, j, iter
    REAL(KIND=8) :: cenmult, coef

    cenmult = (cenweight) / 2.
    coef = 1.0 / (4. + cenweight)

    DO iter=1,it
        !$OMP PARALLEL DO COLLAPSE(2) SCHEDULE(runtime)
        DO j=1,ny
            DO i = 1,nx
                b(i,j) = a(i,j)
            END DO
        END DO
        !$OMP END PARALLEL DO

        !$OMP PARALLEL DO COLLAPSE(2) SCHEDULE(runtime)
        DO j=2,ny-1
            DO i=2,nx-1
                IF (b(i,j-1) .EQ. missing .OR. b(i,j) .EQ. missing .OR. &
                    b(i,j+1) .EQ. missing) THEN
                    a(i,j) = a(i,j)
                ELSE
                    a(i,j) = coef*(b(i,j-1) + cenmult*b(i,j) + b(i,j+1))
                END IF
            END DO
        END DO
        !$OMP END PARALLEL DO

        !$OMP PARALLEL DO COLLAPSE(2) SCHEDULE(runtime)
        DO j=2,ny-1
            DO i=2,nx-1
                IF (b(i-1,j) .EQ. missing .OR. b(i,j) .EQ. missing .OR. &
                    b(i+1,j) .EQ. missing) THEN
                    a(i,j) = a(i,j)
                ELSE
                    a(i,j) = a(i,j) + coef*(b(i-1,j) + cenmult*b(i,j) + b(i+1,j))
                END IF
            END DO
        END DO
        !$OMP END PARALLEL DO
    !        do j=1,ny
    !        do i=1,nx
    !          b(i,j) = a(i,j)
    !        enddo
    !        enddo
    !        do j=2,ny-1
    !        do i=1,nx
    !          a(i,j) = a(i,j) - .99*coef*(b(i,j-1)-2*b(i,j)+b(i,j+1))
    !        enddo
    !        enddo
    !        do j=1,ny
    !        do i=2,nx-1
    !          a(i,j) = a(i,j) - .99*coef*(b(i-1,j)-2*b(i,j)+b(i+1,j))
    !        enddo
    !        enddo
    END DO

    RETURN

END SUBROUTINE DFILTER2D

! Single precision version. If you make a change here, you
! must make the same change below to dfilter2d.

! NCLFORTSTART
SUBROUTINE FILTER2D(a, b, nx, ny, it, missing, cenweight)
    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: a

    INTEGER, INTENT(IN) :: nx, ny, it
    REAL(KIND=4), DIMENSION(nx, ny), INTENT(INOUT) :: a
    REAL(KIND=4), INTENT(IN) :: missing, cenweight
    REAL(KIND=4), DIMENSION(nx, ny), INTENT(INOUT) :: b

! NCLEND

    !REAL(KIND=8), PARAMETER :: COEF = .125
    INTEGER :: i, j, iter
    REAL(KIND=8) :: cenmult, coef

    cenmult = (cenweight) / 2.
    coef = 1.0 / (4. + cenweight)

    !$OMP PARALLEL

    DO iter=1,it
        !$OMP DO COLLAPSE(2) SCHEDULE(runtime)
        DO j=1,ny
            DO i = 1,nx
                b(i,j) = a(i,j)
            END DO
        END DO
        !$OMP END DO

        !$OMP DO COLLAPSE(2) SCHEDULE(runtime)
        DO j=2,ny-1
            DO i=2,nx-1
                IF (b(i,j-1) .EQ. missing .OR. b(i,j) .EQ. missing .OR. &
                    b(i,j+1) .EQ. missing) THEN
                    a(i,j) = a(i,j)
                ELSE
                    a(i,j) = coef*(b(i,j-1) + cenmult*b(i,j) + b(i,j+1))
                END IF
            END DO
        END DO
        !$OMP END DO

        !$OMP DO COLLAPSE(2) SCHEDULE(runtime)
        DO j=2,ny-1
            DO i=2,nx-1
                IF (b(i-1,j) .EQ. missing .OR. b(i,j) .EQ. missing .OR. &
                    b(i+1,j) .EQ. missing) THEN
                    a(i,j) = a(i,j)
                ELSE
                    a(i,j) = a(i,j) + coef*(b(i-1,j) + cenmult*b(i,j) + b(i+1,j))
                END IF
            END DO
        END DO
        !$OMP END DO
    !        do j=1,ny
    !        do i=1,nx
    !          b(i,j) = a(i,j)
    !        enddo
    !        enddo
    !        do j=2,ny-1
    !        do i=1,nx
    !          a(i,j) = a(i,j) - .99*coef*(b(i,j-1)-2*b(i,j)+b(i,j+1))
    !        enddo
    !        enddo
    !        do j=1,ny
    !        do i=2,nx-1
    !          a(i,j) = a(i,j) - .99*coef*(b(i-1,j)-2*b(i,j)+b(i+1,j))
    !        enddo
    !        enddo
    END DO

    !$OMP END PARALLEL

    RETURN

END SUBROUTINE FILTER2D

! NCLFORTSTART
SUBROUTINE DCOMPUTERH(qv, p, t, rh, nx)
    USE wrf_constants, ONLY : EZERO, ESLCON1, ESLCON2, CELKEL, RD, RV, EPS

    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: rh

    INTEGER, INTENT(IN) :: nx
    REAL(KIND=8), DIMENSION(nx), INTENT(IN) :: qv, p, t
    REAL(KIND=8), DIMENSION(nx), INTENT(OUT) :: rh

! NCLEND

    INTEGER :: i
    REAL(KIND=8) :: qvs,es,pressure,temperature

    !$OMP PARALLEL DO PRIVATE(qvs, es, pressure, temperature) SCHEDULE(runtime)
    DO i = 1,nx
        pressure = p(i)
        temperature = t(i)
        ! es  = 1000.*svp1*
        es = EZERO*EXP(ESLCON1*(temperature - CELKEL)/(temperature - ESLCON2))
        ! qvs = ep_2*es/(pressure-es)
        qvs = EPS*es/(0.01D0*pressure - (1.D0 - EPS)*es)
        ! rh = 100*MAX(1., qv(i)/qvs)
        ! rh(i) = 100.*qv(i)/qvs
        rh(i) = 100.D0*MAX(MIN(qv(i)/qvs, 1.0D0), 0.0D0)
    END DO
    !$OMP END PARALLEL DO

    RETURN

END SUBROUTINE DCOMPUTERH

! NCLFORTSTART
SUBROUTINE DGETIJLATLONG(lat_array, long_array, lat, longitude, ii, jj, nx, ny, imsg)
    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: ii,jj

    INTEGER, INTENT(IN) :: nx, ny, imsg
    INTEGER, INTENT(OUT) :: ii, jj
    REAL(KIND=8), DIMENSION(nx,ny), INTENT(IN) :: lat_array, long_array
    REAL(KIND=8) :: lat, longitude
! NCLEND

    REAL(KIND=8) :: longd,latd
    INTEGER :: i,j
    REAL(KIND=8) :: ir,jr
    REAL(KIND=8) :: dist_min,dist

    ! init to missing. was hard-coded to -999 initially.
    ir = imsg
    jr = imsg

    dist_min = 1.D+20
    DO j = 1,ny
        DO i = 1,nx
            latd = (lat_array(i,j) - lat)**2
            longd = (long_array(i,j) - longitude)**2
            ! longd = MIN((long_array(i,j)-longitude)**2, &
            !               (long_array(i,j)+longitude)**2)
            dist = SQRT(latd + longd)
            IF (dist_min .GT. dist) THEN
                dist_min = dist
                ir = DBLE(i)
                jr = DBLE(j)
            END IF
        END DO
    END DO

    ! The original version of this routine returned IR and JR. But, then
    ! the NCL script that called this routine was converting IR and JR
    ! to integer, so why not just return II and JJ?

    ! Also, I'm subtracing 1 here, because it will be returned to NCL
    ! script which has 0-based indexing.

    IF (ir .NE. imsg .AND. jr .NE. imsg) THEN
        ii = NINT(ir) - 1
        jj = NINT(jr) - 1
    ELSE
        ii = imsg
        jj = imsg
    END IF

    ! we will just return the nearest point at present

    RETURN

END SUBROUTINE DGETIJLATLONG

! You need to modify the C-WRAPPER in NCL for this to work

! NCLFORTSTART
SUBROUTINE DCOMPUTEUVMET(u, v, uvmet, longca,longcb,flong,flat, &
                        cen_long, cone, rpd, nx, ny, nxp1, nyp1, &
                        istag, is_msg_val, umsg, vmsg, uvmetmsg)
    IMPLICIT NONE

    ! ISTAG should be 0 if the U,V grids are not staggered.
    ! That is, NY = NYP1 and NX = NXP1.

    !f2py threadsafe
    !f2py intent(in,out) :: uvmet

    INTEGER,INTENT(IN) :: nx, ny, nxp1, nyp1, istag
    LOGICAL,INTENT(IN) :: is_msg_val
    REAL(KIND=8), DIMENSION(nxp1,ny), INTENT(IN):: u
    REAL(KIND=8), DIMENSION(nx,nyp1), INTENT(IN) :: v
    REAL(KIND=8), DIMENSION(nx,ny), INTENT(IN) :: flong
    REAL(KIND=8), DIMENSION(nx,ny), INTENT(IN) :: flat
    REAL(KIND=8), DIMENSION(nx,ny), INTENT(INOUT) :: longca
    REAL(KIND=8), DIMENSION(nx,ny), INTENT(INOUT) :: longcb
    REAL(KIND=8), INTENT(IN) :: cen_long, cone, rpd
    REAL(KIND=8), INTENT(IN) :: umsg, vmsg, uvmetmsg
    REAL(KIND=8), DIMENSION(nx,ny,2), INTENT(OUT) :: uvmet

! NCLEND

    INTEGER :: i,j
    !REAL(KIND=8) :: uk, vk

    ! msg stands for missing value in this code
    ! WRITE (6,FMT=*) ' in compute_uvmet ',NX,NY,NXP1,NYP1,ISTAG

    !$OMP PARALLEL

    !$OMP DO COLLAPSE(2) SCHEDULE(runtime)
    DO j = 1,ny
        DO i = 1,nx

            longca(i,j) = flong(i,j) - cen_long
            IF (longca(i,j).GT.180.D0) THEN
                longca(i,j) = longca(i,j) - 360.D0
            END IF
            IF (longca(i,j).LT.-180.D0) THEN
                longca(i,j) = longca(i,j) + 360.D0
            END IF
            IF (flat(i,j).LT.0.D0) THEN
                longcb(i,j) = -longca(i,j)*cone*rpd
            ELSE
                longcb(i,j) = longca(i,j)*cone*rpd
            END IF

            longca(i,j) = COS(longcb(i,j))
            longcb(i,j) = SIN(longcb(i,j))

        END DO
    END DO
    !$OMP END DO

    ! Note:  Intentionally removed as many IF statements as possible from loops
    ! to improve vectorization.

    IF (istag .EQ. 0) THEN ! Not staggered
        IF (.NOT. is_msg_val) THEN ! No missing values used
            !$OMP DO COLLAPSE(2) SCHEDULE(runtime)
            DO j = 1,ny
                DO i = 1,nx
                    uvmet(i,j,1) = v(i,j)*longcb(i,j) + u(i,j)*longca(i,j)
                    uvmet(i,j,2) = v(i,j)*longca(i,j) - u(i,j)*longcb(i,j)
                END DO
            END DO
            !$OMP END DO
        ELSE ! Missing values used
            !$OMP DO COLLAPSE(2) SCHEDULE(runtime)
            DO j = 1,ny
                DO i = 1,nx
                    IF ((u(i,j) .NE. umsg .AND. v(i,j) .NE. vmsg)) THEN
                        uvmet(i,j,1) = v(i,j)*longcb(i,j) + u(i,j)*longca(i,j)
                        uvmet(i,j,2) = v(i,j)*longca(i,j) - u(i,j)*longcb(i,j)
                    ELSE
                        uvmet(i,j,1) = uvmetmsg
                        uvmet(i,j,2) = uvmetmsg
                    END IF
                END DO
            END DO
            !$OMP END DO
        END IF
    ELSE ! Staggered
        IF (.NOT. is_msg_val) THEN ! No missing values used
            !$OMP DO COLLAPSE(2) SCHEDULE(runtime)
            DO j = 1,ny
                DO i = 1,nx
                    ! This is the more readable version.
                    !uk = 0.5D0*(u(i,j) + u(i+1,j))
                    !vk = 0.5D0*(v(i,j) + v(i,j+1))
                    !uvmet(i,j,1) = vk*longcb(i,j) + uk*longca(i,j)
                    !uvmet(i,j,2) = vk*longca(i,j) - uk*longcb(i,j)

                    uvmet(i,j,1) = (0.5D0*(v(i,j) + v(i,j+1)))*longcb(i,j) + &
                                   (0.5D0*(u(i,j) + u(i+1,j)))*longca(i,j)
                    uvmet(i,j,2) = (0.5D0*(v(i,j) + v(i,j+1)))*longca(i,j) - &
                                   (0.5D0*(u(i,j) + u(i+1,j)))*longcb(i,j)

                END DO
            END DO
            !$OMP END DO
        ELSE ! Missing values used
            !$OMP DO COLLAPSE(2) SCHEDULE(runtime)
            DO j = 1,ny
                DO i = 1,nx
                    IF (u(i,j) .NE. umsg .AND. v(i,j) .NE. vmsg .AND. u(i+1,j) .NE. umsg .AND. v(i,j+1) .NE. vmsg) THEN
                        ! This is the more readable version.
                        !uk = 0.5D0*(u(i,j) + u(i+1,j))
                        !vk = 0.5D0*(v(i,j) + v(i,j+1))
                        !uvmet(i,j,1) = vk*longcb(i,j) + uk*longca(i,j)
                        !uvmet(i,j,2) = vk*longca(i,j) - uk*longcb(i,j)

                        uvmet(i,j,1) = (0.5D0*(v(i,j) + v(i,j+1)))*longcb(i,j) + &
                                       (0.5D0*(u(i,j) + u(i+1,j)))*longca(i,j)
                        uvmet(i,j,2) = (0.5D0*(v(i,j) + v(i,j+1)))*longca(i,j) - &
                                       (0.5D0*(u(i,j) + u(i+1,j)))*longcb(i,j)
                    ELSE
                        uvmet(i,j,1) = uvmetmsg
                        uvmet(i,j,2) = uvmetmsg
                    END IF
                END DO
            END DO
            !$OMP END DO
        END IF
    END IF

    !$OMP END PARALLEL

    RETURN

END SUBROUTINE DCOMPUTEUVMET

! This was originally a routine that took 2D input arrays. Since
! the NCL C wrapper routine can handle multiple dimensions, it's
! not necessary to have anything bigger than 1D here.

! NCLFORTSTART
SUBROUTINE DCOMPUTETD(td, pressure, qv_in, nx)
    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: td

    INTEGER, INTENT(IN) :: nx
    REAL(KIND=8), DIMENSION(nx), INTENT(IN) :: pressure
    REAL(KIND=8), DIMENSION(nx), INTENT(IN) :: qv_in
    REAL(KIND=8), DIMENSION(nx), INTENT(OUT) :: td

! NCLEND

    REAL(KIND=8) :: qv,tdc

    INTEGER :: i

    !$OMP PARALLEL DO PRIVATE(i,qv,tdc) SCHEDULE(runtime)
    DO i = 1,nx
        qv = MAX(qv_in(i), 0.D0)
         ! vapor pressure
         tdc = qv*pressure(i)/(.622D0 + qv)

         ! avoid problems near zero
         tdc = MAX(tdc, 0.001D0)
         td(i) = (243.5D0*LOG(tdc) - 440.8D0)/(19.48D0 - LOG(tdc))
    END DO
    !$OMP END PARALLEL DO

    RETURN

END SUBROUTINE DCOMPUTETD

! NCLFORTSTART
SUBROUTINE DCOMPUTEICLW(iclw, pressure, qc_in, nx, ny, nz)
    USE wrf_constants, ONLY : G

    IMPLICIT NONE

    !f2py threadsafe
    !f2py intent(in,out) :: iclw

    INTEGER, INTENT(IN) :: nx,ny,nz
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: pressure
    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: qc_in
    REAL(KIND=8), DIMENSION(nx,ny), INTENT(OUT) :: iclw

! NCLEND

    REAL(KIND=8) :: qclw, dp
    REAL(KIND=8), PARAMETER :: GG = 1000.D0/G
    INTEGER i,j,k

    iclw = 0

    DO j = 3,ny - 2
        DO i = 3,nx - 2
            DO k = 1,nz
                qclw = MAX(qc_in(i,j,k), 0.D0)
                IF (k.EQ.1) THEN
                    dp = pressure(i,j,k-1) - pressure(i,j,k)
                ELSE IF (k.EQ.nz) then
                    dp = pressure(i,j,k) - pressure(i,j,k+1)
                ELSE
                    dp = (pressure(i,j,k-1) - pressure(i,j,k+1))/2.D0
                END IF
                iclw(i,j) = iclw(i,j) + qclw*dp*GG
            END DO
        END DO
    END DO

    RETURN

END SUBROUTINE DCOMPUTEICLW


module ofdft_fft

! Routines for orbital free density functional theory, discretization is
! done using FFT.
!
! We use 3D reciprocal grid (:, :, :). Variables in reciprocal (Fourier) space
! have G appended to them. E.g. Ven is the electron-nucleus potential in real
! space, VenG is the same potential in reciprocal space.

use types, only: dp
use constants, only: pi, i_, Ha2eV
use optimize, only: bracket, brent, parabola_vertex
use ofdft, only: f
use xc, only: xc_pz, xc_pz2
use utils, only: stop_error, assert, clock
use fourier, only: fft3_inplace, ifft3_inplace, &
    fft1_inplace => fft_pass_inplace, ifft1_inplace
use integration, only: integrate_trapz_1
implicit none
private
public reciprocal_space_vectors, free_energy, free_energy_min, &
    radial_potential_fourier, real2fourier, fourier2real, integralG, &
    logging_info, integral, real_space_vectors, radial_density_fourier, &
    vtk_save, update_fletcher_reeves, update_polak_ribiere, &
    poisson_kernel

! Update types for nonlinear conjugate gradient method:
integer, parameter :: update_fletcher_reeves = 1
integer, parameter :: update_polak_ribiere   = 2

logical :: logging_info = .true.

interface real_space_vectors
    module procedure real_space_vectors_1d
    module procedure real_space_vectors_3d
end interface

interface reciprocal_space_vectors
    module procedure reciprocal_space_vectors_1d
    module procedure reciprocal_space_vectors_3d
end interface

interface integral
    module procedure integral3d, integral1d
end interface

interface integralG
    module procedure integralG_complex, integralG_real
    module procedure integralG_real_1d
end interface

interface real2fourier
    module procedure real2fourier_complex, real2fourier_real
    module procedure real2fourier1d_complex, real2fourier1d_real
end interface

interface fourier2real
    module procedure fourier2real_complex, fourier2real_real
    module procedure fourier2real1d_complex, fourier2real1d_real
end interface

integer :: fft_counter
real(dp) :: fft_time = 0

contains

subroutine real_space_vectors_3d(L, X)
! Calculates the real space vectors in the box [0, L]^3
real(dp), intent(in) :: L(:)
! X(i, j, k, :) is the 3D position vector of the point with index (i, j, k)
real(dp), intent(out) :: X(:, :, :, :)
integer :: Ng(3), i, j, k
Ng = [size(X,1), size(X,2), size(X,3)]
forall(i=1:Ng(1), j=1:Ng(2), k=1:Ng(3))
    X(i, j, k, :) = ([i, j, k] - 1) * L / Ng
end forall
end subroutine

subroutine real_space_vectors_1d(L, X)
! Calculates the real space vectors in the box [0, L]
real(dp), intent(in) :: L
! X(i) is the 1D position vector of the point with index (i)
real(dp), intent(out) :: X(:)
integer :: Ng, i
Ng = size(X, 1)
forall(i=1:Ng) X(i) = (i-1) * L / Ng
end subroutine

subroutine reciprocal_space_vectors_3d(L, G, G2)
real(dp), intent(in) :: L(:)
! G(:, :, :, i) where i=1, 2, 3 are the x, y, z components
! G2(:, :, :) are the squares of G
real(dp), intent(out) :: G(:, :, :, :), G2(:, :, :)
integer :: Ng(3), i, j, k
Ng = shape(G2)
! The FFT frequencies are ordered as follows
!
!     G1D(1:Ng) = 2*pi*k/L,
!
! where
!
!     k = 0, 1, ..., Ng/2-1, -Ng/2, -Ng/2+1, ..., -1 if Ng is even,
!     k = 0, 1, ..., (Ng-1)/2, -(Ng-1)/2, -(Ng-1)/2+1, ..., -1 if Ng is odd.
!
! E.g. for
!
!     Ng=8 we get k = 0, 1, 2, 3, -4, -3, -2, -1.
!     Ng=9 we get k = 0, 1, 2, 3, 4, -4, -3, -2, -1.
!
! 'k' can be generated using
!
!     forall(i=1:Ng) k(i) = i-1 - Ng*nint((i-0.5_dp)/Ng)    ! if Ng is even
!     forall(i=1:Ng) k(i) = i-1 - Ng*nint((i-1.5_dp)/Ng)    ! if Ng is odd
!
! Note that the Ng/2 frequency can go both ways: +/- Ng/2. Below we take it with
! the "+" sign:
!
!     k = 0, 1, ..., Ng/2-1, +Ng/2, -Ng/2+1, ..., -1 if n is even
!
! E.g. for
!
!     Ng=8 we get k = 0, 1, 2, 3, 4, -3, -2, -1.
!     Ng=9 we get k = 0, 1, 2, 3, 4, -4, -3, -2, -1.
!
! These can be generated using (this holds for 'Ng' both even and odd):
!
!     forall(i=1:Ng) k(i) = i-1 - Ng*nint((i-1.5_dp)/Ng)
!
! That is the formula that we use below:
forall(i=1:Ng(1), j=1:Ng(2), k=1:Ng(3))
    G(i, j, k, :) = 2*pi/L * ([i,j,k] - 1 - Ng*nint(([i,j,k]-1.5_dp)/Ng))
end forall
G2 = G(:,:,:,1)**2 + G(:,:,:,2)**2 + G(:,:,:,3)**2
end subroutine

subroutine reciprocal_space_vectors_1d(L, G, G2)
real(dp), intent(in) :: L
! G(:) are the x components
! G2(:) are the squares of G
real(dp), intent(out) :: G(:), G2(:)
integer :: Ng, i
Ng = size(G, 1)
forall(i=1:Ng) G(i) = 2*pi/L * (i-1-Ng*nint((i-1.5_dp)/Ng))
G2 = G**2
end subroutine

subroutine real2fourier_complex(x, xG)
! Calculates Discrete Fourier Transform, with the same normalization as the
! Fourier Transform for periodic systems (which is Fourier Series).
complex(dp), intent(in) :: x(:, :, :)
complex(dp), intent(out) :: xG(:, :, :)
integer :: Ng
real(dp) :: t1, t2
t1 = clock()
Ng = size(x, 1)
xG = x
call fft3_inplace(xG) ! Calculates sum_{n=0}^{N-1} e^{-2*pi*i*k*n/N}*x(n)
xG = xG / size(x)     ! The proper normalization is to divide by N
t2 = clock()
fft_counter = fft_counter + 1
fft_time = fft_time + t2-t1
end subroutine

subroutine real2fourier1d_complex(x, xG)
! Calculates Discrete Fourier Transform, with the same normalization as the
! Fourier Transform for periodic systems (which is Fourier Series).
complex(dp), intent(in) :: x(:)
complex(dp), intent(out) :: xG(:)
integer :: Ng
real(dp) :: t1, t2
t1 = clock()
Ng = size(x, 1)
xG = x
call fft1_inplace(xG) ! Calculates sum_{n=0}^{N-1} e^{-2*pi*i*k*n/N}*x(n)
xG = xG / size(x)     ! The proper normalization is to divide by N
t2 = clock()
fft_counter = fft_counter + 1
fft_time = fft_time + t2-t1
end subroutine

subroutine real2fourier_real(x, xG)
! Calculates Discrete Fourier Transform, with the same normalization as the
! Fourier Transform for periodic systems (which is Fourier Series).
real(dp), intent(in) :: x(:, :, :)
complex(dp), intent(out) :: xG(:, :, :)
integer :: Ng
real(dp) :: t1, t2
t1 = clock()
Ng = size(x, 1)
xG = x
call fft3_inplace(xG) ! Calculates sum_{n=0}^{N-1} e^{-2*pi*i*k*n/N}*x(n)
xG = xG / size(x)     ! The proper normalization is to divide by N
t2 = clock()
fft_counter = fft_counter + 1
fft_time = fft_time + t2-t1
end subroutine

subroutine real2fourier1d_real(x, xG)
! Calculates Discrete Fourier Transform, with the same normalization as the
! Fourier Transform for periodic systems (which is Fourier Series).
real(dp), intent(in) :: x(:)
complex(dp), intent(out) :: xG(:)
integer :: Ng
real(dp) :: t1, t2
t1 = clock()
Ng = size(x, 1)
xG = x
call fft1_inplace(xG) ! Calculates sum_{n=0}^{N-1} e^{-2*pi*i*k*n/N}*x(n)
xG = xG / size(x)     ! The proper normalization is to divide by N
t2 = clock()
fft_counter = fft_counter + 1
fft_time = fft_time + t2-t1
end subroutine

subroutine fourier2real_complex(xG, x)
! Calculates Inverse Discrete Fourier Transform. xG must follow the same
! normalization as defined by real2fourier(), i.e. in the following calls, 'x2'
! will be equal to 'x':
! call real2fourier(x, xG)
! call fourier2real(xG, x2)
complex(dp), intent(in) :: xG(:, :, :)
complex(dp), intent(out) :: x(:, :, :)
real(dp) :: t1, t2
t1 = clock()
x = xG
call ifft3_inplace(x) ! Calculates sum_{k=0}^{N-1} e^{2*pi*i*k*n/N}*X(k)
! The result is already normalized
t2 = clock()
fft_counter = fft_counter + 1
fft_time = fft_time + t2-t1
end subroutine

subroutine fourier2real1d_complex(xG, x)
! Calculates Inverse Discrete Fourier Transform. xG must follow the same
! normalization as defined by real2fourier(), i.e. in the following calls, 'x2'
! will be equal to 'x':
! call real2fourier(x, xG)
! call fourier2real(xG, x2)
complex(dp), intent(in) :: xG(:)
complex(dp), intent(out) :: x(:)
real(dp) :: t1, t2
t1 = clock()
x = xG
call ifft1_inplace(x) ! Calculates sum_{k=0}^{N-1} e^{2*pi*i*k*n/N}*X(k)
! The result is already normalized
t2 = clock()
fft_counter = fft_counter + 1
fft_time = fft_time + t2-t1
end subroutine

subroutine fourier2real_real(xG, x)
! Calculates Inverse Discrete Fourier Transform. xG must follow the same
! normalization as defined by real2fourier(), i.e. in the following calls, 'x2'
! will be equal to 'x':
! call real2fourier(x, xG)
! call fourier2real(xG, x2)
complex(dp), intent(in) :: xG(:, :, :)
real(dp), intent(out) :: x(:, :, :)
complex(dp) :: tmp(size(xG,1), size(xG,2), size(xG,3))
real(dp) :: t1, t2
t1 = clock()
tmp = xG
call ifft3_inplace(tmp) ! Calculates sum_{k=0}^{N-1} e^{2*pi*i*k*n/N}*X(k)
x = real(tmp, dp)       ! The result is already normalized
if (maxval(abs(aimag(tmp))) > 1e-12_dp) then
    if (logging_info) then
        print *, "INFO: fourier2real() max imaginary part:", maxval(aimag(tmp))
    end if
end if
! TODO: make this strict:
!call assert(maxval(abs(aimag(tmp))) < 1e-1_dp)
t2 = clock()
fft_counter = fft_counter + 1
fft_time = fft_time + t2-t1
end subroutine

subroutine fourier2real1d_real(xG, x)
! Calculates Inverse Discrete Fourier Transform. xG must follow the same
! normalization as defined by real2fourier(), i.e. in the following calls, 'x2'
! will be equal to 'x':
! call real2fourier(x, xG)
! call fourier2real(xG, x2)
complex(dp), intent(in) :: xG(:)
real(dp), intent(out) :: x(:)
complex(dp) :: tmp(size(xG,1))
real(dp) :: t1, t2
t1 = clock()
tmp = xG
call ifft1_inplace(tmp) ! Calculates sum_{k=0}^{N-1} e^{2*pi*i*k*n/N}*X(k)
x = real(tmp, dp)       ! The result is already normalized
if (maxval(abs(aimag(tmp))) > 1e-12_dp) then
    if (logging_info) then
        print *, "INFO: fourier2real() max imaginary part:", maxval(aimag(tmp))
    end if
end if
! TODO: make this strict:
!call assert(maxval(abs(aimag(tmp))) < 1e-1_dp)
t2 = clock()
fft_counter = fft_counter + 1
fft_time = fft_time + t2-t1
end subroutine

real(dp) function integral3d(L, f) result(r)
real(dp), intent(in) :: L, f(:, :, :)
r = sum(f) * L**3 / size(f)
end function

real(dp) function integral1d(L, f) result(r)
real(dp), intent(in) :: L, f(:)
r = sum(f) * L / size(f)
end function

real(dp) function integralG_complex(L, fG) result(r)
real(dp), intent(in) :: L
complex(dp), intent(in) :: fG(:, :, :)
complex(dp) :: s
s = sum(fG) * L**3
r = real(s, dp)
if (abs(aimag(s)) > 1e-12_dp) then
    if (logging_info) then
        print *, "INFO: integralG() imaginary part:", aimag(s)
    end if
end if
!if (abs(aimag(s)) > 1e-5_dp) then
!    print *, "aimag(s) =", aimag(s)
!    call stop_error("integralG(): Complex part is not negligible.")
!end if
end function

real(dp) function integralG_real(fG, L) result(r)
real(dp), intent(in) :: L, fG(:, :, :)
r = sum(fG) * L**3
end function

real(dp) function integralG_real_1d(fG, L) result(r)
real(dp), intent(in) :: L, fG(:)
r = sum(fG) * L
end function

subroutine calc_dFdtheta(L, T_au, Ven, C1, C2, C3, psi, eta, theta, dFdtheta)
! Calculates dF[psi(theta)]/dtheta efficiently (no FFT needed, just integrals)
! See the function find_theta() for an example how to prepare the C1, C2 and C3
! constants and how to use this subroutine.
real(dp), intent(in) :: L, T_au, psi(:, :, :), eta(:, :, :)
real(dp), intent(in) :: C1(:, :, :), C2(:, :, :), C3(:, :, :)
real(dp), intent(in) :: theta
real(dp), intent(in) :: Ven(:, :, :)
real(dp), intent(out) :: dFdtheta
real(dp), dimension(size(psi, 1), size(psi, 2), size(psi, 3)) :: psi_, ne

real(dp), dimension(size(Ven,1), size(Ven,2), size(Ven,3)) :: y, &
    Vee, Vxc, dF0dn
real(dp) :: beta, dydn
psi_ = psi*cos(theta)+eta*sin(theta)
ne = psi_**2

call xc_pz2(ne, Vxc)
beta = 1/T_au
y = pi**2 / sqrt(2._dp) * beta**(3._dp/2) * ne
if (any(y < 0)) call stop_error("Density must be positive")
dydn = pi**2 / sqrt(2._dp) * beta**(3._dp/2)
dF0dn = 1 / beta * f(y) + ne / beta * f(y, deriv=.true.) * dydn
Vee = C1*cos(theta)**2 + C2*sin(2*theta) + C3*sin(theta)**2

dFdtheta = integral(L, (-psi*sin(theta)+eta*cos(theta)) * &
    (dF0dn + Vee + Ven + Vxc) * L**3 * psi_)
end subroutine

subroutine calc_F(L, T_au, Ven, C1, C2, C3, psi, eta, theta, F_)
! Calculates F[psi(theta)] efficiently (no FFT needed, just integrals)
! See the function find_theta() for an example how to prepare the C1, C2 and C3
! constants and how to use this subroutine.
real(dp), intent(in) :: L, T_au, psi(:, :, :), eta(:, :, :)
real(dp), intent(in) :: C1(:, :, :), C2(:, :, :), C3(:, :, :)
real(dp), intent(in) :: theta
real(dp), intent(in) :: Ven(:, :, :)
real(dp), intent(out) :: F_
real(dp), dimension(size(psi, 1), size(psi, 2), size(psi, 3)) :: psi_, ne, y, &
    Vee, exc_density, Vxc
real(dp) :: beta
psi_ = psi*cos(theta)+eta*sin(theta)
ne = psi_**2

call xc_pz(ne, exc_density, Vxc)
beta = 1/T_au
y = pi**2 / sqrt(2._dp) * beta**(3._dp/2) * ne
if (any(y < 0)) call stop_error("Density must be positive")
Vee = C1*cos(theta)**2 + C2*sin(2*theta) + C3*sin(theta)**2

F_ = integral(L, ne / beta * f(y) + (Ven+Vee/2+exc_density)*ne)
end subroutine

subroutine poisson_kernel(n, neG, G2, VeeG)
! Calculates VeeG = 4*pi*neG / G2, but skips the (1,1,1) term
integer, intent(in) :: n
complex(dp), intent(in) :: neG(n)
real(dp), intent(in) :: G2(n)
complex(dp), intent(out) :: VeeG(n)
VeeG(1) = 0
VeeG(2:) = 4*pi*neG(2:) / G2(2:)
end subroutine

subroutine precalc_C1C2C2(G2, psi, eta, C1, C2, C3)
real(dp), dimension(:, :, :), intent(in) :: G2, psi, eta
real(dp), dimension(:, :, :), intent(out) :: C1, C2, C3
complex(dp), dimension(size(psi,1), size(psi,2), size(psi,3)) :: neG, VeeG
call real2fourier(psi**2, neG)
call poisson_kernel(size(neG), neG, G2, VeeG)
call fourier2real(VeeG, C1)

call real2fourier(psi*eta, neG)
call poisson_kernel(size(neG), neG, G2, VeeG)
call fourier2real(VeeG, C2)

call real2fourier(eta**2, neG)
call poisson_kernel(size(neG), neG, G2, VeeG)
call fourier2real(VeeG, C3)
end subroutine

real(dp) function free_energy_theta(L, G2, T_au, VenG, psi, eta, theta) &
        result(F_)
! Shows how to use calc_F. Typically you want to precalculate C1, C2, C3 and
! store them.
real(dp), intent(in) :: L, G2(:, :, :), T_au, psi(:, :, :), eta(:, :, :)
complex(dp), intent(in) :: VenG(:, :, :)
real(dp), intent(in) :: theta
real(dp), dimension(size(VenG,1), size(VenG,2), size(VenG,3)) :: Ven, C1, C2, C3
call fourier2real(VenG, Ven)
call precalc_C1C2C2(G2, psi, eta, C1, C2, C3)
call calc_F(L, T_au, Ven, C1, C2, C3, psi, eta, theta, F_)
end function

subroutine find_theta(L, G2, T_au, VenG, psi, eta)
! It prints the values of dF/dtheta using efficient evaluation
real(dp), intent(in) :: L, G2(:, :, :), T_au, psi(:, :, :), eta(:, :, :)
complex(dp), intent(in) :: VenG(:, :, :)
real(dp) :: theta, dFdtheta
real(dp), dimension(size(VenG,1), size(VenG,2), size(VenG,3)) :: Ven, C1, C2, C3
integer :: i
real(dp) :: t1, t2
call fourier2real(VenG, Ven)
call precalc_C1C2C2(G2, psi, eta, C1, C2, C3)
print *, "dFdtheta:"
theta = 0
call cpu_time(t1)
do i = 1, 100
    call calc_dFdtheta(L, T_au, Ven, C1, C2, C3, psi, eta, theta, dFdtheta)
    print *, theta, dFdtheta
    theta = theta + pi/2/100
end do
call cpu_time(t2)
print *, "time: ", t2-t1
end subroutine

subroutine free_energy(L, G2, T_au, VenG, ne, Eee, Een, Ts, Exc, Etot, dFdn, &
    calc_value, calc_derivative)
real(dp), intent(in) :: L, G2(:, :, :), T_au, ne(:, :, :)
complex(dp), intent(in) :: VenG(:, :, :)
real(dp), intent(out) :: Eee, Een, Ts, Exc, Etot
logical, intent(in) :: calc_value, calc_derivative

! dFdn returns "delta F / delta n", the functional derivative with respect to
! the density "n". Use the relation
!     d/dpsi = 2 psi d/dn
! to obtain the derivative with respect to psi (i.e. multiply Hpsi by 2*psi).
real(dp), intent(out) :: dFdn(:, :, :)

real(dp), dimension(size(VenG,1), size(VenG,2), size(VenG,3)) :: y, F0, &
    exc_density, Ven_ee, Vxc, dF0dn
complex(dp), dimension(size(VenG,1), size(VenG,2), size(VenG,3)) :: &
    neG, VeeG
real(dp) :: beta, dydn
call assert(calc_value .or. calc_derivative)

call real2fourier(ne, neG)

call poisson_kernel(size(neG), neG, G2, VeeG)

beta = 1/T_au
! The density must be positive, the f(y) fails for negative "y". Thus we use
! ne.
y = pi**2 / sqrt(2._dp) * beta**(3._dp/2) * ne
if (any(y < 0)) call stop_error("Density must be positive")

call xc_pz(ne, exc_density, Vxc)

if (calc_value) then
    ! Hartree energy
    !Eee = integralG2(L, real(VeeG)*real(neG)+aimag(VeeG)*aimag(neG)) / 2
    Eee = integralG(L, VeeG*conjg(neG)) / 2
    ! Electron-nucleus energy
    !Een = integralG2(L, real(VenG)*real(neG)+aimag(VenG)*aimag(neG))
    Een = integralG(L, VenG*conjg(neG))

    ! Kinetic energy using Perrot parametrization
    F0 = ne / beta * f(y)
    Ts = integral(L, F0)
    ! Exchange and correlation potential
    Exc = integral(L, exc_density * ne)
    Etot = Ts + Een + Eee + Exc
end if

if (calc_derivative) then
    ! Calculate the derivative
    dydn = pi**2 / sqrt(2._dp) * beta**(3._dp/2)
    ! F0 = ne / beta * f(y)
    ! d F0 / d n =
    dF0dn = 1 / beta * f(y) + ne / beta * f(y, deriv=.true.) * dydn

    call fourier2real(VenG+VeeG, Ven_ee)

    dFdn = dF0dn + Ven_ee + Vxc
end if
end subroutine

subroutine radial_potential_fourier(R, V, L, Z, VenG, V0)
! Takes radial potential (given at origin) defined by:
!   V(R) on a grid R
!   Z/r for r > maxval(R)
! and calculates a 3D Fourier transform of it into the VenG grid. 'L' is the
! length of the box.
real(dp), intent(in) :: R(:), V(:), L, Z
real(dp), intent(out) :: VenG(:, :, :)
real(dp), intent(out) :: V0
integer :: Ng, i, j, k, idx
real(dp) :: Rp(size(R)), dk, Rc, w, Vk(0:3*(size(VenG, 1)/2+1)**2)
! Rp is the derivative of the mesh R'(t), which for uniform mesh is equal to
! the mesh step (rmax-rmin)/N:
Rp = (R(size(R)) - R(1)) / (size(R)-1)
Rc = R(size(R))
Ng = size(VenG, 1)
call assert(size(VenG, 2) == Ng)
call assert(size(VenG, 3) == Ng)
dk = 2*pi/L

! V0 =      \int V(r) - Z/r d^3x
!    = 4*pi*\int_0^Rc (V(r) - Z/r) r^2 dr
!    = 4*pi*\int_0^Rc V(r) r^2 dr - Z*Rc^2 / 2
! So these two are equivalent, we use the one that is faster:
!V0 = -4*pi*integrate(Rp, R**2*(V-Z/R))
V0 = -4*pi*(integrate(Rp, R**2*V) - Z*Rc**2/2)

! We prepare the values of the radial Fourier transform on a 3D grid
Vk(0) = 0
do i = 1, 3*(Ng/2+1)**2
    w = sqrt(real(i, dp))*dk
    Vk(i) = 4*pi/(L**3 * w**2) * (w*integrate(Rp, R*sin(w*R)*V) + Z*cos(w*Rc))
end do

! We fill out the 3D grid using the values from Vk
do i = 1, Ng
do j = 1, Ng
do k = 1, Ng
    idx = (i-1-Ng*nint((i-1.5_dp)/Ng))**2 + (j-1-Ng*nint((j-1.5_dp)/Ng))**2 &
            + (k-1-Ng*nint((k-1.5_dp)/Ng))**2
    VenG(i, j, k) = Vk(idx)
end do
end do
end do
end subroutine

subroutine radial_density_fourier(R, ne, L, neG)
! Takes radial density (given at origin) defined by:
!   ne(R) on a grid R
!   0 for r > maxval(R)
! and calculates a 3D Fourier transform of it into the neG grid. 'L' is the
! length of the box.
real(dp), intent(in) :: R(:), ne(:), L
real(dp), intent(out) :: neG(:, :, :)
integer :: Ng, i, j, k, idx
real(dp) :: Rp(size(R)), dk, Rc, w, nk(0:3*(size(neG, 1)/2+1)**2)
! Rp is the derivative of the mesh R'(t), which for uniform mesh is equal to
! the mesh step (rmax-rmin)/N:
Rp = (R(size(R)) - R(1)) / (size(R)-1)
Rc = R(size(R))
Ng = size(neG, 1)
call assert(size(neG, 2) == Ng)
call assert(size(neG, 3) == Ng)
dk = 2*pi/L

! We prepare the values of the radial Fourier transform on a 3D grid
nk(0) = 0
do i = 1, 3*(Ng/2+1)**2
    w = sqrt(real(i, dp))*dk
    nk(i) = 4*pi/(L**3 * w**2) * w*integrate(Rp, R*sin(w*R)*ne)
end do

! We fill out the 3D grid using the values from nk
do i = 1, Ng
do j = 1, Ng
do k = 1, Ng
    idx = (i-1-Ng*nint((i-1.5_dp)/Ng))**2 + (j-1-Ng*nint((j-1.5_dp)/Ng))**2 &
            + (k-1-Ng*nint((k-1.5_dp)/Ng))**2
    neG(i, j, k) = nk(idx)
end do
end do
end do
end subroutine

real(dp) pure function integrate(Rp, f) result(s)
real(dp), intent(in) :: Rp(:), f(:)
s = integrate_trapz_1(Rp, f)
end function

subroutine free_energy_min(Nelec, Natom, L, G2, T_au, VenG, ne, energy_eps, &
        Eee, Een, Ts, Exc, Etot, cg_iter)
! Minimize the electronic free energy using the initial condition 'ne'. Returns
! the ground state in 'ne'. The free energy is returned in Etot, and it's
! components are returned in Eee, Een, Ts and Exc. The relation is:
!
!     Etot = Ts + Een + Eee + Exc
!
real(dp), intent(in) :: Nelec ! Number of electrons
integer, intent(in) :: Natom ! Number of atoms
real(dp), intent(in) :: L, G2(:, :, :), T_au, energy_eps
real(dp), intent(inout) :: ne(:, :, :)
complex(dp), intent(in) :: VenG(:, :, :)
real(dp), intent(out) :: Eee, Een, Ts, Exc, Etot
integer, intent(out) :: cg_iter ! # of CG iterations needed to converge

integer :: Ng
real(dp), allocatable :: free_energies(:)
real(dp), allocatable, dimension(:, :, :) :: Hpsi, &
    psi, psi_, psi_prev, ksi, ksi_prev, phi, phi_prime, eta
integer :: iter, max_iter
real(dp) :: mu, last2, last3, brent_eps, free_energy_, &
    gamma_d, gamma_n, theta, theta_a, theta_b, theta_c, fa, fb, fc
real(dp) :: f2
real(dp) :: psi_norm
integer :: update_type
!real(dp) :: A, B
real(dp) :: t1, t2, t11, t12
real(dp), dimension(size(VenG,1), size(VenG,2), size(VenG,3)) :: Ven, C1, C2, C3
logical :: func_use_fft

brent_eps = 1e-3_dp
max_iter = 2000
update_type = update_polak_ribiere
func_use_fft = .true.

last2 = 0
last3 = 0

Ng = size(ne, 1)

allocate(Hpsi(Ng, Ng, Ng))
allocate(psi(Ng, Ng, Ng))
allocate(psi_(Ng, Ng, Ng))
allocate(psi_prev(Ng, Ng, Ng))
allocate(phi(Ng, Ng, Ng))
allocate(phi_prime(Ng, Ng, Ng))
allocate(ksi(Ng, Ng, Ng))
allocate(ksi_prev(Ng, Ng, Ng))
allocate(eta(Ng, Ng, Ng))

psi = sqrt(ne)
psi_norm = integral(L, psi**2)
print *, "Initial norm of psi:", psi_norm
psi = sqrt(Nelec / psi_norm) * psi
psi_norm = integral(L, psi**2)
print *, "norm of psi:", psi_norm
! This returns H[n] = delta F / delta n, we save it to the Hpsi variable to
! save space:
call free_energy(L, G2, T_au, VenG, psi**2, Eee, Een, Ts, Exc, free_energy_, &
    Hpsi, calc_value=.false., calc_derivative=.true.)
! Hpsi = H[psi] = delta F / delta psi = 2*H[n]*psi, due to d/dpsi = 2 psi d/dn
Hpsi = Hpsi * 2*psi
mu = 1._dp / Nelec * integral(L, 0.5_dp * psi * Hpsi)
ksi = 2*mu*psi - Hpsi
phi = ksi
phi_prime = phi - 1._dp / Nelec *  integral(L, phi * psi) * psi
eta = sqrt(Nelec / integral(L, phi_prime**2)) * phi_prime
theta = pi/2
!print *, "Summary of energies [a.u.]:"
!print "('    Ts   = ', f14.8)", Ts
!print "('    Een  = ', f14.8)", Een
!print "('    Eee  = ', f14.8)", Eee
!print "('    Exc  = ', f14.8)", Exc
!print *, "   ---------------------"
!print "('    Etot = ', f14.8, ' a.u.')", free_energy_
allocate(free_energies(max_iter))
gamma_n = 0
do iter = 1, max_iter
    t1 = clock()
    fft_counter = 0
    fft_time = 0
    !Hpsi = Hpsi/(2*psi)
    ! Formula (36) in Jiang & Yang
    !A = integral(L, psi*Hpsi*psi) - integral(L, eta*Hpsi*eta)
    !B = 2*integral(L, eta*Hpsi*psi)
    !print *, "theta? =", 0.5_dp * atan(B/A)
    theta_a = 0
    theta_b = mod(theta, 2*pi)
    t11 = clock()
    if (iter == 1) then
        if (func_use_fft) then
            call bracket(func_fft, theta_a, theta_b, theta_c, fa, fb, fc, 100._dp, 20, verbose=.false.)
            call brent(func_fft, theta_a, theta_b, theta_c, brent_eps, 50, theta, &
            free_energy_, verbose=.false.)
        else
            call fourier2real(VenG, Ven)
            call precalc_C1C2C2(G2, psi, eta, C1, C2, C3)
            call bracket(func_nofft, theta_a, theta_b, theta_c, fa, fb, fc, 100._dp, 20, verbose=.false.)
            call brent(func_nofft, theta_a, theta_b, theta_c, brent_eps, 50, theta, &
            free_energy_, verbose=.false.)
        end if
    else
        call bracket(func_fft, theta_a, theta_b, theta_c, fa, fb, fc, 100._dp, 20, verbose=.false.)
        call parabola_vertex(theta_a, fa, theta_b, fb, theta_c, fc, theta, f2)
    end if
    t12 = clock()
    ! TODO: We probably don't need to recalculate free_energy_ here:
    psi_prev = psi
    psi = cos(theta) * psi + sin(theta) * eta
    call free_energy(L, G2, T_au, VenG, psi**2, Eee, Een, Ts, Exc, &
        free_energy_, Hpsi, calc_value=.true., calc_derivative=.true.)
!    print *, "Iteration:", iter
!    psi_norm = integral(L, psi**2)
!    print *, "Norm of psi:", psi_norm
!    print *, "mu =", mu
!    print *, "|ksi| =", sqrt(gamma_n)
!    print *, "theta =", theta
!    print *, "Summary of energies [a.u.]:"
!    print "('    Ts   = ', f14.8)", Ts
!    print "('    Een  = ', f14.8)", Een
!    print "('    Eee  = ', f14.8)", Eee
!    print "('    Exc  = ', f14.8)", Exc
!    print *, "   ---------------------"
    free_energies(iter) = free_energy_ / Natom
    if (iter > 1) then
        last2 = maxval(free_energies(iter-1:iter)) - &
            minval(free_energies(iter-1:iter))
    end if
    if (iter > 2) then
        last3 = maxval(free_energies(iter-2:iter)) - &
            minval(free_energies(iter-2:iter))
    end if
    print "('# ', i3, ' Etot/atom = ', f18.8, ' eV; last2 = ', es10.2, ' last3 = ',es10.2)", &
        iter, free_energy_ * Ha2eV / Natom, last2 * Ha2eV, last3 * Ha2eV
    if (iter > 3) then
        if (last3 < energy_eps) then
            ne = psi**2
            Etot = free_energy_
            cg_iter = iter
            return
        end if
    end if
    Hpsi = Hpsi * 2*psi ! d/dpsi = 2 psi d/dn
    mu = 1._dp / Nelec * integral(L, 0.5_dp * psi * Hpsi)
    ksi_prev = ksi
    ksi = 2*mu*psi - Hpsi
    select case(update_type)
        case(update_fletcher_reeves) ! Fletcher-Reeves
            gamma_n = integral(L, ksi**2)
        case(update_polak_ribiere)   ! Polak-Ribiere
            gamma_n = max(integral(L, ksi*(ksi-ksi_prev)), 0._dp)
        case default
            call stop_error("Unknown update type.")
    end select
    gamma_d = integral(L, ksi_prev**2)
    phi = ksi + gamma_n / gamma_d * phi
    phi_prime = phi - 1._dp / Nelec * integral(L, phi * psi) * psi
    eta = sqrt(Nelec / integral(L, phi_prime**2)) * phi_prime
    t2 = clock()
    print "('time: ', f6.3, ' fft: ', f6.3, ' (', f5.2, '%) # fft: ', i3, ' line search', f6.3)", &
        t2-t1, fft_time, fft_time / (t2-t1) * 100, fft_counter, t12-t11
end do
call stop_error("free_energy_minimization: The maximum number of iterations exceeded.")

contains

    real(dp) function func_nofft(theta) result(energy)
    real(dp), intent(in) :: theta
    call calc_F(L, T_au, Ven, C1, C2, C3, psi, eta, theta, energy)
    end function

    real(dp) function func_fft(theta) result(energy)
    real(dp), intent(in) :: theta
    psi_ = cos(theta) * psi + sin(theta) * eta
    call free_energy(L, G2, T_au, VenG, psi_**2, Eee, Een, Ts, Exc, &
        energy, Hpsi, calc_value=.true., calc_derivative=.false.)
    end function

end subroutine

subroutine forces(L, G, nG, Ven0G)
real(dp), intent(in) :: L, G(:, :, :, :)
complex(dp), intent(in) :: nG(:, :, :), Ven0G(:, :, :)
! Ven0G .... the potential returned from:
! radial_density_fourier(R, V, L, Z, Ven0G)
! x, y, z ... the ion position
real(dp) :: x, y, z, Fx, Fy, Fz
complex(dp) :: tmp(size(Ven0G, 1), size(Ven0G, 2), size(Ven0G, 3))
x = 1
y = 0
z = 0
tmp = i_ * nG * Ven0G * exp(-i_*(G(:,:,:,1)*x+G(:,:,:,2)*y+G(:,:,:,3)*z))
! Force: only take the real component of this:
Fx = real(sum(G(:,:,:,1) * L**3 * tmp), dp)
Fy = real(sum(G(:,:,:,2) * L**3 * tmp), dp)
Fz = real(sum(G(:,:,:,3) * L**3 * tmp), dp)
end subroutine

subroutine vtk_save(filename, Xn, f)
character(len=*), intent(in) :: filename
real(dp), intent(in) :: Xn(:, :, :, :), f(:, :, :)

integer :: Nx, Ny, Nz
integer :: u

Nx = size(Xn, 1)
Ny = size(Xn, 2)
Nz = size(Xn, 3)
call assert(size(Xn, 4) == 3)

call assert(size(f, 1) == Nx)
call assert(size(f, 2) == Ny)
call assert(size(f, 3) == Nz)

open(newunit=u, file=filename, status="replace")
write(u, "(a)") "# vtk DataFile Version 2.0"
write(u, "(a)") "Function at a real space uniform grid"
write(u, "(a)") "ASCII"
write(u, "(a)") "DATASET RECTILINEAR_GRID"
write(u, "(a, i8, i8, i8)") "DIMENSIONS", Nx, Ny, Nz

write(u, "(a, ' ', i8, ' ', a)") "X_COORDINATES", Nx, "float"
write(u, *) Xn(:, 1, 1, 1)
write(u, "(a, ' ', i8, ' ', a)") "Y_COORDINATES", Ny, "float"
write(u, *) Xn(1, :, 1, 2)
write(u, "(a, ' ', i8, ' ', a)") "Z_COORDINATES", Nz, "float"
write(u, *) Xn(1, 1, :, 3)

write(u, "(a, ' ', i10)") "POINT_DATA", size(f)
write(u, "(a)") "SCALARS MyFunction float"
write(u, "(a)") "LOOKUP_TABLE default"
write(u, *) f

close(u)
end subroutine

end module


module fqs_quaternion

    use fqs_types,        only: wp
    use fqs_utility_base, only: check_division_scalar

    implicit none

    private

    type, public :: quat_t
        real(wp) :: w = 0.0_wp
        real(wp) :: x = 0.0_wp
        real(wp) :: y = 0.0_wp
        real(wp) :: z = 0.0_wp
    contains
        generic, public :: operator(*) => multiply
        generic, public :: operator(*) => multiply_scalar_left
        generic, public :: operator(*) => multiply_scalar_right
        generic, public :: operator(/) => divide_by_scalar

        procedure, pass(self) :: multiply
        procedure, pass(self) :: multiply_scalar_left
        procedure, pass(self) :: multiply_scalar_right
        procedure, pass(self) :: divide_by_scalar

        procedure, pass(self) :: inv  => inverse
        procedure, pass(self) :: conj => conjugate
        procedure, pass(self) :: norm => normal
        procedure, pass(self) :: sum2 => sum_of_squares
        procedure, pass(self) :: to_array => to_array
    end type quat_t

    interface quat_t
        procedure :: new_from_values
    end interface quat_t

contains

    ! quat_t class constructors
    ! -------------------------------------------------------------------------

    elemental function new_from_values(w,x,y,z) result(q)
        real(wp), intent(in) :: x, y, z, w
        type(quat_t) :: q
        q%w = w
        q%x = x
        q%y = y
        q%z = z
    end function new_from_values

    ! quaterion_t class methods
    ! -------------------------------------------------------------------------

    elemental function multiply(self,p) result(q)
        class(quat_t), intent(in) :: self
        type(quat_t),  intent(in) :: p
        type(quat_t) :: q
        q%w = self%w * p%w - self%x * p%x - self%y * p%y - self%z * p%z
        q%x = self%w * p%x + self%x * p%w + self%y * p%z - self%z * p%y
        q%y = self%w * p%y - self%x * p%z + self%y * p%w + self%z * p%x
        q%z = self%w * p%z + self%x * p%y - self%y * p%x + self%z * p%w
    end function multiply

    elemental function multiply_scalar_right(self,scalar) result(q)
        class(quat_t), intent(in) :: self
        real(wp),      intent(in) :: scalar
        type(quat_t) :: q
        q%w = scalar * (self % w)
        q%x = scalar * (self % x)
        q%y = scalar * (self % y)
        q%z = scalar * (self % z)
    end function multiply_scalar_right

    elemental function multiply_scalar_left(scalar,self) result(q)
        class(quat_t), intent(in) :: self
        real(wp),      intent(in) :: scalar
        type(quat_t) :: q
        q%w = scalar * (self % w)
        q%x = scalar * (self % x)
        q%y = scalar * (self % y)
        q%z = scalar * (self % z)
    end function multiply_scalar_left

    elemental function divide_by_scalar(self,scalar) result(q)
        class(quat_t), intent(in) :: self
        real(wp),      intent(in) :: scalar
        type(quat_t) :: q
        q%w = self%w / scalar
        q%x = self%x / scalar
        q%y = self%y / scalar
        q%z = self%z / scalar
    end function divide_by_scalar

    elemental function inverse(self) result(inv)
        class(quat_t), intent(in) :: self
        type(quat_t)              :: inv

        ! Local variables
        type(quat_t)              :: conj
        real(wp)                  :: sum2
        logical                   :: ok_w
        logical                   :: ok_x
        logical                   :: ok_y
        logical                   :: ok_z

        conj = self % conj()
        sum2 = self % sum2()
        ok_w  = check_division_scalar(conj % w, sum2)
        ok_x  = check_division_scalar(conj % x, sum2)
        ok_y  = check_division_scalar(conj % y, sum2)
        ok_z  = check_division_scalar(conj % z, sum2)
        if (ok_w .and. ok_x .and. ok_y .and. ok_z) then
            inv = ( self % conj() ) / sum2
        else
            inv = quat_t(0.0_wp, 0.0_wp, 0.0_wp, 0.0_wp)
        end if
    end function inverse

    elemental function conjugate(self) result(conj)
        class(quat_t), intent(in) :: self
        type(quat_t)              :: conj
        conj%w =  self%w
        conj%x = -self%x
        conj%y = -self%y
        conj%z = -self%z
    end function conjugate

    elemental function normal(self) result(norm)
        class(quat_t), intent(in) :: self
        real(wp)                  :: norm
        norm = sqrt( self % sum2() )
    end function normal

    elemental function sum_of_squares(self) result(sum2)
        class(quat_t), intent(in) :: self
        real(wp)                  :: sum2
        sum2 = (self % w)**2 + (self % x)**2 + (self % y)**2 + (self % z)**2
    end function sum_of_squares

    function to_array(self) result(a)
        class(quat_t), intent(in) :: self
        real(wp)                  :: a(4)
        a(1) = self % w
        a(2) = self % x
        a(3) = self % y
        a(4) = self % z
    end function to_array

end module fqs_quaternion


use FEMDomainClass
implicit none

type(FEMDomain_) :: Cube
real(real64),allocatable :: displacement(:),YoungModulus(:),PoissonRatio(:)
real(real64),allocatable :: traction(:)
integer(int32),allocatable :: nodeList(:)
real(real64) :: disp_x

disp_x = 0.50d0

call Cube%create("Cube3D",x_num=2, y_num=2,z_num=2)
call Cube%resize(x=2.0d0,y=2.0d0,z=2.0d0)

! select nodes by range
NodeList = Cube%getNodeList(xmax=0.0d0)

call Cube%vtk("cube")

Displacement = zeros(cube%nn()*cube%nd())
YoungModulus = zeros(cube%ne())
PoissonRatio = zeros(cube%ne())

Displacement(3*NodeList(:)-2) = disp_x
call print(reshape(Displacement,cube%nn(),cube%nd()))

YoungModulus = 1.0d0
PoissonRatio = 0.00d0
Traction     = Cube%TractionVector(&
    displacement=displacement,&
    YoungModulus=YoungModulus,&
    PoissonRatio=PoissonRatio)

call print(reshape(Traction,cube%nn(),cube%nd()))
call print(sum(Traction(3*NodeList(:)-2 ) ) )
call Cube%move(x=disp_x,NodeList=NodeList)
call Cube%vtk("cube")
call Cube%vtk("SelectedNode",NodeList=Nodelist)

end