    module assimilation
!  SYNOPSIS
!    use assimilation
!
!  DESCRIPTION
!  The routines in this package implement Brian Hunt's and Ed Ott's
!  versions of the LEKF data assimilation scheme.
!
!  REQUIRED PACKAGES
!
!  grid_model, diagnostics - in main source directory
!  lapackmt, blasmt - in system object file library
!
!  ROUTINES IMPLEMENTED
!
!   subroutine assimilate(loc,nsol,w,diag)
!     and other private support routines.
!  The only publicly callable routine in this module is ASSIMILATE.
!
!  AUTHOR
!  Eric Kostelich, Dept. of Math. & Stat., Arizona State U.,
!  Tempe, AZ 85287-1804, kostelich@asu.edu
!
!  MODIFICATION HISTORY
!  08/08/05: Corrected normalization of Eb^T in Ed's method (the factor is
!    missing in the Szunyogh et al. paper); added optional reporting of
!    the singular values of (k-1)I+Eb^T*inv(R)*Eb in Brian's method and its
!    analogue in Ed's method.
!  08/06/05: Added enhanced variance inflation to original Ott method.
!  07/16/05: Added the original Ott et al. method as a backup in case
!    Brian's method fails for any reason.  Corrected computation E^b.
!  06/26/05: Rewrote enhanced inflation in terms of multiples of the largest
!    eigenvalue of Brian's F; rewrote EIGEN_FIXUP.
!  06/24/05: altered enhanced inflation to work on selected variables only;
!    rewrote module to re-try assimilation without enhanced inflation when
!    numerical problems arise; corrected error in multiplicative inflation.
!  06/21/05: added enhanced variance inflation.
!  03/11/05: Removed check for "singular" Pa~ in factor of a check of the
!    forward error bound: we conclude that the computed value of Pa~ is
!    faulty if the relative forward error is larger than 1.0e-04.
!  01/19/05: Changed variable names to be consistent with the July 30, 2004
!    version of Brian's writeup describing the more efficient algorithm.
!    Added the analysis update method of Zhang, Snyder, and Sun.
!    Incorporated changes to model grid handling; this version no longer needs
!    to worry about managing the analysis grid.
!  11/23/04: Added update to analysis ensemble and ensemble mean.  Added
!    heuristic check for obviously garbled values and a new return code to
!    indicate such.
!  09/27/04: Internal changes to accommodate the change of indexing in
!    the model grid.  This new version should be more cache-friendly.
!  09/12/04: Removed ASSIM_ORIG and modified Brian's method to return
!    only the matrix E that gives the appropriate linear combinations
!    of the background ensembles.  This matrix can be used to generate
!    the analysis ensemble and to update the surface boundary conditions,
!    as well as the unassimilated variables.
!    Also removed support for adding "random" components to the background.
!    This version of the routine is preparatory to the eventual "4-d"
!    implementation.
!  08/04/04: Revamped ASSIM_ORIG to be more consistent with the final
!    Tellus version of the first LEKF paper.  Summary of changes:
!    (1) Pa^ is computed using Eq. (23) of the paper.
!    (2) A is computed using Eq. (A-19).
!    (3) Xa^ is computed using Eq. (34).
!    Also made miscellaneous corrections to the comments in ARRIM_ORIG.
!  07/30/04: Added the Eispack eigenvalue/vector routines as a backup
!    when DSYEVR fails to converge.  Fixed the simple covariance scheme in
!    Brian's method and in the old method.
!  06/28/04: added extra checks to ensure that variable and observation
!    types agree.
!  06/25/04: optimization to operate on just the variables comprising the
!    center column of the patch in Brian's method.
!  06/16/04: added original method as a backup to Brian's.
!  06/11/04: first implementation.
!
!  Required packages:
!
    use grid_obs
    use diagnostics
    use blas
!
    implicit none
    private
    public::assimilate
    contains
!------------------------------------------------------------------------------
    subroutine assimilate(loc,nsol,w,diag)
!  ASSIMILATE - Perform the Ott et al. data assimilation step for one spatial
!  patch of data.  Attempt Brian's optimized method first, and if that
!  fails, then use the original version.
!  Arguments:
!  LOC  :  (x,y,z) grid coordinates defining the center of the patch.
!  NSOL : number of ensemble members
!  W := NSOL x NSOL array giving linear combination of background covariance
!     needed to give the analysis at the center of the patch.  See Brian's
!     writeup.
!  DIAG := diagnostics include the return code DIAG%INFO, which is
!     0 unless an error occurs.
!     Error indications (each of which is positive):
!     INFLATION_FAIL: could not compute inv(I+F/eps) for enhanced var inflation
!     PA_NO_EV:  could not compute square root of analysis covariance matrix
!     P_NEG_EV:  error: negative eigenvalue from covariance matrix
!     NO_OBS_UPDATE:  no observations available to do the update.
!     Also includes information about reciprocal condition number of Pa~
!     and number of observations in the patch.
!
    type(patch_index),intent(in)::loc
    integer,intent(in)::nsol
    real(D),intent(out)::w(nsol,nsol)
    type(patch_diag),intent(inout)::diag
!
!  Local variables
!
    integer::patchsz,nc
!
    patchsz=patchsize(loc)
    call center_column_height(loc,nc)
    call brian(loc,nsol,patchsz,nc,w,diag)
    if(diag%info.ne.0.and.diag%info.ne.NO_OBS_UPDATE) then
       diag%brian_failed=.true.
       call ed(loc,nsol,patchsz,nc,w,diag,nsol-1)
    else
       diag%brian_failed=.false.
    endif
    return
    end subroutine assimilate
!---------------------------------------------------------------------------
    subroutine brian(loc,nsol,patchsz,nc,w,diag)
