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

! Macro
#define IF_ERROR_RETURN(action) if (status> 0) then; TRACEBACK; action; return; end if
#define TRACEBACK write (gol,'("in ",a," (",a,i6,")")') rname, __FILE__, __LINE__ ; call goErr
#define IF_NOTOK_RETURN(action) if (status/=0) then; TRACEBACK; action; return; end if

! integer, parameter :: sp = selected_real_kind(6, 37)
! integer, parameter :: dp = selected_real_kind(15, 307)

!=====================================================================
! 
! Central NO2 column data assimilation module
!
! Interface call by TM5 (in assim_interface_mod.f90):
!
!         call AssimNO2( ntimestep )
!
! The subroutine AssimNO2 performs a data assimilation time step.
!
!                                        Henk Eskes, KNMI, January 2003
!                                           Version 1.0 - Februari 2003
!
!                                      J.D. Maasakkers, KNMI, July 2013
!                                              Prototype v3 - July 2013
!
!                       Henk Eskes, Mark ter Linden, KNMI, Mar-Sep 2015
!                               Interface to TM5-mp-CB05 and to TROPOMI
!======================================================================

module MAssimilation

  ! printing, error handling
  use GO,               only : gol, goPr, goErr
  ! NO2 AMF lookup table
  use MTAmf,            only : TAmfLut

  implicit none

  private

  public :: AssimNO2

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

  ! integer to count the number of calls to the routine "AssimNO2" 
  integer :: callcount  = 0
  integer :: assimcount = 0

  ! module name
  character(len=*), parameter  :: mname = 'MAssimilation'

  ! keep filenames in memory
  character(256)               :: filename_S5P_L2_NO2_filelist
  ! character(256)               :: terrain_filename
  character(256)               :: amf_filename
  logical                      :: divideByAMFGeo
  ! output directory for OmF netcdf files
  character(256)               :: outputdir

  ! NO2 AMF lookup table (keep in memory)
  type(TAmfLut)                :: amfLut

  ! If true, call the retrieval for both forecast and analysis fields
  logical  ::  computeRetrievalForAnalysis 

  ! Switch: optional to write OmF information to separate netcdf file
  logical  ::  doWriteOmF

  ! Debug: extra retrieval output
  logical  ::  debugMode

contains


  subroutine AssimNO2( rcFile, ntimestep, TM5Data, nObs, noxscaling, status )
    !======================================================================
    !
    !   Assimilation of NO2 slant column data
    !   Main routine from where all assimilation routines are called
    !
    !   rcFile     : name of the rc (settings) file
    !   ntimestep  : timestep in seconds between calls to AssimNO2
    !   TM5Data    : All TM5 fields and grid info from the interface module
    !   nObs       : nr. of observations assimilated in this timestep
    !   noxscaling : an/fc ratio for the TM5 NOx transported field
    !			                   Henk Eskes, KNMI, 2000-2002-2015
    !======================================================================

    ! --- TM5 interfaces ---

    ! TM5 model arrays and variables:
    use MTM5Data,               only : TTM5Data 
    ! TM5 date and time routines:
    use dims,			only : itau
    ! use datetime,      		only : tau2date
    use GO,                     only : TrcFile, Init, Done, ReadRc

    ! --- TROPOMI L2 interfaces ---

    ! Observational input (slant columns) is stored in this structure
    use MTObsTrack,             only : TObsTrack, ObsTrackDeallocate
    ! Storage for retrieval results (vertical columns)
    use MTObsFcInfo,            only : TObsFcInfo, ObsFcInfoAllocate, ObsFcInfoDeallocate
    ! module to read OMI NO2 slant column data (by tracks)
    use MObsRead,               only : ObsFileGet
    ! module to write OMI NO2 vertical column retrieval data (by tracks)
    use TM5IF_module,           only : TM5IF_obsfcinfo_write

    ! --- DOMINO interfaces ---

    ! compute model forecast and observation operator
    use MObservationOperator,   only : InitOMFHx, GetOMFHx
    ! construct superobservations
    use MOmFSuper,              only : GetOMFHxSuper, TObsInfoReduced, ObsInfoReduced_deallocate
    ! Kalman Filter routines
    use oi_module,              only : AssimilateColumn
    ! Terrain height module
    ! use terrain,                only : terrain_read
    ! Error module, use the model vertical column forecast error estimate 
    use MNO2RetrievalError,     only : forecastError
    ! NO2 AMF lookup table
    use MAmfRead,               only : ReadAmfLut
    ! Write OmF data to file 
    use MOmF_output,            only : OmF_Output
    ! Cloud retrieval correction
    ! use cloudcor,		only : cloudcor_read
    
    implicit none

    ! --- in-out ---

    character(*), intent(in)         ::  rcFile
    integer, intent(in)              ::  ntimestep
    type(TTM5Data), intent(inout)    ::  TM5Data 
    integer, intent(out)             ::  nObs
    real, dimension(TM5Data%im, TM5Data%jm, TM5Data%lm), intent(inout)  ::  noxscaling
    integer, intent(out)             ::  status

    ! ---  local variables ---

    character(len=*), parameter ::  rname = mname//'/AssimNO2'

    type(TrcFile)    ::  rcF

    ! the ratio between analysis and forecast (3D field)
    real, dimension(:,:,:), allocatable :: AnFcRatio
    ! array of 3D partial NO2 columns per gridbox 
    real, dimension(:,:,:), allocatable :: no2pcf
    ! array of 2D NO2 columns per lat-lon gridbox 
    real, dimension(:,:), allocatable   :: no2col_fc, no2col_an

    ! the 2D model column forecast error distribution 
    real, dimension(:,:), allocatable   :: NO2VCForecastError 

    ! ---------------------------------------------------------	
    ! Storage of track data: NO2 and cloud info
    type(TObsTrack)       :: no2Tr
    ! observations and model forecast
    type(TObsFcInfo)      :: obsFcInfo, obsAnInfo
    ! clustered superobservations
    type(TObsInfoReduced) :: obsFcInfoReduced, obsAnInfoReduced

    ! local
    integer                          :: i, l
    real                             :: slc,slcsig,tropfraction
    logical                          :: obsFound, isForecast
    integer,parameter                :: omf_tape=83552

    ! OMIFilelist !JDM
    ! integer            :: unit, ierror, system
    ! character(len=160) :: listfile, command

    ! --- begin code ---

    callcount = callcount + 1

    write(*,1111) callcount
