program enkf_eco
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Program which uses the ensemble Kalman Filter to assimilate   !
! in the 1D version of the ecosystem model.                     !
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

   use mod_dimensions
   use mod_states
   use mod_observations
   use mod_variances
   use m_cholesky
   use m_physpar
   use m_mkmxlayer
   use m_reference
   use m_dump
   use m_load
   use m_tecfld
   use m_pseudo
   use m_load_ens
   use m_dump_ens
   use m_measure
   use m_integrate
   use m_ens_stat
   use m_random
   use m_analysis
   use m_set_random_seed2
   implicit none
   type(states), allocatable ::  mem(:)  ! Contains the ensemble of model states

   type(states) ref(ndim)     ! Holds the full reference solution      
   type(states) bguess        ! Holds the initial values (best guess solution)
                              ! used when generating the init. ensemble
   type(states) estimate(ndim)! The solution estimate
   type(states) variance(ndim)! The variance estimate

   type(states) forw          ! Temporary model state
   type(local_states) bnd0,bnd0w    ! Boundary conditions
   type(local_states) bndH,bndHw    ! Boundary conditions
   real workA(kdim,antvar)
   real workB(antvar)
   real cov(kdim,kdim), chov(kdim,kdim) ! Cov matrix and its Sqrt
   type(states) ave           ! Temporary model state
   type(states) var           ! Temporary model state
   type(states) Mean          ! Ensemble mean
   real, allocatable :: cov1(:,:), chov1(:,:) 
   type(variances) vars
   
   real, allocatable :: S(:,:) ! S(3*kdim,nrsamp)   

   real MM(ndim)             ! Mixed layer depth

   type(observations) d(nrobs,nrmes_t) , d_temp
   type(lstates) meas_action
   type(local_states) mse

   real, allocatable :: U0(:,:),sig0(:),VT0(:,:)
   real, allocatable :: U(:,:), sig(:),VT1(:,:),A0(:,:)
   real, allocatable, dimension(:)   :: work
   real scale           ! 1/nrsamp   
! ME
   type(local_states) mseres
! ME end
   type(local_states) totrms
   type(local_states) totsd

   real,parameter :: range = 50.      ! 50 m decorrelation length (in z)
   real tfin                          ! final integration time
   logical iniref    ! genereate new ref sol or not
   logical iniens
   real dt,dz
   real totdepth
   integer mode, i_xp
   integer, parameter :: n_xp = 10 ! repeat 10 times experiment

   integer i,j,k,m,n,iens,na,nb,nstep,nrsamp
   
   

   logical ex

   character(len=1) tag1
   character(len=3) tag3
   real t0,t1

   integer ierr,lwork
  
 

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Reading the parameters from the file "infile"
   call physpar(tfin,dt,totdepth,dz,iniref,iniens,vars,mode,meas_action,nrsamp)
   print *,'physpar done!'

   allocate (mem(nrsamp))
   allocate (cov1(nrsamp,nrsamp))
   allocate (chov1(nrsamp,nrsamp))
   allocate (S(3*kdim,nrsamp))

   call mkmxlayer(MM,dt)
   print *,'mxlayer done!'

   call set_random_seed2

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Sqrt spatial covariance matrix
  print *, 'initial cov ...'
  do i=1,kdim
     do j=1,i
        cov(i,j)=gaussian(float(abs(i-j))*dz/range)
        cov(j,i)=cov(i,j)
     end do
  end do
                                                                                                                    
  print *, 'beginning cholesky(cov) ...'
  chov=cov
  call cholesky(chov,kdim) 

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

!! We need this matrix to be of dimension (nrsamp,nrsamp) to be able to use it flater when computing othogonal matrix.
 


! Sqrt spatial covariance matrix
  print *, 'initial cov ...'
  do i=1,nrsamp
     do j=1,i
        cov1(i,j)=gaussian(float(abs(i-j))*dz/range)
        cov1(j,i)=cov1(i,j)
     end do
  end do
                                                                                                                    
  print *, 'beginning cholesky(cov1) ...'
  chov1=cov1
  call cholesky(chov1,nrsamp)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
do i_xp=1, n_xp

! Set initial and boundary conditions
! Boundary cond (no flux at surface)
   bnd0=0.0

! Boundary cond (specified value at bottom)
   bndH%N=10.0
   bndH%P=0.0
   bndH%H=0.0

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Reference solution
   print *,'iniref ',iniref
   inquire(file='refsol.uf',exist=ex)
   if ((iniref).or.(.not.ex)) then 
      print *,'Computing reference solution'
      call reference(ref,bguess,MM,dz,totdepth,dt,tfin,vars,bnd0,bndH,chov)
      call dump('refsol.uf',ref)
   else
      print *,'Loading reference solution from file'
      call load('refsol.uf',ref)
      !print *, ' Init Ref N(kdim) ' , ref(1)%N(kdim)
   endif
   print *,'done!'

   totsd=0. ; totrms=0.


