!$id: $
! Copyright (c) 2011-2014 Pierre De Mey, pdm789@gmail.com
!
! This program is free software: you can redistribute it and/or modify
! it under the terms of the GNU Lesser General Public License
! as published by the Free Software Foundation, either version
! 3 of the License, or (at your option) any later version.
!
! This program is distributed in the hope that it will be useful,
! but WITHOUT ANY WARRANTY; without even the implied warranty of
! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
! GNU Lesser General Public License for more details.
!
! You should have received a copy of the GNU Lesser General Public
! License along with this software.  If not, see <http://www.gnu.org/licenses/>.

! ----------------------------------------------------------------------
SUBROUTINE sangoma_arm( &
  nstate, nens, nobs, ndof, Af, Df, R, arm_spect, arm, arm_rep, status &
  ) BIND( C, name="sangoma_arm_" )
! ----------------------------------------------------------------------

! AUTHOR
! Pierre De Mey

! REVISION HISTORY
! 2011-06: P De Mey arm.f Bologna Summer School 2011
! 2013-06: P De Mey arm.f Bologna Summer School 2013
! 2014-03: P De Mey sangoma_arm.F90 initial SANGOMA release

! PURPOSE
! Calculate array modes and associated quantities

! INPUT
! nstate	state size
! nens		Ensemble size
! nobs		number of observations
! ndof		=MIN(nens,nobs) number of d.o.f.s of problem
! Af		forecast ensemble anomalies, defined as Af(nstate,nens)
! Df		same as af in data space, defined as Df(nobs,nens)
!		note: df can be directly generated by the model, or
!		linearly calculated as H*Af, H(nobs,nstate) = obs.op.
! R		observation error covariance matrix, def as R(nobs,nobs)

! OUTPUT
! arm_spect	array mode spectrum, defined as arm_spect(ndof)
! arm		array modes, defined as arm(nobs,ndof)
! arm_rep	modal representers, defined as arm_rep(nstate,ndof)
! status	status flag (0=success)

! NOTES
! 1. The actual precision of REAL is to be provided by the compiler.  Only
!    KIND=8 will work with the current version (promotion of REAL to DOUBLE
!    PRECISION) because of the use of Dxxxxx BLAS/LAPACK calls.  This can
!    be changed in a later version.
! 2. State space and data space are n-DIMENSIONal + (optionally) time.
!    - Each state-space sample and data-space sample can contain 
!      information from several instants if desired.
!    - Modal representers can span space *and* time.

  USE, INTRINSIC :: ISO_C_BINDING
  USE sangoma_base, ONLY: REALPREC, INTPREC

  IMPLICIT NONE

! args

  INTEGER(INTPREC), INTENT(in)  :: nstate
  INTEGER(INTPREC), INTENT(in)  :: nens
  INTEGER(INTPREC), INTENT(in)  :: nobs
  INTEGER(INTPREC), INTENT(in)  :: ndof
  REAL(REALPREC), INTENT(in)    :: Af(nstate,nens)
  REAL(REALPREC), INTENT(in)    :: Df(nobs,nens)
  REAL(REALPREC), INTENT(in)    :: R(nobs,nobs)
  REAL(REALPREC), INTENT(out)   :: arm_spect(ndof)
  REAL(REALPREC), INTENT(out)   :: arm(nobs,ndof)
  REAL(REALPREC), INTENT(out)   :: arm_rep(nstate,ndof)
  INTEGER(INTPREC), INTENT(out) :: status

! local

  REAL :: alpha
  INTEGER :: k, l

! Lapack

  CHARACTER :: uplo = "U"
  INTEGER :: lwork, info
  REAL, DIMENSION(:), ALLOCATABLE :: work
  INTEGER, DIMENSION(:), ALLOCATABLE :: ipiv

! ALLOCATABLEs

  REAL, DIMENSION(:,:), ALLOCATABLE :: S, T, R_inv, R_sqrtinv, VT, w

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

  status = 1
  arm_spect = 0.0
  arm = 0.0
  arm_rep = 0.0

  IF ( KIND(Af) /= 8 ) RETURN ! we use REAL(KIND=8) Lapack calls
  IF ( KIND(S) /= 8 ) RETURN ! we use REAL(KIND=8) Lapack calls