1111   format(' ------- assimno2 ',i4,'  -------')

    ! -------------- initialise -------------------

    if ( callcount == 1 ) then

       ! open rcfile:
       call Init( rcF, rcfile, status )
       IF_NOTOK_RETURN(status=1)

       ! Read the location of the file which contains a list of L2 NO2 files
       call ReadRc( rcF, 'domino.inputlist', filename_S5P_L2_NO2_filelist, status, default='list' )
       IF_ERROR_RETURN(status=1)

       ! Read the location of the orography file
       ! call ReadRc( rcF, 'domino.terrain_filename', terrain_filename, status, default='dem.nc' ) 

       ! Read the location of the AMF lookup table file
       call ReadRc( rcF, 'domino.amf_file', amf_filename, status, default='no2_amf_lut.nc' )

       ! Divide the AMF by AMFgeo after reading 
       call ReadRc( rcF, 'domino.amf.divide.by.amfgeo', divideByAMFGeo, status, default=.false. )

       ! Optional to call observation operator after the assimilation step 
       call ReadRc( rcF, 'domino.compute.retrieval.for.analysis', computeRetrievalForAnalysis, status, default=.false. )

       ! Write OmF statistics to netcdf file
       call ReadRc( rcF, 'domino.write.OmF', doWriteOmF, status, default=.true. )

       ! Debug mode yes/no
       call ReadRc( rcF, 'domino.debug.mode', debugMode, status, default=.false. )

       ! Output directory (for OmF netcdf output)
       call ReadRc( rcF, 'output.dir', outputdir, status, default='.' )

       ! close rcfile:
       call Done( rcF, status )
       IF_NOTOK_RETURN(status=1)

       ! Fill error specification matrices
       write(*,112) ForecastError