!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Initial ensemble
   print *,'iniens= ',iniens
   inquire(file='iniens.uf',exist=ex)
   if ((iniens).or.(.not.ex)) then
      print *,'Computing initial ensemble'
      do iens=1,nrsamp
         call pseudo(worka,kdim,antvar,chov)

         mem(iens)%N=max(0.0, bguess%N+sqrt(vars%ini%N)*workA(:,1) )
         mem(iens)%P=max(0.0, bguess%P+sqrt(vars%ini%P)*workA(:,2) )
         mem(iens)%H=max(0.0, bguess%H+sqrt(vars%ini%H)*workA(:,3) )
      enddo
      call dump_ens('iniens.uf',mem,nrsamp)
   else
      print *,'Loading initial ensemble from file'
      call load_ens('iniens.uf',mem,nrsamp)
   endif
   print *,'done!'
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

!NEW PART WITH GENERATION OF START ENSEMBLE OF "NRSAMP_INI" MEMBERS AND THEN RESAMPLE "NRSAMP" MEMBERS .	

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!Compute ensemble mean
!   scale = 1.0/float(nrsamp_ini)
!   Mean=0.0
 !  var=0.0
  ! do i=1,nrsamp_ini
   !   Mean = Mean + mem(i)
   !      var=var+mem(i)*mem(i)
   ! enddo

 !  Mean=scale*Mean
  !       var=var*scale
   


! Substract mean from ensemble members

    
  
 !   do i=1,nrsamp_ini
  !    mem(i)%N=mem(i)%N-Mean%N
   !   mem(i)%P=mem(i)%P-Mean%P
    !  mem(i)%H=mem(i)%H-Mean%H

    
    
  !    var%N = var%N + mem(i)%N**2    
   !   var%P = var%P + mem(i)%P**2
    !  var%H = var%H + mem(i)%H**2
   
   
    
   !  mem(i)%N=var%N*mem(i)%N
    ! mem(i)%P=var%P*mem(i)%P
     !mem(i)%H=var%H*mem(i)%H

!enddo


!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

!Make an othogonal VT1 used as linear combination for final ensemble
   
 !   allocate (A0(nrsamp,nrsamp), U(nrsamp,nrsamp) , sig(nrsamp) , VT1(nrsamp,nrsamp))
  !  lwork=2*max(3*nrsamp+3*kdim,5*nrsamp)
   ! allocate(work(lwork))
    !call pseudo(A0,nrsamp,nrsamp,chov1)
    !call dgesvd('N', 'S', nrsamp, nrsamp, A0, nrsamp, sig, U, nrsamp, VT1, nrsamp, work, lwork, ierr)


 !   if(ierr/=0) print *, 'ierr',ierr  
  !  deallocate(A0, sig, U, work)




! Compute SVD of oversized ensemble
 
 ! allocate (U0(3*kdim,nrsamp_ini))
  !allocate (sig0(nrsamp_ini))
  !allocate(VT0(1,1))
  !lwork=2*max(3*nrsamp_ini+3*kdim,5*nrsamp_ini)
  !allocate(work(lwork))
  !sig0=0.0 
 !call dgesvd('S', 'N', 3*kdim, nrsamp_ini, mem, 3*kdim,sig0, U0, 3*kdim, VT0, nrsamp_ini , work, lwork, ierr)
  !deallocate(work,VT0)

 !  if (ierr /= 0) then
  !   print *,'analysis: ierr from call dgesvd 0= ',ierr; stop
   !endif


 !  do i=1,nrsamp_ini
  !   mem(i)%N=mem(i)%N/var%N
   !  mem(i)%P=mem(i)%P/var%P
    ! mem(i)%H=mem(i)%H/var%H
  
  !   enddo


! Generate first nrsamp members and add mean

  

 ! do i=1,nrsamp
  ! do j=1,nrsamp


  !   mem(i)%N=(U0(1:kdim,i)*sig0(i)/sqrt(float(nrsamp))*VT1(j,i))+Mean%N
   !  mem(i)%P=(U0(kdim+1:2*kdim,i)*sig0(i)/sqrt(float(nrsamp))*VT1(j,i))+Mean%P
   !  mem(i)%H=(U0(2*kdim+1:3*kdim,i)*sig0(i)/sqrt(float(nrsamp))*VT1(j,i))+Mean%H
      
    
 

 !enddo
!enddo


!deallocate(U0,sig0,VT1)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
        