! calculate R^-1(nobs,nobs)

  ALLOCATE ( R_inv(nobs,nobs) ); R_inv = R
  ALLOCATE ( ipiv(nobs) )

  lwork = -1
  ALLOCATE ( work(1) )
  CALL dsytrf( uplo,nobs,R_inv,nobs,ipiv,work,lwork,info )
  IF ( info /= 0 ) THEN
    DEALLOCATE ( R_inv,ipiv,work )
    status = 100000+info
    RETURN
  END IF
  lwork = NINT( work(1) )
  DEALLOCATE ( work )

  ALLOCATE ( work(lwork) )
  CALL dsytrf( uplo,nobs,R_inv,nobs,ipiv,work,lwork,info )
  IF ( info /= 0 ) THEN
    DEALLOCATE ( R_inv,ipiv,work )
    status = 200000+info
    RETURN
  END IF
  DEALLOCATE ( work )

  ALLOCATE ( work(nobs) )
  CALL dsytri( uplo,nobs,R_inv,nobs,ipiv,work,info )
  IF ( info /= 0 ) THEN
    DEALLOCATE ( R_inv,ipiv,work )
    status = 300000+info
    RETURN
  END IF
  DEALLOCATE ( work )

  DEALLOCATE ( ipiv )

! calculate R^-1/2 as R_sqrtinv(nobs,nobs)

  ALLOCATE ( R_sqrtinv(nobs,nobs) )
  CALL dsysqr(R_inv,nobs,nobs,R_sqrtinv,info)
  IF ( info /= 0 ) THEN
    DEALLOCATE ( R_inv,R_sqrtinv,ipiv,work )
    status = 400000+info
    RETURN
  END IF
  DEALLOCATE ( R_inv )

! calculate S (e.g. Sakov et al., 2010)
! S(nobs,nens) = alpha*R^⁻1/2(nobs,nobs)*Df(nobs,nens)

  alpha = 1.0/SQRT(FLOAT(nens-1))
  ALLOCATE ( S(nobs,nens) )
  CALL dgemm( "N","N",nobs,nens,nobs,alpha,R_sqrtinv,nobs,Df,nobs,0.0,S,nobs )
  DEALLOCATE ( R_sqrtinv )

! calculate singular values and vectors of S
! - the left singular vectors of S are the eigenvectors of the scaled  
!   representer matrix F = S * S^T, therefore they are the array modes
! - the eigenvalues of F (array mode spectrum) are the squares of the 
!   singular values of S  

  ALLOCATE ( T(nobs,nens) ); T = S ! dgesvd is destructive and we still need S

  ALLOCATE ( VT(ndof,nens) )
  lwork = -1
  ALLOCATE ( work(1) )
  CALL dgesvd( "S","S",nobs,nens,T,nobs,arm_spect,arm,nobs, &
    VT,ndof,work,lwork,info )
  IF ( info /= 0 ) THEN
    DEALLOCATE ( T,VT,work )
    status = 500000+info
    RETURN
  END IF
  lwork = NINT(work(1))
  DEALLOCATE ( work )

  ALLOCATE ( work(lwork) )
  CALL dgesvd( "S","S",nobs,nens,T,nobs,arm_spect,arm,nobs, &
    VT,ndof,work,lwork,info )
  IF ( info /= 0 ) THEN
    DEALLOCATE ( T,VT,work )
    status = 600000+info
    RETURN
  END IF

  DEALLOCATE ( T,VT,work )

  DO k = 1, ndof
    arm_spect(k) = arm_spect(k)**2
    IF ( arm(1,k) < 0.0 ) arm(:,k) = -arm(:,k)
  END DO