112    format('  InitForecastError: Initialising forecast error to ',f6.2,' 1e15 molec/cm2')

       ! Read settings for the observation operator
       call InitOMFHx( rcFile )
       
       ! Here we read the 90Arc terrain height database to be able to 
       ! calculate the average terrain height of each pixel. !JDMTH
       ! call terrain_read ( terrain_filename )

       ! Read the air-mass factor lookup table from file
       print *,' Reading NO2 AMF from file : ',trim(amf_filename)
       print *,'          divide by amfgeo : ',divideByAmfGeo
       call ReadAmfLut ( amf_filename, divideByAmfGeo, amfLut, status )

       print *,' S5P list file : ',trim(filename_S5P_L2_NO2_filelist)

    end if

    ! 2D forecast error field
    allocate ( NO2VCForecastError(TM5Data%im,TM5Data%jm) )
    NO2VCForecastError(:,:) = ForecastError

    ! ---------- Allocate storage and read the track matching this model time ----------

    call ObsFileGet( filename_S5P_L2_NO2_filelist, itau-ntimestep/2,itau+ntimestep/2, no2Tr, obsFound )

    ! At this point a track of NO2 and cloud data, "no2Tr" has 
    ! been read or the track data record "no2Tr" is empty: obsFound = .false.

    if ( obsFound ) then

       write(*,'(a,i6)') 'assimNO2: No. of pixels in track ',no2Tr%count
       nObs = no2Tr%count

       ! allocate 3D partial column field and 2D column fields
       allocate ( no2pcf(TM5Data%im,TM5Data%jm,TM5Data%lm) )
       allocate ( no2col_fc(TM5Data%im,TM5Data%jm) )
       allocate ( no2col_an(TM5Data%im,TM5Data%jm) )
       
       ! Store analysis/forecast ratio, = 1 before the analysis
       noxscaling(:,:,:) = 1.0
       
       ! get 3D field of partial NO2 columns (forecast)
       call GetNO2PartialColumnField ( TM5Data, noxscaling, no2pcf, no2col_fc )
    
       ! allocate storage for retrieval output
       call ObsFcInfoAllocate ( obsFcInfo, no2Tr%count, TM5Data%lm )

       ! Calculate model forecast and observation operator 
       ! store information for individual observations in "ObsFcInfo"
       isForecast = .true.
       call GetOMFHx( amfLut, no2Tr, TM5Data, no2pcf, obsFcInfo, isForecast, outputdir, debugMode )
       
       ! Write the retrieval data to the level-2 file, based on TM5 forecast 
       call TM5IF_obsfcinfo_write( no2Tr, obsFcInfo, TM5Data )
       
       ! cluster observations into superobservations
       call GetOMFHxSuper( TM5Data%im, TM5Data%jm, TM5Data%lm, no2Tr, obsFcInfo, obsFcInfoReduced, 'f', computeRetrievalForAnalysis )

       write(*,1244) obsFcInfo%count,obsFcInfoReduced%count,TM5Data%date
1244   format(' assimno2: ',i5,' (',i4,') observ.,  Timestep ',i4,5i3)

       ! ----- Assimilation step -----

       assimcount = assimcount + 1
          
       ! allocate the helper array
       allocate ( anFcRatio(TM5Data%im,TM5Data%jm,TM5Data%lm) )
       
       print *, 'assimno2: assimilating observations'
       call AssimilateColumn ( TM5Data%im, TM5Data%jm, TM5Data%lm, NO2VCForecastError, obsFcInfoReduced, no2pcf, anFcRatio )
       
       ! Store analysis/forecast ratio
       noxscaling(:,:,:) = anFcRatio(:,:,:)
          
       ! remove helpers
       deallocate ( anFcRatio )

       ! ----- Calculate model analysis and observation operator -----

       if ( computeRetrievalForAnalysis ) then	  

          ! allocate storage for retrieval output
          call ObsFcInfoAllocate ( obsAnInfo, no2Tr%count, TM5Data%lm )
       
          ! get 3D field of partial NO2 columns: for analysis, noxscaling <> 1.0
          call GetNO2PartialColumnField ( TM5Data, noxscaling, no2pcf, no2col_an )

          ! store information for individual observations in "ObsFcInfo"
          isForecast = .false.
          call GetOMFHx ( amfLut, no2Tr, TM5Data, no2pcf, obsAnInfo, isForecast, outputdir, debugMode )

          ! cluster observations into superobservations         
          !   note: this routine allocates the storage for the superobs
          call GetOMFHxSuper ( TM5Data%im, TM5Data%jm, TM5Data%lm, no2Tr, obsAnInfo, obsAnInfoReduced, 'a', computeRetrievalForAnalysis )

       end if
          
       ! ----- Write innovation statistics to netcdf file -----
       if ( doWriteOmF ) then
          if ( computeRetrievalForAnalysis ) then
             call OmF_Output( no2Tr, obsFcInfo, obsFcInfoReduced, NO2VCForecastError, TM5Data%im, TM5Data%jm, TM5Data%lon, TM5Data%lat, no2col_fc, TM5Data%date, outputdir, status, obsAnInfo, obsAnInfoReduced, no2col_an )
          else          
             call OmF_Output( no2Tr, obsFcInfo, obsFcInfoReduced, NO2VCForecastError, TM5Data%im, TM5Data%jm, TM5Data%lon, TM5Data%lat, no2col_fc, TM5Data%date, outputdir, status )
          end if
       end if
       
       ! empty the orbit stack
       no2Tr%norbitparts = 0
       no2Tr%count       = 0

       ! remove 2D (im,jm) forecast error field
       deallocate ( NO2VCForecastError )

       ! remove storage TROPOMI L2 NO2 observations
       call ObsTrackDeallocate ( no2Tr )

       ! remove storage retrieval results
       call ObsFcInfoDeallocate ( obsFcInfo )
       if ( computeRetrievalForAnalysis ) then	  
          call ObsFcInfoDeallocate ( obsAnInfo )
       end if

       ! remove storage superobservations
       call ObsInfoReduced_deallocate ( obsFcInfoReduced )
       if ( computeRetrievalForAnalysis ) then	  
          call ObsInfoReduced_deallocate ( obsAnInfoReduced )
       end if

       ! storage for partial column field and column fields
       deallocate ( no2pcf )
       deallocate ( no2col_fc )
       deallocate ( no2col_an )
       
    else  ! no observations found

       ! write(*,'(a,t40)') '  assimno2:',TM5Data%date
       nObs = 0

    end if
    
    write(*,1112) 