!END OF NEW PART

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Generate measurements by measuring the reference solution
   call measure(d,ref,vars,meas_action)
   

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! Time stepping
! -compute statistics of initial ensemble
   call ens_stat(mem,ave,var,nrsamp)
   estimate(1)=ave
   variance(1)=var

   do m=1,nrmes_t+1
      na=(m-1)*nres+1
      nb=m*nres+1
      print *,'na,nb,nres,nrmes_t ',na,nb,nres,nrmes_t
      do nstep=na+1,nb
         t0=dt*float(nstep-2)
         t1=dt*float(nstep-1)
         print '(a,2I5,2f8.2)','m,nstep,t0,t1:',m,nstep,t0,t1
         do iens=1,nrsamp 
            bnd0w=bnd0
            call random(workb,3)
            bndHw%N=bndH%N+sqrt(vars%bH%N)*workb(1)
            bndHw%P=bndH%P+sqrt(vars%bH%P)*workb(2)
            bndHw%H=bndH%H+sqrt(vars%bH%H)*workb(3)
            forw=mem(iens)


            call integrate(forw,MM,bnd0w,bndHw,totdepth,t0,t1,dt)

            call pseudo(worka,kdim,antvar,chov)

            mem(iens)%N=max(0.0, forw%N+sqrt(dt)*sqrt(vars%dyn%N)*worka(:,1))
            mem(iens)%P=max(0.0, forw%P+sqrt(dt)*sqrt(vars%dyn%P)*worka(:,2))
            mem(iens)%H=max(0.0, forw%H+sqrt(dt)*sqrt(vars%dyn%H)*worka(:,3))
         enddo
         call ens_stat(mem,ave,var,nrsamp)
         estimate(nstep)=ave
         variance(nstep)=var
      enddo
      if (nstep < ndim) then
         if (mode == 1) then
            print *,'calling analysis',m
            print *,'Meas. at t1 ',t1
            call prepare_and_analysis(mem,d(:,m),nrsamp)
         endif
      endif
   enddo

   write(tag1,'(i1.1)')mode
   write(tag3,'(i3.3)')nrsamp
   call tecfld('output_'//tag1,ref,estimate,variance,dz,dt)

   open(10,file='mse_'//tag3//'.dat')
   do i=1,ndim  
      mse%N=sum(variance(i)%N(:))/float(kdim)
      mse%P=sum(variance(i)%P(:))/float(kdim)
      mse%H=sum(variance(i)%H(:))/float(kdim)
      write(10,'(f12.4,4f12.4)')float(i-1)*dt,sqrt(mse%N), sqrt(mse%P), &
                             sqrt(mse%H), sqrt(mse%N+mse%P+mse%H)
      totsd%N = totsd%N + mse%N
      totsd%P = totsd%P + mse%P
      totsd%H = totsd%H + mse%H
   enddo
   close(10)
   
! ME for the residuals:
   open(10,file='mseres_'//tag3//'.dat')
   do i=1,ndim
     mseres%N=sum((estimate(i)%N(:)-ref(i)%N(:))**2)/float(kdim)
     mseres%P=sum((estimate(i)%P(:)-ref(i)%P(:))**2)/float(kdim)
     mseres%H=sum((estimate(i)%H(:)-ref(i)%H(:))**2)/float(kdim)
     write(10,'(f12.4,4f12.4)')float(i-1)*dt,sqrt(mseres%N), &
              sqrt(mseres%P), sqrt(mseres%H), &
              sqrt(mseres%N+mseres%P+mseres%H)
      totrms%N = totrms%N + mseres%N
      totrms%P = totrms%P + mseres%P
      totrms%H = totrms%H + mseres%H
   enddo
   close(10)
! ME res. end.

enddo

totsd%N = sqrt(totsd%N / float(ndim*n_xp) )
totsd%P = sqrt(totsd%P / float(ndim*n_xp) )
totsd%H = sqrt(totsd%H / float(ndim*n_xp) )

totrms%N = sqrt(totrms%N / float(ndim*n_xp) )
totrms%P = sqrt(totrms%P / float(ndim*n_xp) )
totrms%H = sqrt(totrms%H / float(ndim*n_xp) )

open(11, file='totsd.dat',position='append')
write(11,'(i5.3,4g15.5)') nrsamp, totsd%N, totsd%P, totsd%H, totsd%N+totsd%P+totsd%H
close(11)

open(12, file='totrms.dat',position='append')
write(12,'(i5.3,4g15.5)') nrsamp, totrms%N, totrms%P, totrms%H, totrms%N+totrms%P+totrms%H
close(12)

deallocate (mem)
deallocate (cov1)
deallocate (chov1)
deallocate (S)

end program enkf_eco