! calculate and save modal representers
! arm_rep(nstate,ndof) = alpha*Af(nstate,nens)*S^T(nens,nobs)*arm(nobs,ndof)

  ALLOCATE ( w(nens,ndof) )
  CALL dgemm( "T","N",nens,ndof,nobs,1.0,S,nobs,arm,nobs,0.0,w,nens )
  CALL dgemm( &
    "N","N",nstate,ndof,nens,alpha,Af,nstate,w,nens,0.0,arm_rep,nstate )
  DEALLOCATE ( w,S )

! normal end

  status = 0

CONTAINS

! ----------------------------------------------------------------------
SUBROUTINE dsysqr ( A, n, ldA, B, info )
! ----------------------------------------------------------------------

! AUTHOR
! Pierre De Mey

! REVISION HISTORY
! 2011-06: P De Mey initial release
! 2014-03: P De Mey contained in sangoma_arm.F90

! PURPOSE
! Calculate the square root B of a positive semidefinite symmetric matrix A

! INPUT
! A		positive semidefinite symmetric matrix A(ldA,n)
! n		the order of the matrix A
! ldA		the leading dimension of the array A

! OUTPUT
! B		square root of A, B(ldA,n)
! info		status flag (0=success)

! NOTE
! algorithm:
!   [V,D] = eig(A)
!   sqrtm(A) = V*diag(sqrt(diag(D)))*V'
! uses BLAS/LAPACK calls and notations
! please note that if something goes wrong in dsyev(), xerbla() will
! cite DSYEV as being the source of the problem, not DSYSQR

  IMPLICIT NONE

! args

  INTEGER, INTENT(in) :: n, ldA
  REAL, DIMENSION(ldA,n), INTENT(in) :: A
  REAL, DIMENSION(ldA,n), INTENT(out) :: B
  INTEGER, INTENT(out) :: info

! local

  INTEGER :: i, k, l

! Lapack

  CHARACTER :: jobz = "V", uplo = "U"
  CHARACTER(LEN=6) :: srname = "DSYSQR"
  INTEGER :: lwork
  REAL, DIMENSION(:), ALLOCATABLE :: work

! ALLOCATABLEs

  REAL, DIMENSION(:), ALLOCATABLE :: spect
  REAL, DIMENSION(:,:), ALLOCATABLE :: V, mat1, mat2

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

! check args

  IF ( n <= 0 ) THEN
    info = 2
    CALL xerbla(srname,info)
    RETURN
  ELSE IF ( ldA < n ) THEN
    info = 3
    CALL xerbla(srname,info)
    RETURN
  END IF

! solve eigenproblem

  ALLOCATE ( V(ldA,n),spect(n) )
  V = A

  lwork = -1
  ALLOCATE ( work(1) )
  CALL dsyev( jobz,uplo,n,V,ldA,spect,work,lwork,info )
  IF ( info /= 0 ) THEN
    DEALLOCATE ( V,spect,work )
    CALL xerbla( srname,info )
    RETURN
  END IF
  lwork = nint(work(1))
  DEALLOCATE ( work )

  ALLOCATE ( work(lwork) )
  CALL dsyev( jobz,uplo,n,V,ldA,spect,work,lwork,info )
  DEALLOCATE ( work )
  IF ( info /= 0 ) THEN
    DEALLOCATE ( V,spect )
    CALL xerbla( srname,info )
    RETURN
  END IF

! form diag(sqrt(diag(D)))

  ALLOCATE ( mat1(n,n) )
  mat1 = 0.0
  DO i = 1, n; IF ( spect(i) > 0 ) mat1(i,i) = SQRT(spect(i)); END DO
  DEALLOCATE ( spect )

! calculate square root

  ALLOCATE ( mat2(n,n) )
  CALL dgemm( "N","T",n,n,n,1.0,mat1,n,V,ldA,0.0,mat2,n )
  DEALLOCATE (mat1)
  CALL dgemm( "N","N",n,n,n,1.0,V,ldA,mat2,n,0.0,B,ldA )
  DEALLOCATE ( mat2,V ) 

! normal end

  info = 0

END SUBROUTINE dsysqr

END SUBROUTINE sangoma_arm