1112 format(' ------- assimno2 done  -------')
    
  end subroutine AssimNO2
  




  subroutine GetNO2PartialColumnField( TM5Data, analysisScaling, no2pcf, no2col )
    !
    ! extract a 3D partial column field from the TM model field "no2"
    ! input  :  "TM5Data" structure with TM5 fields 
    !           "analysisScaling" (=1 for forecast, <>1 for analysis),
    !               scaling factors resulting from the analysis 
    ! output :  "no2pcf" in 10^15 molecules per cm^2 for each 3D grid box
    !           "no2col" total column in 10^15 molecules per cm^2 
    !
    ! Adjusted in DOMINO-3 because the interface passes mixing ratio's 
    ! while the original code is based on the mass and "rm" arrays
    !                                                     Henk Eskes, 2015 

    ! TM5 model arrays and variables:
    use MTM5Data,               only : TTM5Data 
    use physconstants,          only : pdiff2moleccm2

    implicit none

    ! -- in/out --
    type(TTM5Data), intent(in)                                   ::  TM5Data 
    real,dimension(TM5Data%im,TM5Data%jm,TM5Data%lm),intent(in)  ::  analysisScaling
    real,dimension(TM5Data%im,TM5Data%jm,TM5Data%lm),intent(out) ::  no2pcf
    real,dimension(TM5Data%im,TM5Data%jm),intent(out)            ::  no2col

    ! -- local --
    integer :: i, j, l
    real    :: pbot, ptop
    real,dimension(:), allocatable  :: no2col_zonal
    real    :: no2col_mean

    !
    ! Convert from volume mixing ratio to molecules cm^-2 
    ! using the pressure drop over a layer (in Pa)
    !
    allocate ( no2col_zonal(TM5Data%jm) )
    no2col_zonal(:) = 0.0
    no2col(:,:) = 0.0

    do i = 1, TM5Data%im
       do j = 1, TM5Data%jm
          do l = 1, TM5Data%lm
             pbot = TM5Data%hyai(l)   + TM5Data%hybi(l) * TM5Data%ps(i,j) 
             ptop = TM5Data%hyai(l+1) + TM5Data%hybi(l+1) * TM5Data%ps(i,j) 
             no2pcf(i,j,l) = pdiff2moleccm2 * ( pbot - ptop ) * TM5Data%no2(i,j,l) * analysisScaling(i,j,l)
             no2col_zonal(j) = no2col_zonal(j) + no2pcf(i,j,l)
             no2col(i,j) = no2col(i,j) + no2pcf(i,j,l)
             ! if ( i==100 .and. j==100 ) print*,'l, pbot, ptop, no2, scaling ',l,pbot, ptop, TM5Data%no2(i,j,l), analysisScaling(i,j,l)
          end do
       end do
    end do
    no2col_zonal(:) = no2col_zonal(:) / real(TM5Data%im)
    no2col_mean     = sum(no2col_zonal) / real(TM5Data%jm)

    if ( isnan(no2col_mean) ) then 
       print '(a)', "ERROR in GetNO2PartialColumnField: NaN found in the TM5 NO2 field"
       stop "stop because of NaN"
    end if

!    print *, 'no2col_zonal -90, -60, .. 90 = ',no2col_zonal(1),no2col_zonal(30),no2col_zonal(60),no2col_zonal(90),no2col_zonal(120),no2col_zonal(150),no2col_zonal(180)
!    print *, 'no2col at (59,7) and (34,38) ',sum(no2pcf(59,7,:)),sum(no2pcf(34,38,:))
!    print *, 'no2pcf profile at 5E, 53N ',no2pcf(185,143,:)

    print '(a,f10.3)', ' GetNO2PartialColumnField, mean NO2 column = ', no2col_mean

    deallocate ( no2col_zonal )

  end subroutine GetNO2PartialColumnField



end module MAssimilation
