program cov_window
!
!---   Calculate error covariance (analysis or forecast) ----------
!
!   INPUT: sqrtPx  = square-root error covariance (columns only)
!
!  OUTPUT: Px(i1:i2,j1:j2)   = sub-matrix (window, block) of full error covariance matrix
!
!--------------------------------------------------------------------

#ifdef MPI_USE
  include "mpif.h"
#endif

   integer,parameter::xdim=20              ! dimensions (x1d_cntrl)
   integer,parameter::isqrtPx=21           ! sqrtPx file #
   integer,parameter::iPx=51               ! Px file #

!-----
         real,dimension(:),allocatable :: sqrtPx
         real,dimension(:,:),allocatable :: covPx,covPx_loc
         real,dimension(:),allocatable :: sum_loc,sumx
         integer,allocatable :: jlen(:),jdisp(:)

     integer :: N_LOC,jsta,jend
     integer :: NPROC,MPIRANK
     integer :: NENS
     integer :: NENS_START
     integer :: NALL
     integer :: idim
     character*11 Pxname
     character*7 Px

     NAMELIST /ENSEMBLE_SIZE/ NENS_START,NENS
     NAMELIST /WINDOW/ i_s,i_e,j_s,j_e
!====================================================================

#ifdef MPI_USE
!
!     START MPI
!
      CALL MPI_INIT(IERR)
      IF (IERR .NE. MPI_SUCCESS) STOP 'FAILED TO INIT MPI'
      CALL MPI_COMM_SIZE(MPI_COMM_WORLD, NPROC, IERR)
      CALL MPI_COMM_RANK(MPI_COMM_WORLD, MPIRANK, IERR)
      CALL MPI_BARRIER(MPI_COMM_WORLD,IERR)
#else
      write(*,*) "This is a NO MPI run"
      MPIRANK=0
      NPROC=1
#endif

       if(MPIRANK.eq.0) then
       write(*,*) "start cov_window"
       end if

!-- read ensemble size

       REWIND 15
       READ(15,ENSEMBLE_SIZE)

          if(MPIRANK.eq.0) then
          write(*,*) "cov_window: NENS_START, NENS=",NENS_START,NENS
          end if

!- Total number of degrees of freedom
          NALL=NENS-NENS_START+1

!-------- Read window parameters -----------
       REWIND 16
       READ(16,WINDOW)

          if(MPIRANK.eq.0) then
          write(*,*) "get_cov: Covariance window: I_start,I_end=",i_s,i_e
          write(*,*) "get_cov: Covariance window: J_start,J_end=",j_s,j_e
          end if

!====================================================================
!-- read square-root error covariance (column) 

    call getenv("sqrtPx",Px)
    if(MPIRANK.eq.0) then
    write(*,*) "sqrtPx=",Px
    end if

!====================================================================
!==
!==  define local dimensions and indexes
!==
       allocate(jdisp(0:NPROC-1)) ; allocate(jlen(0:NPROC-1))

            do irank=0,NPROC-1
      CALL para_range(NENS_START,NENS,NPROC,irank,jsta,jend)
            jlen(irank)=jend-jsta+1
            jdisp(irank)=jsta-NENS_START
            end do
#ifdef MPI_USE
      CALL MPI_BARRIER(MPI_COMM_WORLD,IERR)
#endif
      CALL para_range(NENS_START,NENS,NPROC,MPIRANK,jsta,jend)

!==
         N_LOC=jlen(MPIRANK)
!==
!
         if(MPIRANK.eq.0) then
           do nn=0,NPROC-1
           write(*,*) "nn=",nn," jdisp=",jdisp(nn)," jlen=",jlen(nn)
           end do
          end if

!=============== dimensions =================
      CLOSE(xdim)
      OPEN(xdim,file='x1d_cntrl',status='unknown',form='unformatted')
      READ(xdim) idim
      CLOSE(xdim,status='KEEP')
       if(MPIRANK.eq.0) then
       write(*,*) "get_cov: idim=",idim
       end if

 2000 format(a7,i4.4)

         if(MPIRANK.eq.0) then
         write(*,*) "idim=",idim
         end if
!====================================================================
!-- get the output string 


      allocate(sqrtPx(1:idim))
      allocate(covPx_loc(i_s:i_e,j_s:j_e))
      covPx_loc(:,:)=0.0


      DO N=jsta,jend

      WRITE(Pxname,2000)  Px,N
      write(*,*) N," Pxname=",Pxname
      CLOSE(isqrtPx)
      OPEN(isqrtPx,file=Pxname,status='unknown',form='unformatted')
      READ(isqrtPx) NN
      READ(isqrtPx) sqrtPx
      CLOSE(isqrtPx,status='KEEP')

      write(*,*) N," read sqrtPx min,max=",minval(sqrtPx),maxval(sqrtPx)

          do j=j_s,j_e
         do i=i_s,i_e
        covPx_loc(i,j)=covPx_loc(i,j)+sqrtPx(i)*sqrtPx(j)
         end do
          end do

      END DO

      write(*,*) MPIRANK," covPx_loc min,max=",minval(covPx_loc),maxval(covPx_loc)

!--------------------------------------------
       ldim=(i_e-i_s+1)*(j_e-j_s+1)
       allocate(sum_loc(1:ldim))
       allocate(sumx(1:ldim))
       
        kk=0
          do j=j_s,j_e
         do i=i_s,i_e
        kk=kk+1
        sum_loc(kk)=covPx_loc(i,j)
         end do
          end do

#ifdef MPI_USE
             CALL MPI_REDUCE &
   (sum_loc,sumx,ldim,MPI_REAL4,MPI_SUM,0,MPI_COMM_WORLD,IERR)
#else
    sumx=sum_loc
#endif


      if(MPIRANK.eq.0) then
      allocate(covPx(i_s:i_e,j_s:j_e))
        kk=0
          do j=j_s,j_e
         do i=i_s,i_e
        kk=kk+1
        covPx(i,j)=sumx(kk)
         end do
          end do

!== choose format (for GRADS ???)
      CLOSE(iPx)
      OPEN(iPx,file='covPx',status='unknown',form='unformatted')
      WRITE(iPx) covPx
      CLOSE(iPx,status='KEEP')

      do i=i_s,i_e
      write(*,*) i," covPx=",covPx(i,i_s)
      end do
      deallocate(covPx)
      end if

      deallocate(sum_loc) ; deallocate(sumx)
      deallocate(sqrtPx)  ; deallocate(covPx_loc)

       if(MPIRANK.eq.0) then
       write(*,*) "end cov_window"
       end if

!====================================================================

#ifdef MPI_USE
      CALL MPI_FINALIZE(IERR)
#endif


stop
end program cov_window