! First include the set of model-wide compiler flags 
#include "tm5.inc"

module oi_module

  ! use dims
  ! use MatInvertLib,  only : MatInvert_lapack

  ! uses two LAPACK routines:
  !  call dgetrf( nmsi, nmsi, bplo, nmsi, ipiv, ierror )
  !  call dgetrs( 'N', nmsi, 1, bplo, nmsi, ipiv, vincr, nmsi, ierror )


  implicit none

  private

  public :: AssimilateColumn 
 
  ! public :: qanratio_latest
  ! real,dimension(im,jm,lm) :: qanratio_latest

  ! Do not allow analysis ratio's below this limit (excluding also negatives)
  real,parameter :: QAnRatioLimit=0.5

  ! correlation lookup tables
  ! blut : the global B correlation matrix
  ! rho: the distance-dependent correlation
  ! rho_conc: the fractional concentration-difference correlation
  real, dimension(:,:,:), allocatable :: blut
  real,dimension(0:2100)       :: rho 
  real,dimension(0:2000)       :: rho_conc 
  
  ! lower limit for forecast subcolumn in layer (10^15 molec cm^-2) 
  ! to prevent arithmetic exceptions
  real,parameter               :: subcolumnLimit = 1.0e-7

  ! flag for first call to "AssimilateColumn"
  logical                      :: firstcall = .true.
  
contains


  subroutine AssimilateColumn ( im, jm, lm, fcerr, obsInfoReduced, qfc, qanratio )
    !=======================================================================
    ! 
    ! AssimilateColumn - OI assimilation
    !
    !   Uses the optimal interpolation formula with a homogemeous
    !   horizontal model covariance. 
    !
    !   im, jm, lm = 3D grid dimensions
    !   fcerr      = 2D field of forecast column errors (1e15 molec/cm2)
    !   obsInfoReduced = contains observations, 
    !                    forecast and observation operators
    !   qfc        = 3D forecast field of partial columns of NO2
    !
    !   qanratio   = (output) 3D field of ratios between the 
    !                analysis and forecast
    !
    !              Henk Eskes and Folkert Boersma, September 5, 2003.
    !   
    !
    !=======================================================================
    use MOmFSuper,	only : TObsInfoReduced 
    implicit none

    ! in/out
    ! clustered superobservations
    integer, intent(in)                  :: im, jm, lm
    type(TObsInfoReduced),intent(in)     :: obsInfoReduced
    real,dimension(im,jm),intent(in)     :: fcerr
    real,dimension(im,jm,lm),intent(in)  :: qfc
    real,dimension(im,jm,lm),intent(out) :: qanratio		

    ! correlation: shape and length
    !      'g': Gaussian shape 
    !      'e': exponential shape
    !      't': Thiebaux 2nd order autoregressive shape
    real,parameter         :: corlen = 600.0   ! correlation length in km (eg 650.0)
    character(1),parameter :: corshape = 't'

    ! correlation length in concentration space 
    !   % concentration difference for 1/e correlation decrease 
    real,parameter         :: conc_corlen = 30.0  ! 20% concentration diff 

    ! local arrays
    real, dimension(:,:,:), allocatable   :: qcorrl, htvincr
    integer, dimension(:,:), allocatable  :: iobsmask
    real, dimension(:), allocatable       :: odifms, smms, soms, vincr
    integer, dimension(:), allocatable    :: ilonms, jlatms
    real, dimension(:,:), allocatable     :: Hmat, qfcmat

    ! local
    integer                        :: i, j, l, ipix, nms, nms1, imax, il
    integer                        :: nSmallAnRatio, nLargeAnRatio, nWrongQFC

    !
    ! fill the model correlation lookup table (only once!)
    !
    if( firstcall )then
       call Inioilut ( im, jm, corshape, corlen, conc_corlen )
       write(*,733) corshape,corlen,conc_corlen