!  BRIAN - Brian Hunt's implementation of the Ott et al. data assimilation
!  step for one spatial patch of data.
!  This version avoids direct eigenvector decompositions of the background
!  covariance matrix.
!
!  Arguments:
!  LOC  :  (x,y,z) grid coordinates defining the center of the patch.
!  NSOL  : number of ensemble members
!  PATCHSZ  : patch size; number of dynamical variables in the patch
!  NC  : height of center column at logical center of patch
!  W := NSOL x NSOL array giving linear combination of background covariance
!     needed to give the analysis at the center of the patch.  See Brian's
!     writeup.
!  DIAG := diagnostics include the return code DIAG%INFO, which is
!     0 unless an error occurs.
!     Error indications (each of which is positive):
!     INFLATION_FAIL: could not compute inv(I+F/eps) for enhanced var inflation
!     PA_NO_EV:  could not compute square root of analysis covariance matrix
!     P_NEG_EV:  error: negative eigenvalue from covariance matrix
!     NO_OBS_UPDATE:  no observations available to do the update.
!     Also includes information about reciprocal condition number of Pa~
!     and number of observations in the patch.
!
    type(patch_index),intent(in)::loc
    integer,intent(in)::nsol,patchsz,nc
    real(D),intent(out)::w(nsol,nsol)
    type(patch_diag),intent(inout)::diag