733    format(' assimilateColumn: lookup table filled, ',&
              'shape = ',a1,', corlen = ',f7.1,/,&
              '    concentration correlation dC = ',f7.1,' %' )
       firstcall = .false.
    endif

    ! number of super-observations
    nms = obsInfoReduced%count
    nms1 = max(1,nms)

    ! Allocate the local arrays
    allocate (  qcorrl(im,jm,lm)  )
    allocate (  htvincr(im,jm,lm) )
    allocate (  iobsmask(im,jm)   )
    allocate (  odifms(nms1)      )
    allocate (  smms(nms1)        )
    allocate (  soms(nms1)        )
    allocate (  vincr(nms1)       )
    allocate (  ilonms(nms1)      )
    allocate (  jlatms(nms1)      )
    allocate (  Hmat(lm,nms1)     )
    allocate (  qfcmat(lm,nms1)   )
    !
    !
    ! fill measurement space data arrays - mismatch, coordinates and sigma's
    !
    iobsmask(:,:) = 0 ! mask for cells containing observations
    do ipix = 1, nms
       odifms(ipix) = obsInfoReduced%slcobs(ipix) - obsInfoReduced%slcfc(ipix)
       i = obsInfoReduced%obsi(ipix)  ! measurement longitude
       j = obsInfoReduced%obsj(ipix)  ! measurement latitude
       ilonms(ipix) = i 
       jlatms(ipix) = j 
       iobsmask(i,j) = 1
       smms(ipix)   = fcerr(i,j)      ! model uncertainty at observation
       soms(ipix)   = obsInfoReduced%sigobs(ipix)   ! measurement uncertainty
       Hmat(1:lm,ipix) = obsInfoReduced%hobs(1:lm,ipix) ! observation operator          
       qfcmat(1:lm,ipix) = qfc(i,j,1:lm)
    end do
    !
    nSmallAnRatio = 0
    nLargeAnRatio = 0
    nWrongQFC     = 0

    if ( nms .gt. 0 ) then     ! there are measurements
       !
       ! solve "vincr", the solution of
       !   ( H B H^T + O ) vincr = Q_obs - H Q_mod
       !   H: model-measurement interpolation matrix
       !   B: model covariance matrix
       !   O: observation covariance
       !   Q: total column
       !   ( Q_obs - H Q_mod = "odif" )
       !
       call Solvincr ( im, lm, nms, odifms, ilonms, jlatms, smms, soms, Hmat, qfcmat, vincr)

       ! interpolate vincr to the model grid
       call HTy ( im, jm, lm, nms, vincr, ilonms, jlatms, Hmat, htvincr )

       ! multiply with the covariance matrix on the model grid to obtain the
       ! OI field correction
       call Bx ( im, jm, lm, htvincr, iobsmask, fcerr, corlen, qfc, qcorrl)

       !
       ! determine analysed field
       !    q_an = q_mod + qcorrl
       !    qcorrl = B H^T ( H B H^T + O )^-1 ( Q_obs - H Q_mod )
       !
       ! analysed correction ratio = (qfc+qcorrl)/qfc
       ! check for small and negative ratio's
       !

       do l = 1, lm
          do j = 1, jm
             ! circumvent unphysical polar cells
            imax = im
            !if( (j == 1) .or. (j == jm) ) imax = 1 	!JDMP
            do i = 1, imax
              ! If qfc becomes very small, set qanratio equal 1.0 
              ! to prevent Arithmetic Exceptions (Folkert Boersma, September 3, 2003)
              if ( qfc(i,j,l)  <= subcolumnLimit ) then
                qanratio(i,j,l) = 1.0
                nWrongQFC = nWrongQFC + 1
                ! print*,'ERROR, small qfc value',i,j,l,qfc(i,j,l),qfc(i,j,l)+qcorrl(i,j,l)
              else                               
                qanratio(i,j,l) = (qfc(i,j,l)+qcorrl(i,j,l)) / qfc(i,j,l)
              end if
       
              ! do not allow small and negative ratio's
              if ( qanratio(i,j,l) < QAnRatioLimit ) then
                ! set ratio equal to specified limit
                qanratio(i,j,l) = QAnRatioLimit
                nSmallAnRatio = nSmallAnRatio + 1
              end if

              ! avoid extreme relative analysis increments
              if ( qanratio(i,j,l) > 5.0 ) then
                ! set ratio equal to specified limit
                qanratio(i,j,l) = 5.0
                nLargeAnRatio = nLargeAnRatio + 1
              end if

            end do
          end do
       end do

       !do i=2,im             ! polar caps						    !JDMP
       !  qanratio(i,1,1:lm)  = qanratio(1,1,1:lm) 				    !JDMP
       !  qanratio(i,jm,1:lm) = qanratio(1,jm,1:lm) !JDMP
       !end do 												    !JDMP
       !     
    else    ! nms=0
       ! no measurements, analysis = forecast
       qanratio(:,:,:) = 1.0
    endif
    !
    if ( nSmallAnRatio > 0 ) write(*,'(a,i5,a)') 'assimilateColumn: WARNING ',nSmallAnRatio, &
            ' AnRatio values < 0.5 occurred and corrected'
    if ( nLargeAnRatio > 0 ) write(*,'(a,i5,a)') 'assimilateColumn: WARNING ',nLargeAnRatio, &
            ' AnRatio values > 5 occurred and corrected'
    if ( nWrongQFC > 0 ) write(*,'(a,i5,a)') 'assimilateColumn: WARNING ',nWrongQFC, &
            ' forecast values qfc with very small, zero, or negative qfc occurred'


    ! De-allocate the local arrays
    deallocate (  qcorrl   )
    deallocate (  htvincr  )
    deallocate (  iobsmask )
    deallocate (  odifms   )
    deallocate (  smms     )
    deallocate (  soms     )
    deallocate (  vincr    )
    deallocate (  ilonms   )
    deallocate (  jlatms   )
    deallocate (  Hmat     )
    deallocate (  qfcmat   )


  end subroutine AssimilateColumn



  subroutine Inioilut ( im, jm, choice, corlen, conc_corlen )
    !=======================================================================
    !
    ! OI (Optimal Interpolation) Utility routine:
    ! Fill the model correlation lookup tables
    !      "rho" - distance-dependent correlation function
    !      "rho_conc" - concentration-dependent correlation function
    !      "blut" - correlation lookup table for two grid points on the globe
    !           blut(idif,j1,j2)
    !           idif: longit. index distance between points (0..im/2)
    !           j1,j2: latitude indices of two points 
    !
    ! choice [in] - single character describing correlation function: 
    !    'e' : exponential distance dependence, with linear tail
    !    'g' : Gaussian distance dependence
    !    't' : Thiebaux autoregressive correlation
    ! corlen [in] - correlation length in km (1/e length)
    ! conc_corlen [in] - concentration correlation length in % difference
    !
    ! The array dimensions "im" and "jm" are taken from the file "constants"
    ! The lookup tables B and rho are permanent: see above
    !
    !=======================================================================
    use misctools, 	only : dist

    implicit none

    ! in/out     
    integer, intent(in)     :: im, jm
    character(1),intent(in) :: choice
    real,intent(in)         :: corlen,conc_corlen

    ! local
    integer      :: i,idif,j1,j2,icorlen,imax
    real         :: long1,long2,lati1,lati2,dst,b,ccl

    ! begin code

    ! allocate lookup table
    allocate ( blut( 0:im/2, jm, jm ) )

    ! Fill lookup table for distance dependent function - 10 km step
    !
    if ( corlen .lt. 0.000 ) stop 'ERROR Inioilut: negative correlation length'
    !if( corlen .lt. 100.0 ) stop 'WARNING Inioilut: correlation length < 100 km'
    !
    rho(0) = 1.0
    rho(1:2100) = 0.0
    icorlen = nint(corlen/10.0)
    if (icorlen == 0) stop 'ERROR in ass_oi_subr.f90; icorlen =0.'
    if ( 5*icorlen .gt. 2100 ) stop 'ERROR Inioilut: correlation length too large'
    !
    if ( choice .eq. 'e' ) then  ! exponential correlation
       if ( icorlen .gt. 0 ) then
          do i = 0, 4*icorlen-1
             rho(i) = exp(-real(i)/real(icorlen))
          end do
          do i = 4*icorlen, 5*icorlen-1
             rho(i) = exp(-4.0)*(1.0-real(i-4*icorlen)/real(icorlen))
          end do
       end if
    else
       if( choice .eq. 'g' )then  ! Gaussian correlation
          do i = 0, 5*icorlen-1
             rho(i) = exp(-(real(i)/real(icorlen))**2)
          end do
       else
          if( choice .eq. 't' )then  ! Thiebaux autoregressive correlation
             do i = 0, 5*icorlen-1
                rho(i) = (1.0+2.0*real(i)/real(icorlen)) * exp( -2.0*real(i)/real(icorlen) )
             end do
          else
             stop 'ERROR Inioilut: "choice" invalid'
          end if
       end if
    end if
    !
    ! Correlation between concentrations c1 and c2:
    ! correlation = rho_conc( nint(abs(1000*(c1-c2)/(c1+c2))) )
    !
    if ( conc_corlen > 50.0 ) stop 'IniOILut: ERROR conc_corlen > 50%, inconsistent with functional form'
    ccl = 5.0*conc_corlen  ! 1/e reduction for "corlen_conc" % difference
    imax = 900             ! correlations = 0 when abs(c1-c2)/(c1+c2) > 0.9
    do i = 0, imax
       rho_conc(i) = exp( - (real(i)/ccl)**2 )
    end do
    rho_conc(imax:2000) = 0.0
    rho_conc(0) = 1.0
    !
    ! Fill covariance matrix table for all distinct distances on the sphere
    !
    ! calculate lat,lon of the two points
    long1 = 0.0
    do idif = 0, im/2
       long2 = real(idif)*360.0/real(im)
       do j1 = 1, jm
          lati1 = -90.0 + (real(j1)-0.5)*180.0/real(jm)
          if ( j1 .eq. 1  ) lati1 = -90.0			
          if ( j1 .eq. jm ) lati1 =  90.0		
          do j2 = j1, jm
             lati2 = -90.0 + (real(j2)-0.5)*180.0/real(jm)
             if ( j2 .eq. 1  ) lati2 = -90.0			
             if ( j2 .eq. jm ) lati2 =  90.0	
             ! determine distance between points (dst in km)
             call dist(long1,lati1,long2,lati2,dst)
             ! use lookup table for rho
             if ( nint(dst/10.0) > 2100 ) stop 'Inioilut: coding ERROR, distance'
             b = rho(nint(dst/10.0))
             blut(idif,j1,j2) = b
             blut(idif,j2,j1) = b
             if ( (b<0.0) .or. (b>1.0) ) stop 'Inioilut: coding ERROR, b'
          end do
       end do
    end do
    !
    ! polar caps exception: polar dummy cells are uncorrelated
    !
    !do idif=1,im/2				!JDMP
    !   blut(idif,1,1)   = 0.0					!JDMP
    !   blut(idif,jm,jm) = 0.0	!JDMP
    !   blut(idif,1,jm)  = 0.0			!JDMP
    !   blut(idif,jm,1)  = 0.0			!JDMP
    !enddo
    !
  end subroutine Inioilut


  subroutine Solvincr &
       ( im, lm, nmsi, odifms, ilonmsi, jlatmsi, smmsi, somsi, Hmat, qfcmat, vincr )
    !=======================================================================
    !
    ! OI (Optimal Interpolation) Utility routine:
    !   Solves the vector "vincr", the solution of
    !      ( B + O ) vincr = Q_obs - H Q_mod
    !      B: model covariance matrix, defined at measurement positions
    !      O: observation covariance
    !      Q: observed quantity
    !   all variables are defined in the space of the observations
    !
    !   im     = number of model longitudes (in)
    !   lm     = number of model levels (in)
    !   nmsi   = number of measurements (in)
    !   odifms = Q_obs - H Q_mod, the set of departures  (in)
    !   ilonmsi,jlatmsi = grid coordinates of the measurements  (in)
    !   smmsi,somsi =  model and obs uncertainties (sigma standard dev's)  (in)
    !   Hmat   = observation operator matrix
    !   qfcmat = 3D model forecast profiles at observations
    !   vincr  = see above (out)
    !
    ! The subroutine usue the lookup tables "blut" and "rho"
    ! The array dimensions "im" and "lm" are taken from the module "dimension"
    !
    !=======================================================================
    implicit none

    ! in/out
    integer,intent(in)                    :: im, lm
    integer,intent(in)                    :: nmsi
    real,dimension(nmsi),intent(in)       :: odifms, smmsi, somsi
    real,dimension(lm,nmsi),intent(in)    :: Hmat, qfcmat
    integer,dimension(nmsi),intent(in)    :: ilonmsi, jlatmsi
    real,dimension(nmsi),intent(out)      :: vincr

    ! local
    integer   :: l, m, n, ierror, itriag, i1mini2
    real      :: test, hht, correlsum, c1, c2, ccorrel
    !
    real, dimension(:,:), allocatable  :: bplo       ! B+O matrix

    integer, parameter :: obsMaxSuper = 50000
    ! Lapack storage
    integer, dimension(obsMaxSuper) :: ipiv

    ! begin code

    ! allocate bplo matrix
    allocate ( bplo(nmsi,nmsi) )

    !
    ! Fill model + observation covariance matrix ( H B H^T + O )
    ! B is assumed diagonal in altitude, and altitude independent
    !
    do m = 1, nmsi
       do n = m, nmsi
          i1mini2 = abs(ilonmsi(m) - ilonmsi(n))
          if ( i1mini2 .gt. im ) stop 'ERROR solvincr: ilonmsi array incorrect'
          if ( i1mini2 .gt. im/2 ) i1mini2 = im - i1mini2		
          hht = 0.0
          correlsum = 0.0 ! sum of vertical correlation factors
          do l = 1, lm
             ! Add max statement to avoid division by 0, Folkert Boersma
             c1 = max ( qfcmat(l,m), subcolumnLimit )
             c2 = max ( qfcmat(l,n), subcolumnLimit )
             ccorrel = rho_conc( nint(abs(1000.0*(c1-c2)/(c1+c2))) )
             hht = hht + Hmat(l,m)*Hmat(l,n)*ccorrel*(c1+c2)
             correlsum = correlsum + (c1+c2)
          end do
          if ( correlsum == 0. ) stop 'ERROR in ass_oi_subr.f90; correlsum = 0.'
          hht = hht / correlsum
          bplo(m,n) = smmsi(m)*smmsi(n)*hht*blut(i1mini2,jlatmsi(m),jlatmsi(n))
          if ( n .eq. m ) then
             bplo(m,n) = bplo(m,n) + somsi(m)*somsi(m)
          else
             bplo(n,m) = bplo(m,n)
          end if
       end do
    end do
    !
    ! Solve  ( B + O ) v_incr = odifms
    ! GAUSS - Solves Ax = vec where A is a square matrix
    !
    vincr(1:nmsi) = odifms(1:nmsi)

    ! test   = 1.e-2
    ! itriag = 0
    ! call Gauss ( bplo, vincr, nmsi, nmsi, test, ierror, itriag )
    !
    ! using lapack
    !
    call dgetrf( nmsi, nmsi, bplo, nmsi, ipiv, ierror )
    if ( ierror /= 0 ) then
       print *,' LAPACK sgetrf error = ',ierror
       stop
    end if
    call dgetrs( 'N', nmsi, 1, bplo, nmsi, ipiv, vincr, nmsi, ierror )
    if ( ierror /= 0 ) then
       print *,' LAPACK sgetrf error = ',ierror
       stop
    end if
    !
    if( ierror .ne. 0 ) then
       write(*,'(a,i5)') 'ERROR solvincr, gauss: ierror = ',ierror
       stop 
    end if
    
    ! deallocate bplo matrix
    deallocate ( bplo )

  end subroutine Solvincr


  subroutine HTy ( im, jm, lm, nms, v, ilonmsi, jlatmsi, Hmat, htv )
    !=======================================================================
    !
    ! OI (Optimal Interpolation) Utility routine:
    !   Interpolates an observation-space input vector "v" to the model grid
    !   (multiplies "v" with H^T, the transpose of the observation operator)
    !
    !  im, jm, lm = model array dimensions  (in)
    !  nms    = number of measurements  (in)
    !  v      = vector in observation space  (in)
    !  ilonmsi, jlatmsi  = the model grid positions of the observations  (in)
    !  Hmat   = observation matrix (in)
    !  htv    = resulting model-space vector  (out)
    !
    !=======================================================================
    implicit none

    ! in/out
    integer, intent(in)                     :: im, jm, lm
    integer, intent(in)                     :: nms
    real, dimension(nms), intent(in)        :: v
    integer, dimension(nms), intent(in)     :: ilonmsi, jlatmsi
    real, dimension(lm,nms), intent(in)     :: Hmat
    real, dimension(im,jm,lm), intent(out)  :: htv

    ! local
    integer :: i, j, ims

    ! initialise htv
    htv(:,:,:) = 0.0

    ! multiply "v" with Hmat
    do ims = 1, nms
       i = ilonmsi(ims)
       j = jlatmsi(ims)
       htv(i,j,1:lm) = Hmat(1:lm,ims)*v(ims)
    end do
    !
  end subroutine HTy


  subroutine Bx ( im, jm, lm, htvincr, iobsmask, smod, corlen, qfc, qcorrl )
    !=======================================================================
    !
    ! OI (Optimal Interpolation) Utility routine:
    ! Operates with the model covariance matrix B on the input vector "htvincr"
    !
    ! im,jm,lm: dimension of the lattice and vectors (in)
    ! htvincr: input vector (from OI analysis)  (in)
    ! iobsmask: equals 1 when grid box contains observations (in)
    ! smod: model column error standard deviation field  (in)
    ! corlen: model error correlation length in km  (in)
    ! qfc: the 3D forecast field (in)
    ! qcorrl: product of B with htvincr - the 3D correction field (out)
    !
    !=======================================================================
    implicit none

    ! in/out
    integer, intent(in)                    :: im, jm, lm
    real, dimension(im,jm,lm), intent(in)  :: htvincr,qfc
    real, dimension(im,jm), intent(in)     :: smod
    real, intent(in)                       :: corlen
    integer, dimension(im,jm), intent(in)  :: iobsmask
    real, dimension(im,jm,lm), intent(out) :: qcorrl

    ! local
    real, parameter       :: rearth = 6378.5
    real, parameter       :: pi = 3.1415927
    !
    real, dimension(:), allocatable    :: fac
    real, dimension(:,:), allocatable  :: qfccol    
    !
    integer               :: i, j, l, jdel, jmax, jmin, idel, i2, j2
    integer               :: imini2, i2index, idel1, idel2
    real                  :: lat, fac2, qfcnorm, c1, c2, ccorrel, tcorrel

    ! lookup tables for model covariance, blut,rho (see above)    

    ! begin code

    ! allocate helper arrays
    allocate ( fac(lm) )     
    allocate ( qfccol(im,jm) )    ! forecast columns of NO2 

    ! initialize output field to zero
    qcorrl(:,:,:) = 0.0
    !
    ! calculate model forecast columns
    do i = 1, im
       do j = 1, jm
          qfccol(i,j) = 0.0
          do l = 1, lm
            qfccol(i,j) = qfccol(i,j) + qfc(i,j,l)
          end do
       end do
    end do
    !
    ! distribute OI increments with the model covariance matrix.
    ! (i,j)   coordinate of measurement gridbox
    ! (i2,j2) coordinate of points in neighbourhood of (i,j)
    !
    jdel = 1 + int(5.0*corlen*real(jm)/(pi*rearth))
    do j = 1, jm
       jmax = min(j+jdel,jm)
       jmin = max(j-jdel,1)
       lat = -0.5*pi+(real(j)-0.5)*pi/real(jm)
       idel = 1+int(5.0*corlen*real(im)/(2.0*pi*rearth*cos(lat)))
       idel1 = min(idel,im/2-1)
       idel2 = min(idel,im/2)
       do i = 1, im
          if ( iobsmask(i,j) == 1 ) then    ! there is an observation
             do l = 1, lm
                fac(l) = htvincr(i,j,l)*smod(i,j)
             end do
             do i2index = i-idel1, i+idel2
                i2 = mod(i2index+im-1,im) + 1    ! array index el 1..im
                imini2 = abs(i2index-i)
                do j2 = jmin, jmax
                   fac2 = smod(i2,j2)*blut(imini2,j,j2)
                   qfcnorm = 1.0/(qfccol(i,j)+qfccol(i2,j2))
                   do l = 1, lm
                      ! Add max statement to aviod division by 0, Folkert Boersma
                      c1 = max ( qfc(i,j,l), subcolumnLimit )
                      c2 = max ( qfc(i2,j2,l), subcolumnLimit )
                      ccorrel = rho_conc( nint(abs(1000.0*(c1-c2)/(c1+c2))) )
                      tcorrel = fac(l)*fac2*ccorrel*(c1+c2)*qfcnorm
                      qcorrl(i2,j2,l) = qcorrl(i2,j2,l) + tcorrel
                   end do
                end do
             end do
          end if
       end do
    end do

    ! deallocate helper arrays
    deallocate ( fac )     
    deallocate ( qfccol ) 

  end subroutine Bx


end module oi_module