!
!  Local storage - preparatory steps (1-3).
!
!  Background ensemble: For implementation efficiency, the background difference
!  array (Xb in Brian's notation) is NSOL x PATCHSIZE (K x M).
!  Brian assumes M x K.  TYPELIST(J,K) holds the dynamical type of the
!  value in XBACK(J,K) and XBAR(K).
!
    real(D)::xbar(patchsz)  ! background mean over the patch
    real(D)::xback(nsol,patchsz)  ! background ensemble differences over patch
    integer::typelist(nsol,patchsz)
    real(D)::cmean(nc)  ! patch mean at logical center
    real(D)::center(nsol,nc)  ! analysis at center of patch
!
!  Observations; the arrays are allocated in COPY_FROM_OBS.
!
    real(D),allocatable::c(:,:)  ! NSOL x NOBS array = Eb^T*inv(R)
    real(D)::cyb(nsol)  !  C * (YOBS - H(XBAR)) in writeup
    real(D)::ceb(nsol,nsol)  ! C * Eb
    real(D),allocatable::yobs(:)  ! local copy of observations
    real(D),allocatable::ebtr(:,:)  ! H(x(k)) - HBAR
    real(D),allocatable::hbar(:)  ! H(XBAR) = Y^b in writeup
    real(D),allocatable::rdiag(:)  ! obs covariance matrix R (assumed diagonal)
    !!real(D),allocatable::hx(:,:)  ! Sun only; avoids compiler segfault
!
!  Analysis ensemble
!
    type(inflation_info)::inflate
    real(D)::pa(nsol,nsol)  ! analysis covar mtx; see Brian's writeup.
    real(D)::s
!
!  Miscellaneous
!
    real(D)::badval
    integer::nobs,info,j,error,var,varlist(NAVARS)
    logical::debug_dump
!
    info=0 
    debug_dump=loc.eq.dump_loc
!
!  Step 1.  Get differences of background ensemble from background mean
!  over the patch.  XBACK (=Xb^T in Brian's writeup) is NOT normalized by 
!  the square root of the rank of the ensemble, as it is in the original
!  Ott el al. paper.
!
    call copy_from_ensemble(loc,xback,nsol,patchsz,typelist)
    if(ensemble_diffs_precomputed()) then
       call copy_from_mean(loc,xbar,patchsz)
    else
       do j=1,patchsz
          xbar(j)=sum(xback(:,j))/nsol
          xback(:,j)=xback(:,j)-xbar(j)
       enddo
    endif
!
!  Check that the columns of Xb (=rows of xback) sum to zero
!  if great verbosity is requested.
!
    if(verbosity().ge.VERB_CHATTY) then
       write(LOGFILE,100,advance='no') loc
       s=0.0
       do j=1,patchsz
          s=max(s,abs(sum(xback(:,j))))
       enddo
       write(LOGFILE,101) ' ||colsum(xback)||: ',s
    endif
100 format('brian: loc:',4i5)
101 format(a,es11.3)
!
!  Step 2.  Get observations Y and respective interpolations from
!  the model grid.  EBTR=(Eb)^T and YB = H*XBAR, using Brian's notation.
!  If we have no observations in the local patch, then just return the
!  background.
!
    call copy_from_obs(loc,yobs,hbar,rdiag,nsol,nobs,dh=ebtr)  ! HX: Sun only
    diag%nobs=nobs
    if(nobs.eq.0) then
       info=NO_OBS_UPDATE
       goto 911  ! return background
    endif
!
    if(debug_dump) then
       open(DUMPUNIT,file='assim.dump',status='replace',recl=1024)
       write(DUMPUNIT,500) patchsz,nsol,nobs
500    format(3i8)
       write(DUMPUNIT,501) 0,xbar
       do j=1,patchsz
          write(DUMPUNIT,501) j,xback(:,j)
       enddo
501    format(i5,40es24.15/('>    ',40es24.15))
    endif
!
!  Step 3.
!  Form C = Eb^T*inv(R).  If R is variable, then it may be more efficient
!  to solve R*transpose(C)=H~ for C than to compute inverse(R) directly.
!
    allocate(c(nsol,nobs))
    do j=1,nobs
       c(:,j)=ebtr(:,j)/rdiag(j)
    enddo
!
!  Step 4.  Form Pa~T either with or without enhanced variance inflation.
!  If enhanced inflation causes problems (singular analysis covariance
!  matrix or garbage values in the analysis), then abandon it and attempt
!  the calculaton without enhanced inflation.
!
    ceb=matmul(c,transpose(ebtr))  ! C*Eb = Eb^T*inv(R)*Eb
    yobs=yobs-hbar  ! y0 - yb
    cyb=matmul(c,yobs)  ! C(y0-yb)
    call inflation(loc,inflate,varlist)
    info=0
    if(inflate%enhanced.gt.0.0) then  ! perform enhanced inflation
       call enhanced(ceb,nsol,patchsz,xback,typelist,varlist, &
          inflate%enhanced,pa,diag)
       if(diag%info.eq.0) call update(loc,nsol,patchsz,nc,xback,xbar,pa,cyb, &
          inflate,w,diag%edim,cmean,center,info)
    endif
!
!  If we have problems, re-try without enhanced inflation.
!
    if(info.ne.0.or.inflate%enhanced.le.0.0) then
       do j=1,nsol
          ceb(j,j)=ceb(j,j)+(nsol-1)  ! (k-1)*I + C*Eb
       enddo
       if(allocated(diag%iceb_sing)) call singular_vals(ceb,diag%iceb_sing,nsol)
       call invert(ceb,nsol,pa,diag%rcondice,info)  ! Pa~ = inv((k-1)/rhoI+C*Eb)
       call update(loc,nsol,patchsz,nc,xback,xbar,pa,cyb,inflate, &
          w,diag%edim,cmean,center,info)
    endif
    diag%info=info
    if(debug_dump) then
       write(DUMPUNIT,'(a)') 'pa'
       do j=1,nsol
          write(DUMPUNIT,501) j,pa(j,:)
       enddo
    endif
!
!  If there are no errors, then we're done!
!
    if(debug_dump) close(DUMPUNIT)
    if(info.eq.0) return
!
!  If numerical problems have caused us to compute junk, then attempt to
!  report some useful data by re-trying the failed operation to get
!  some statistics.
!
    if(info.eq.INF_OR_NAN.or.info.eq.BAD_VALUE) then
       call copy_to_analysis(loc,nc,nsol,center,error,var,badval)
       if(error.eq.INF_OR_NAN) then
          write(IERR,110) 'inf or nan',TYPEOF(var),loc
       else
          write(IERR,110,advance='no') 'bad value',TYPEOF(var),loc
          write(IERR,111) badval
       endif
       info=GARBAGE_VALUES
110    format('assim: ',a,': var: ',a1,' at: ',3i3)
111    format(' value: ',es13.4)
    endif
!
911 continue
    diag%info=info
    call return_background(loc,xbar,xback,nsol,patchsz,nc,w)
    return
    end subroutine brian
!---------------------------------------------------------------------------
    subroutine update(loc,nsol,patchsz,nc,xback,xbar,pa,cyb,inflate,w,edim, &
       cmean,center,info)
!  UPDATE - Perform Steps 5-9 of Brian's update procedure.
!  LOC  : location of the logical center of the patch
!  NSOL  : number of ensemble solutions
!  PATCHSZ : number of dynamical variables in the patch.
!  NC  : number of variables at the logical center of the patch.
!  XBACK  : (J,K)th element holds Jth component of Kth ensemble diff
!  XBAR  : ensemble mean over the patch.
!  PA  : estimated analysis covariance matrix
!  CYB  : C * (YOBS - H(XBAR)); see Brian's writeup.
!  INFLATE : info on multiplicative and Snyder-Zhang variance inflation
!  W  :=  NSOL x NSOL matrix giving the linear combination of ensemble
!     differences that generate the analysis differences.
!  EDIM := a surrogate for the E-dimension over the patch.  The E-dimemsion
!     is defined in terms of the eigenvalues of the background covariance
!     matrix; we use the analysis covariance matrix instead.
!  CMEAN := analysis mean at the logical center of the patch.
!  CENTER := Kth row holds Kth analysis member at logical center of patch.
!  INFO := set to zero unless an error occurs.  If an error occurs, then
!     the output variables are undefined.
!  Eric: this calling interface could be simplified by having a get_center_col
!  routine in GRID_MODEL; also have RETURN_BACKGROUND do all the work so
!  we don't need to pass around CENTER and CMEAN.
!
    type(patch_index),intent(in)::loc
    integer,intent(in)::nsol,patchsz,nc
    real(D),intent(in)::xback(nsol,patchsz)  ! ensemble diffs (over patch)
    real(D),intent(in)::xbar(patchsz)  ! background mean (over patch)
    real(D),intent(in)::pa(nsol,nsol)  ! estimated covariance mtx Pa~
    real(D),intent(in)::cyb(nsol)   ! C * (YOBS - H(XBAR)) in writeup
    type(inflation_info),intent(in)::inflate
    real(D),intent(out)::w(nsol,nsol)  ! the update matrix W
    real(D),intent(out)::edim  ! surrogate for E-dimension on patch
    real(D),intent(out)::cmean(nc)   ! analysis mean at center column
    real(D),intent(out)::center(nsol,nc)  ! analysis at center of patch
    integer,intent(out)::info
!
!  Local variables
!
    real(D)::xa(nsol,nsol)  !  Xa~; analysis covariance in tilde space
    real(D)::cdiff(nsol,nc)  !  analysis diff at logical center of patch
    real(D)::lambda(nsol)  ! eigenvalues of Pa~
    real(D)::v(nsol)  ! Pa~ *[C*(YOBS-YB)]; step 6 of Brian's writeup
    integer::j,offset,effrank
!
!  Step 5.  Form Xa~=sqrt((k-1)*Pa~).  Do multiplicative variance inflation
!  here if requested.
!
    w=pa*(nsol-1)*(inflate%covar+1.0)  ! W is work array here
    call sym_sqrt(w,nsol,xa,lambda,effrank,info)  ! Xa~ in Brian's writeup
    if(info.ne.0) then   ! can't compute Xa~
       info=PA_NO_EV
       return
    endif
!
!  Use the eigenvalues of Pa~ as a surrogate for the background's E-dimension.
!
    edim=sum(lambda)**2/dot_product(lambda,lambda)
!
!  Step 6. Form v = Pa~ [C(y-y^b)] and add to each column of Xa~, giving W.
!
    v=matmul(pa,cyb)
    do j=1,nsol
       w(:,j)=xa(:,j)+v
    enddo
!
!  Step 7.  Update the background with the analysis.  Since we don't need
!  the entire patch, it's much faster to right-multiply Xb with W
!  only for those rows of Xb that correspond to the dynamical variables
!  in the center of the patch.  XBACK is the transpose of Brian's Xb,
!  so we compute W^T*XBACK.  CDIFF holds the analysis perturbations.
!
    offset=patch_offset(loc,0,0)
    cdiff=matmul(transpose(w),xback(:,offset:offset+nc-1))
!
!  Step 8.  Update the analysis ensemble and the ensemble mean.
!  If requested, first apply the modified update of Zhang, Snyder, and Sun,
!  MWR 132 (May 2004), 1238-1253.  The analysis that we use is a homotopy
!  from the background differences to the analysis differences computed
!  according to our scheme; if alpha=0, then we use our scheme without
!  modification; if alpha=1, we'd return the background spread; if alpha=0.5,
!  then the analysis dfferences are the average of the background and
!  what we have obtained above.
!
    do j=1,nc
       center(:,j)=cdiff(:,j)+xbar(offset+j-1)
       cmean(j)=sum(center(:,j))/nsol
    enddo
    if(inflate%alpha.gt.0.0) then
       do j=1,nc
          center(:,j)=cmean(j)+(1.0-inflate%alpha)*cdiff(:,j)+ &
             inflate%alpha*xback(:,offset+j-1)
       enddo
    endif
!
!  Step 9.  Check that numerical problems haven't caused us to compute junk.
!  If we have, then flag an error.
!
    call copy_to_analysis(loc,nc,nsol,center,info)
    if(info.eq.0) call copy_to_mean(loc,nc,cmean)  ! also store analysis mean
    return
    end subroutine update
!---------------------------------------------------------------------------
    subroutine enhanced(ceb,nsol,patchsz,xback,typelist,var,enh,patilde,diag)
!  ENHANCED - implement enhanced inflation according to Brian's writeup.
!  We regard eigenvalues of F that are less than ENH times the largest
!  eigenvalue as negligible.  Then, we add EPSILON to each of the remaining
!  eigenvalues, where EPSILON is LAMBDA(1)*ENH.  Thus, the inflation at
!  most doubles the size of the nontrivial eigenvalues of F.
!  Smaller values of ENH affect more of the spectrum, but they add smaller
!  increments; larger values affect less of the spectrum but add larger
!  increments.
!
!  Arguments:
!  CEB  : The product (C*Eb).
!  NSOL : number of ensemble solutions; must be > 1
!  PATCHSZ  : number of variables in the patch
!  XBACK  : covariance matrix of all variables in the patch (i.e., the 
!     dynamical variables with the means subtracted off).
!  TYPELIST  : (J,K)th element gives the type of the corresponding variable in
!     XBACK.
!  VAR : list of the variables whose covariances will be used to compute
!     the inflation.  Use 0 for those variables that are omitted;
!     otherwise, some of PRES_TYPE, TEMP_TYPE, UWIND_TYPE, VWIND_TYPE.
!  ENH  : enhanced inflation factor.
!  PATILDE := Brian's Pa~ on success; see writeup.
!  RCOND := reciprocal of estimated condition number of resulting Pa~.
!  DIAG := patch diagnostics; we update the error code and elements relating
!     to the enhanced variance inflation only.  If DIAG%INFO is nonzero on
!     returnkthe contents of PATILDE are undefined.  This can occur 
!     if PATILDE^(-1) is numerically singular.
!
    integer,intent(in)::nsol,patchsz
    real(D),intent(in)::ceb(nsol,nsol),xback(nsol,patchsz)
    real(R)::enh
    integer,intent(in)::typelist(nsol,patchsz),var(NAVARS)
    real(D),intent(out)::patilde(nsol,nsol)
    type(patch_diag),intent(inout)::diag
!
    real(D)::f(nsol,nsol),dx(nsol,patchsz),pinv(nsol,nsol),work(nsol,nsol), &
       u(nsol,nsol),sigma(nsol),g(nsol,nsol),epsilon
    real(D),allocatable::svdwork(:)
    integer::j,lwork,effrank,info
!
    if(enh.le.0.0.or.nsol.le.1) then
       diag%info=INFLATION_FAIL
       return
    endif
!
!  Use only those elements of the covariance array 
!  corresponding to the variables for which the inflation is to be performed.
!
    if(all(var.ne.0)) then
       dx=xback
    else
       dx=0.0 
       do j=1,NAVARS
          if(var(j).gt.0) where(typelist.eq.var(j)) dx=xback
       enddo
    endif
    dx=dx/sqrt(real(nsol-1,kind(dx)))
!
!  Next find the eigenvectors U of F using the singular value decomposition
!  of DX.  (The call to DGESVD destroys the contents of DX, which is why
!  we copy XBACK to it in the event that all dynamical variables participate
!  in the inflation.  (PINV is not referenced in the calls below; see the
!  LAPACK documentation, available online from netlib.org.)
!  Upon successful return, SIGMA**2 is the eigenvalues of F.
!
    call dgesvd('A','N',nsol,patchsz,dx,nsol,sigma,u,nsol,pinv,nsol, &
       work,-1,info)  ! workspace query
    lwork=nint(work(1,1))
    allocate(svdwork(lwork))
    call dgesvd('A','N',nsol,patchsz,dx,nsol,sigma,u,nsol,pinv,nsol, &
       svdwork,lwork,info)
    if(info.lt.0) then
       call fabort('enhanced: illegal value in call to DGESVD')
    else if(info.gt.0) then  ! DGESVD did not converge
       diag%info=INFLATION_FAIL
       return
    endif
    sigma=sigma**2
!
!  Threshold the eigenvalues of F to determine its effective rank and
!  to determine the variance inflation factor EPSILON.
!
    epsilon=enh
    call eigen_fixup(sigma,nsol,effrank,1,info,epsilon)
    diag%effrank=effrank
    diag%smax=sigma(1)
    if(effrank.eq.0) then
       diag%info=INFLATION_FAIL
       return
    endif
    epsilon=enh*sigma(1)
!
!  Compute PINV = inv(I + F/eps) and G = UU^T, where we include only those
!  columns of U that correspond to nontrivial eigendirections of F.
!
    f=matmul(dx,transpose(dx))/epsilon
    do j=1,nsol
       f(j,j)=f(j,j)+1.0
    enddo
    call invert(f,nsol,pinv,diag%rcondpb,info)
    if(info.ne.0) then
       diag%info=INFLATION_FAIL
       return
    endif
    g=matmul(u(:,1:effrank),transpose(u(:,1:effrank)))
!
!  Next obtain inv(P^i), which, from p. 14 of Brian's revised writeup, is
!  (1-k)[inv(I+F/eps) + UU^T - 2I].
!
    pinv=pinv+g
    do j=1,nsol
       pinv(j,j)=pinv(j,j)-2.0
    enddo
    work=ceb+pinv*(1-nsol)
    call invert(work,nsol,patilde,diag%rcondice,info)
    diag%info=merge(0,INFLATION_FAIL,info.eq.0)
    return
    end subroutine enhanced
!-----------------------------------------------------------------------------
    subroutine ed(loc,nsol,patchsz,nc,w,diag,rank)
!  ED - Ed Ott and Brian Hunt's original data assimilation algorithm for
!  one spatial patch of data.
!  The equation numbers follow those in Ott et al.,
!  "A Local Ensemble Kalman Filter for Atmospheric Data Assimilation",
!  Tellus 56A (2004), 415-428.
!
!  Arguments:
!  LOC  :  (x,y,z) grid coordinates defining the center of the patch.
!  NSOL : number of ensemble members
!  PATCHSZ : number of dynamical variables in the patch
!  NC  : center column height (at logical center of patch)
!  W := NSOL x NSOL array giving linear combination of background covariance
!     needed to give the analysis at the center of the patch.  See Brian's
!     writeup.
!  DIAG := diagnostics include the return code DIAG%INFO, which is
!     0 unless an error occurs.
!     Error indications (each of which is positive):
!     INFLATION_FAIL: could not compute inv(I+F/eps) for enhanced var inflation
!     PA_NO_EV:  could not compute square root of analysis covariance matrix
!     P_NEG_EV:  error: negative eigenvalue from covariance matrix
!     NO_OBS_UPDATE:  no observations available to do the update.
!     Also includes information about reciprocal condition number of Pa~
!     and number of observations in the patch.
!  RANK : rank of the projections used (usually NSOL-1 in this implementation).
!
    type(patch_index),intent(in)::loc
    integer,intent(in)::nsol,rank,patchsz,nc
    real(D),intent(out)::w(nsol,nsol)
    type(patch_diag),intent(inout)::diag
!
!  Local storage
!
!  Ensemble data over the local patch
!
    real(D)::xback(nsol,patchsz) ! ensemble background diffs
    real(D)::xbar(patchsz) ! background mean over patch
    integer::typelist(nsol,patchsz)  ! dynamical variable types
!
!  Observations
!
    real(D),allocatable::yobs(:) ! local copy of observations
    real(D),allocatable::hbar(:)  ! H(xbar)
    real(D),allocatable::ebtr(:,:) ! Kth col is H(x(k))-HBAR for Kth ensemble.
    real(D),allocatable::rdiag(:)   ! obs covariance matrix (assumed diagonal)
    integer::nobs
!
!  Covariance matrix and eigenvectors
!
    real(D)::q(patchsz,nsol)  ! Orthonormal eigenvectors of Pb.
    real(D)::lambda(nsol)  ! Eigenvalues of Pb, sorted in descending order.
    real(D)::pbhat(rank)  ! Pb^ in eigenvector coords, which is diagonal
    real(D)::renorm  ! normalization factor for Xb, Xa, and hat versions
    integer::info
    type(inflation_info)::inflate
!
!  Data assimilation step
!
    real(D),allocatable::c(:,:)  ! Eb^T*inv(R)
    real(D)::innov(nsol)  ! C*innovation 
    real(D)::iceb(nsol,nsol)  ! I+C*Eb
    real(D)::icebinv(nsol,nsol)  ! inv(I+C*Eb)
    real(D)::rcond  ! reciprocal condition number of ICEB
    real(D)::pahat(rank,rank) ! analysis covariance mtx in eigenvector coords
    real(D)::xbhat(nsol,rank)  ! transpose of Ed's Xb^
    real(D)::xahat(nsol,rank)  ! transpose of Ed's Xa^
    real(D)::xicebinv(rank,nsol) ! work array
    real(D)::dxamn(rank)  ! adjustment to local analysis mean
!
!  Analysis step
!
    real(D)::pdiff(rank,rank) ! Pa^ - Pb^
    real(D)::work(nsol,rank)
    real(D)::a(nsol,nsol)  ! determines Pa^ in Potter's method
    real(D)::center(nsol,nc)  ! analysis at center of patch
    real(D)::trace
!
    integer::j,offset,effrank
!
!  Get background ensemble mean and perturbations therefrom.
!  Henceforth, XBACK(:,k) holds the difference between the emsemble
!  mean and the Kth ensemble solution in the patch whose logical center
!  is LOC.
!
    call copy_from_ensemble(loc,xback,nsol,patchsz,typelist)
    if(ensemble_diffs_precomputed()) then
       call copy_from_mean(loc,xbar,patchsz)
    else
       do j=1,patchsz
          xbar(j)=sum(xback(:,j))/nsol
          xback(:,j)=xback(:,j)-xbar(j)
       enddo
    endif
!
!  Normalize by the square root of one less than the ensemble size, Eq. (5).
!
    renorm=sqrt(real(rank,kind(renorm)))
    xback=xback/renorm
!
!  Get observations, if any.  If there are none, then return the background.
!  The normalization of Eb^T is missing in Szunyogh et al.
!
    call copy_from_obs(loc,yobs,hbar,rdiag,nsol,nobs,dh=ebtr)
    ebtr=ebtr/sqrt(real(nsol-1,kind(ebtr)))
    if(nobs.eq.0) then
       call return_background(loc,xbar,xback,nsol,patchsz,nc,w,renorm)
       diag%info=NO_OBS_UPDATE
       return
    endif
!
!  Find Q, the eigenvector decomposition of P^b, Eq. (11).
!
    call get_pb_eigen(xback,patchsz,nsol,lambda,q,info)
    if(info.ne.0) then  ! SVD did not converge
       call return_background(loc,xbar,xback,nsol,patchsz,nc,w,renorm)
       diag%info=PB_NO_EV
       return
    else
       diag%rcondpb=lambda(rank)/lambda(1)  ! reciprocal condition number of P^b
       diag%edim=sum(sqrt(lambda))**2/sum(lambda)
    endif
!
!  Next do multiplicative covariance inflation if requested.
!
    pbhat=lambda(1:rank)
    call inflation(loc,inflate)
    if(inflate%covar.gt.0.0) then
       pbhat=pbhat*(inflate%covar+1.0)
       xback=xback*sqrt(inflate%covar+1.0)
    endif
    if(verbosity().ge.VERB_CHATTY) write(IOUT,100) loc,inflate%covar
100 format('ed: loc:',3i4,' covar:',f6.3)
!
!  Compute the analysis covariance matrix using Eq. (A-24) of the 
!  Szunyogh et al paper (all matrices in the hat space):
!  Pa = Xb*inv(I + Eb^T*inv(R)*Eb)*(Eb^T*inv(R))*(yobs-H(xbar)).
!  We'll define the matrix C (as in Brian's writeup) to be Eb^T*inv(R).
!
    allocate(c(nsol,nobs))
    do j=1,nobs
       c(:,j)=ebtr(:,j)/rdiag(j)
    enddo
    yobs=yobs-hbar
    innov=matmul(c,yobs)
    iceb=matmul(c,transpose(ebtr))
    do j=1,nsol
       iceb(j,j)=iceb(j,j)+1.0
    enddo
    if(allocated(diag%iceb_sing)) call singular_vals(iceb,diag%iceb_sing,nsol)
    call invert(iceb,nsol,icebinv,rcond,info)
    if(info.ne.0) then
       diag%info=PA_SINGULAR
       return
    endif
    xbhat=matmul(xback,q(:,1:rank))  ! gives transpose of Ed's Xb^
    xicebinv=matmul(transpose(xbhat),icebinv)
    dxamn=matmul(xicebinv,innov)  ! Eq. (A-25) of Szunyogh et al.
    pahat=matmul(xicebinv,xbhat)
!
!  Enhanced variance inflation if requested.
!
    if(inflate%enhanced.gt.0.0) then  ! Pa^ -> Pa^ + (eps*tr(Pa~)/k)*I, Eq. (42)
       trace=pahat(1,1)
       do j=2,rank
          trace=trace+pahat(j,j)
       enddo
       trace=trace*inflate%enhanced/(nsol-1)
       do j=1,rank
          pahat(j,j)=pahat(j,j)+trace
       enddo
    endif
!
!  Compute Y = I + Xb^T*inv(Pb^)*(Pa^-Pb^)*inv(Pb^)*Xb^, Eq. (41)
!
    pdiff=pahat
    do j=1,rank
       pdiff(j,j)=pdiff(j,j)-pbhat(j)  ! since Pb^ is diagonal
    enddo
    do j=1,rank
       work(:,j)=xbhat(:,j)/pbhat(j)  ! Xb^T*inv(Pb^)
    enddo
    xahat=matmul(work,pdiff)   ! XAHAT is work array here
    a=matmul(xahat,transpose(work))
    do j=1,nsol
       a(j,j)=a(j,j)+1.0
    enddo
    call sym_sqrt(a,nsol,w,lambda,effrank,info)
    if(info.ne.0) then
       call return_background(loc,xbar,xback,nsol,patchsz,nc,w,renorm)
       diag%info=SYM_SQRT_FAIL
       return
    endif
!
!  Get the analysis perturbations, Eq. (34).  W (=Y) is symmetric, and our
!  XAHAT and XBHAT are the transposes of Ed's.  Also undo the normalization;
!  see Eq. (35) of Ott et al.
!
    xahat=matmul(w,xbhat)*renorm
    xbar=xbar+matmul(q(:,1:rank),dxamn)  ! new analysis mean (in model coords)
!
!  Remember that XAHAT is the tranpose of Ed's Xa^.  Save time by applying
!  the update only to the center column of the patch.
!
    offset=patch_offset(loc,0,0)
    center=matmul(xahat,transpose(q(offset:offset+nc-1,1:rank))) ! Eq. (26)
    do j=1,nc
       center(:,j)=center(:,j)+xbar(offset+j-1)
    enddo
    call copy_to_analysis(loc,nc,nsol,center,diag%info)
    if(diag%info.eq.0) then
       call copy_to_mean(loc,nc,xbar(offset:offset+nc-1))
    else
       call return_background(loc,xbar,xback,nsol,patchsz,nc,w,renorm)
    endif
    return
    end subroutine ed
!---------------------------------------------------------------------------
    subroutine return_background(loc,xbar,xback,nsol,patchsz,nc,w,c)
!  RETURN_BACKGROUND - return the corresponding elements of the background.
!  Service routine for cases where the data assimilation routine fails.
!  LOC  : location of logical center of patch
!  XBAR  : background mean over the patch
!  XBACK : holds background perturbations over the patch
!  NSOL : number of the ensemble solutions
!  PATCHSZ : number of dynamical variables in the patch
!  W  := set to the identity matrix.
!  C, if present, undoes any scaling that may have been applied to XBACK.
!
    type(patch_index),intent(in)::loc
    integer,intent(in)::nsol,patchsz,nc
    real(D),intent(in)::xback(nsol,patchsz),xbar(patchsz)
    real(D),intent(out)::w(nsol,nsol)
    real(D),intent(in),optional::c
!
    real(D)::cmean(nc),center(nsol,nc)
    integer::offset,j
!
    offset=patch_offset(loc,0,0)
    do j=1,nc
       cmean(j)=xbar(offset+j-1)
       if(present(c)) then
          center(:,j)=xbar(offset+j-1)+xback(:,offset+j-1)*c
       else
          center(:,j)=xbar(offset+j-1)+xback(:,offset+j-1)
       endif
    enddo
    call copy_to_analysis(loc,nc,nsol,center)
    call copy_to_mean(loc,nc,cmean)
    call identity(w,nsol)
    return
    end subroutine return_background
!---------------------------------------------------------------------------
    subroutine sym_sqrt(a,n,b,lambda,effrank,info)
!  Return the symmetric square root of A in B.  The square root is
!  computed as a product of the orthonormal matrix M of eigenvectors of A:
!  B = (sqrt(lambda)*M)*transpose(M), where lambda is
!  the diagonal matrix of eigenvalues of A.
!  We do not consider a singular matrix to be an error here.
!  Packages required:
!
!  Arguments:
!  A  :  N x N matrix, assumed symmetric.
!  B  := N x N positive symmetric square root of A, if INFO=0.
!  LAMBDA := the N eigenvalues of B (which equal the eigenvalues of A,
!    except that we zero out small negative eigenvalues that arise from
!    numerical roundoff).  Entries undefined if INFO != 0.
!  EFFRANK := the effective rank of B (number of unambiguously positive
!    eigenvalues).  Value undefined if INFO != 0.
!  INFO := set to 0 unless the eigenvector decomposition fails.
!
    integer,intent(in)::n
    real(D),intent(in)::a(n,n)
    real(D),intent(out)::b(n,n),lambda(n)
    integer,intent(out)::effrank,info
!
!  Local variables
!
    real(D)::work(n,n)
    integer::j,nfound
!
    b=a
    call get_eigenspace(b,n,lambda,nfound)  ! B holds eigenvectors of A
    if(nfound.lt.n) then
       info=PA_NO_EV
       return
    endif
    call eigen_fixup(lambda,n,effrank,0,info)
    if(info.ne.0) then
       info=P_NEG_EV
       return
    endif
    do j=1,n
       work(:,j)=b(:,j)*sqrt(lambda(j))
    enddo
    b=matmul(work,transpose(b))   ! gives sqrt(B)
    info=0
    return
    end subroutine sym_sqrt
!---------------------------------------------------------------------------
    subroutine get_eigenspace(a,n,lambda,nfound)
!  GET_EIGEN - Overwrite the symmetric matrix A with its eigenvectors.
!  Arguments:
!  N    :   dimension of A, which is NxN.
!  A  :=:  On input, the covariance matrix.  On output, the Kth
!      column is the eigenvector corresponding to the Kth eigenvalue.
!  LAMBDA := On return, the largest NFOUND eigenvalues of A, from
!      largest to smallest in absolute value.
!  NFOUND := the number of successfully computed eigenvalues and eigenvectors.
!
    integer,intent(in)::n
    real(D),intent(inout)::a(n,n)
    real(D),intent(out)::lambda(n)
    integer,intent(out)::nfound
!
!  Local variables
!
    real(D),parameter::ZERO=0.0
    real(D)::work(70*n) ! work array for dsyev
    real(D)::mu(n)    ! holds eigenvalues from dsyev
    real(D)::abstol   ! convergence tolerance for dsyev
    real(D)::z(n,n)   ! Jth column holds Jth eigenvalue from dsyev
    real(D)::abak(n,n) ! backup copy of A in case DSYEVR doesn't converge
    integer::iwork(10*n),isuppz(2*n)  ! more work arrays for dsyev
    integer::j,info
!
!  DSYEVR and RS return the eigenvalues in ascending order.
!  Note: the convergence of DSYEVR and DSYEVX depends on how ABSTOL
!  is set.  Looser tolerances (e.g., ABSTOL=sqrt(machine epsilon)) often
!  cause an internal failure in DSYEVR, as some eigenvectors don't
!  converge.  The higher tolerance appears to be better, based on
!  experiments with the Lorenz model.
!
    abstol=2.0*tiny(a)
    abak=a
    lambda=0.0
    call dsyevr('v','i','u',n,a,n,ZERO,ZERO,1,n, &
       abstol,nfound,mu,z,n,isuppz,work,size(work),iwork,size(iwork),info)
    if(info.lt.0) then   ! bad argument - definitely a bug in caller
       call fabort('get_eigenspace: bad call to dsyevr')
    else if(info.gt.0) then  ! internal error in dsyevr; try EISPACK routines
       call rs(n,n,abak,mu,1,z,work,work(n+1),info)
    endif
    if(info.eq.0) then
       do j=1,nfound
          lambda(j)=mu(nfound-j+1)
          a(:,j)=z(:,nfound-j+1)
       enddo
    else
       nfound=0
    endif
    return
    end subroutine get_eigenspace
!---------------------------------------------------------------------------
    pure subroutine eigen_fixup(lambda,n,effrank,ignore,info,threshold)
!  EIGEN_FIXUP - LAMBDA is a list of eigenvalues that are supposed to be
!  nonzero (e.g., they come from a symmetric real matrix) and that are
!  sorted in descending order (this is not checked).
!  We check for nontrivial negative eigenvalues, in which case INFO flags
!  an error.  Assuming none are found, then EFFRANK is set to the number
!  of eigenvalues that are larger than LAMBDA(1) times the THRESHOLD.
!  All eigenvalues less than this product are set to zero on ordinary
!  return.  If THRESHOLD is not supplied, then we use the square root
!  of the machine precision.
!
!  Arguments:
!  LAMBDA :=: must be sorted in descending order on entry (not checked).
!    On return, the last IGNORE values are replaced by 0, as are
!    any of the first N-IGNORE values that are less than LAMBDA(1)*THRESHOLD. 
!    (If THRESHOLD is absent, then we use the square root of the machine
!    precision.)  On return, LAMBDA(1:EFFRANK) are the only nonzero values.
!  N  : number of eigenvalues in LAMBDA.
!  EFFRANK := set to number of nonzero eigenvalues after processing.
!  IGNORE : the last IGNORE eigenvalues are assumed zero (and are set to 0).
!     For example, if the matrix in question is expected to have rank
!     at most n-1, for instance, then we'd call this function with IGNORE=1.
!     Otherwise, if the matrix in principle can have full rank, then
!     call with IGNORE=0.
!  INFO := set to 0 unless some element of LAMBDA is negative and fails the
!       smallness test, in which case LAMBDA is unchanged.
!  THRESHOLD : if present, eigenvalues smaller than THRESHOLD*LAMBDA(1)
!   are assumed to be negligible.  We use the maximum of THRESHOLD and the
!   square root of the machine precision.
!
    integer,intent(in)::n,ignore
    real(D),intent(inout)::lambda(n)
    integer,intent(out)::effrank
    integer,intent(out)::info
    real(D),intent(in),optional::threshold
!
!  Local variables
!
    real(D)::thresh
!
!  Check whether we have a nontrivial negative eigenvalue (or a negative
!  threshold).  If we do, then flag an error and return.
!
    if(present(threshold)) then
       thresh=lambda(1)*max(threshold,sqrt(epsilon(lambda)))
    else
       thresh=lambda(1)*sqrt(epsilon(lambda))
    endif
    if(lambda(1).lt.0.0.or.any(lambda(1:n-ignore).lt.-thresh)) then
       info=P_NEG_EV
       return
    endif
!
!  Otherwise assume that negative eigenvalues and suitably small positive
!  eigenvalues correspond to numerical roundoff.
!
    if(ignore.gt.0) lambda(n-ignore+1:n)=0.0
    where(lambda.lt.thresh) lambda=0.0
    effrank=count(lambda.gt.0.0)
    info=0
    return
    end subroutine eigen_fixup
!---------------------------------------------------------------------------
    subroutine invert(a,n,ainv,rcond,info)
!  INVERT - utility routine to compute the inverse of the input array A.
!  On successful returnk AINV holds inv(A) and RCOND holds an estimate of the
!  reciprocal condition number of A.  We regard the computation of the
!  inverse as a failure if the estimate of the forward error in any column
!  of AINV exceeds EXCESS_ERR and return a nonzero value of INFO.
!  A   :  matrix to be inverted
!  N   :  size of A
!  AINV := set to the inverse of A unless INFO is nonzero.
!  RCOND := reciprocal of the estimated condition number of A.
!  INFO  := set to 0 unless the inverse cannot be computed with acceptable
!     relative accuracy.
!
    integer,intent(in)::n
    real(D),intent(inout)::a(n,n)
    real(D),intent(out)::ainv(n,n),rcond
    integer,intent(out)::info
!
!  Variables for DGESVX
!
    real(D),parameter::EXCESS_ERR=1.0e-04_D
    character::equed
    real(D)::af(n,n),rowscale(n),colscale(n),berr(n),ferr(n)
    real(D)::rvec(n*4),rhs(n,n)
    integer::iwork(n),ipiv(n)
!
!  We use the expert driver to solve AX=I for X; it computes estimates of the
!  forward and backward error in the solution.
!
    call identity(rhs,n)
    call dgesvx('E','N',n,n,a,n,af,n,ipiv,equed,rowscale, &
       colscale,rhs,n,ainv,n,rcond,ferr,berr,rvec,iwork,info)
    if(info.lt.0) call fabort('invert: dgesvx: illegal argument')
    if(any(ferr.gt.EXCESS_ERR)) info=1
    return
    end subroutine invert
!---------------------------------------------------------------------------
    subroutine get_pb_eigen(xback,patchsz,nsol,lambda,q,info)
!  GET_PB_EIGEN - compute the eigenvector spectrum of the background
!  covariance matrix, XBACK*transpose(XBACK) using the singular value
!  decomposition.  The SVD is chosen because it is expected to be numerically
!  stable even if XBACK is ill-conditioned.  (Note: XBACK has at least
!  one zero singular value in principle, as its rank is at most NSOL-1.)
!  My XBACK is the transpose of Ed's X^b.  The required eigenvector matrix Q
!  of X^b consists of the left singular vectors of transpose(XBACK).
!  Arguments:
!  XBACK : the Kth row holds the Kth ensemble state vector, K<=NSOL,
!     with the ensemble mean subtracted off.
!  PATCHSZ := the number of dynamical variables in the patch.
!  LAMBDA := on successful return, holds the largest NSOL-1 eigenvalues
!     in ascending order by index.
!  Q := on successful return, Kth column holds the eigenvector
!     corresponding to LAMBDA(K).
!  INFO := Set to 0 unless an error occurs.
!
    integer,intent(in)::patchsz,nsol
    real(D),intent(in)::xback(nsol,patchsz)
    real(D),intent(out)::lambda(nsol)
    real(D),intent(out)::q(patchsz,nsol)
    integer,intent(out)::info
!
!  Local storage
!
    real(D),allocatable::work(:)
    integer::lwork
    real(D)::extra(1)
!
!  We assume that XBACK has already been normalized by one less than
!  the ensemble size, Eq. (5) of the Tellus paper.  We compute the left
!  singular vectors XBACK and place them in Q; the right singular vectors
!  are not computed.
!
    q=transpose(xback)
    call dgesvd('O','N',patchsz,nsol,q,patchsz,lambda,extra,patchsz, &
       extra,nsol,extra,-1,info)  ! workspace query
    lwork=nint(extra(1))
    allocate(work(lwork))
    call dgesvd('O','N',patchsz,nsol,q,patchsz,lambda,extra,patchsz, &
       extra,nsol,work,lwork,info)
    if(info.lt.0) then
       call fabort('get_pb_eigen: illegal argument to DGESVD')
    else if(info.eq.0) then
       lambda=lambda**2  ! eigenvals of P^b are squares of singular values
    endif
    return
    end subroutine get_pb_eigen
!---------------------------------------------------------------------------
    pure subroutine identity(a,n)
!  IDENTITY - return the NxN identity matrix.
!
    integer,intent(in)::n
    real(D),intent(out)::a(n,n)
!
    integer::j
!
    a=0.0
    do j=1,n
       a(j,j)=1.0
    enddo
    return
    end subroutine identity
!---------------------------------------------------------------------------
    subroutine singular_vals(c,sing,n)
!  SINGULAR_VALS - find the singular values of the matrix C.  If the SVD
!  fails, then set the singular values to 0.
!
    integer,intent(in)::n
    real(D),intent(inout)::c(n,n)
    real(D),intent(out)::sing(n)
!
    real(D)::copy(n,n),wsize(1)
    integer::lwork,info
    real(D),allocatable::work(:)
!
    call dgesvd('N','N',n,n,copy,n,sing,copy,n,copy,n,wsize,-1,info)
    lwork=nint(wsize(1))
    allocate(work(lwork))
    copy=c
    call dgesvd('N','N',n,n,copy,n,sing,copy,n,copy,n,work,lwork,info)
    if(info.lt.0) then
       call fabort('singular_vals: argument error to DGESVD')
    else if(info.gt.0) then  ! no convergence
       sing=0.0
    endif
    return
    end subroutine singular_vals
!---------------------------------------------------------------------------

    end module assimilation
