#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
#define IF_ERROR_RETURN(action) if (status> 0) then; TRACEBACK; action; return; end if

#include "tm5.inc"

!-----------------------------------------------------------------------
!
!       PM module to calculate particulate matter (PM1, PM2.5, PM10, ...)
!       For station output
!       Implemented by Joost Aan de Brugh
!
!-----------------------------------------------------------------------

module PM ! {{{

  use GO,              only : gol, goErr, goPr

  ! You can specifiy which PM instances you want in your station output.
  ! For instance: PM2.5 and PM10
  ! You will get SO4, BC, POM, SS, DU, NO3, NH4 and water
  ! Where NO3 and NH4 come from EQSAM, and are considered to be in the accumulation mode.

  implicit none
  private

  Public :: PM_Init
  Public :: PM_Done
  Public :: PM_Step
  Public :: PM_Write

  integer                                  :: nsizelimits
  Real,    Dimension(:),       Allocatable :: sizelimits              ! Length: Number of size limits (~2)
  Integer, Dimension(:),       Allocatable :: stations_count          ! Length: Number of stations (~20)
  Real,    Dimension(:,:,:),   Allocatable :: stations_pm             ! Length: Size limits, PM tracers, stations (~2,8,~20)
  Real,    Dimension(:,:,:,:), Allocatable :: stations_pm_tobewritten ! Length: Size limits, PM tracers, stations, one. This is for the HDF, using start=(/0,0,0,io_record) (~2,8,~20,1)
  Real,    Dimension(:,:),     Allocatable :: modfrac                 ! Length: Number of M7 modes, number of size limits (7,~2)

  Integer                                  :: io_record
  Integer                                  :: io_nrecords 

  character(len=*), parameter  ::  mname = 'PM'
  Integer, Parameter           :: npmtracers = 8

  Real, Parameter :: hr2 = 0.5*sqrt(2.0)

! }}}
contains
  
  Subroutine PM_Init(pmsizelimits,status) ! {{{

    Use user_output_station, only : io_hdf,sds_station_pm,n_stat,n_stat2 ! To put in the PM attributes.
    Use mo_aero_m7,          only : nmod
    use file_hdf,            only : WriteAttribute, Init, Done, SetDim
    Use Dims,                Only : itaue,itaui,ndyn_max
    Use Partools,            only : myid,root

    ! --- in/out --------------------------------
    
    Character (len=200)     ::  pmsizelimits
    Integer, Intent(out)           :: status

    ! --- var   ------------------------------

    Integer                        :: ipos,i,cnt,idx,wlidx
    Character(len=200)             :: checkstring
    Logical                        :: last
    Integer                        :: j

    ! --- const ------------------------------

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

    io_nrecords = abs(itaue-itaui)/ndyn_max ! Like in the station file
    ipos = 1
    last = .False.
    cnt = 0
    Do While (.Not. last)
      i = index(pmsizelimits(ipos:),' ')
      If (i == 0) Then
        last = .True. ! This last sizelimit has to be counted.
      Else 
        checkstring = pmsizelimits(ipos:ipos+i-1)
        if(trim(checkstring) == '') Exit ! If you have trailing spaces, you do not want to count this one.
        ipos = ipos + i 
      End If
      cnt = cnt + 1
    End Do
    ! Okay, we counted the number of sizelimits, not very nice, but I do not know how it can be done nice, like 
    nsizelimits = cnt
    ! Allocate all stuff. The stations_count could have been allocated beforehand, but this looks more logical, having all Allocate-statements in one batch.
    Allocate(sizelimits(cnt))
    Allocate(stations_count(n_stat+n_stat2))
    Allocate(stations_pm(nsizelimits,npmtracers,n_stat+n_stat2))
    Allocate(stations_pm_tobewritten(nsizelimits,npmtracers,n_stat+n_stat2,1))
    Allocate(modfrac(nmod,nsizelimits))
    stations_count = 0
    stations_pm = 0.0

    ! Fill up the sizelimits array. I will be surprised if it works.
    ipos = 1
    Do idx=1,cnt
      i = index(pmsizelimits(ipos:),' ')
      If (i == 0) Then
        Read(pmsizelimits(ipos:),*) sizelimits(idx)
      Else 
        Read(pmsizelimits(ipos:ipos+i-1),*) sizelimits(idx)
        ipos = ipos + i 
      End If
    End Do

    If (myid .eq. root) Then
      io_record = 0
      call Init(sds_station_pm, io_hdf, 'PM', (/nsizelimits,npmtracers,n_stat+n_stat2,io_nrecords/), 'real(8)', status)
      IF_NOTOK_RETURN(status=1)
      call SetDim(sds_station_pm, 0, 'nsizelimits', 'Size limit', (/ (j, j=1,nsizelimits) /) , status)
      IF_NOTOK_RETURN(status=1)
      call SetDim(sds_station_pm, 1, 'npmtracers', 'PM tracer number', (/ (j, j=1,npmtracers) /) , status)
      IF_NOTOK_RETURN(status=1)
      call SetDim(sds_station_pm, 2, 'n_stations', 'station number', (/ (j, j=1,n_stat+n_stat2) /) , status)
      IF_NOTOK_RETURN(status=1)
      call SetDim(sds_station_pm, 3, 'n_records', 'time slot #', (/ (j, j=1,io_nrecords) /) , status)
      IF_NOTOK_RETURN(status=1)
      call WriteAttribute(sds_station_pm, 'Unit', 'ug/m3, with size limits in um', status)
      IF_NOTOK_RETURN(status=1)
      Call WriteAttribute(sds_station_pm, 'Size limits', sizelimits,status)
      IF_NOTOK_RETURN(status=1)
      Call WriteAttribute(sds_station_pm, 'Tracers', 'SO4, BC, POM, SS, DU, NO3, NH4, H2O',status)
      IF_NOTOK_RETURN(status=1)

    End If
    ! Now the micrometers are written down, we can convert micrometers diameter to meters radius. One micrometer diameter is 7.e-6 meter radius.
    sizelimits = sizelimits * 5.e-7

    ! ok
    status = 0

  End Subroutine PM_Init
  ! ================================================================ }}}
  subroutine PM_Done(status ) ! {{{

    Use user_output_station, only : sds_station_pm

    Use Partools, only : myid,root

    Use file_hdf, only : Done

    Implicit None

    ! --- in/out --------------------------------

    integer, intent(out)           ::  status

    ! --- var   ------------------------------

    ! --- const ------------------------------

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

    ! --- begin --------------------------------
      
    ! Deallocate a hell of a lot of disastrous things.
    Deallocate(sizelimits)
    Deallocate(stations_count)
    Deallocate(stations_pm)
    Deallocate(stations_pm_tobewritten)
    Deallocate(modfrac)

    If (myid .eq. root) Then
      call Done(sds_station_pm, status)
      IF_NOTOK_RETURN(status=1)
    End If

    ! ok
    status = 0

  end subroutine PM_Done ! }}}
  
  Subroutine PM_Step(region,status) ! {{{

    Use user_output_station, only : stat,n_stat,n_stat2
    Use tracer_data,         only : mass_dat
    ! I need to convert (/kg air) to (/m3 air)
    Use chem_param,          only : iso4nus, iso4ais, ibcais, ipomais, iso4acs, ibcacs, ipomacs, issacs, iduacs, iso4cos, ibccos, &
         ipomcos, isscos, iducos, ibcaii, ipomaii, iduaci, iducoi, ino3_a, inh4
    Use mo_aero_m7,          only : nmod, cmedr2mmedr, sigma
    Use m7_data,             only : h2o_mode, rw_mode

    ! To rewrite ug / cell to ug/m3
    Use global_data,         only : region_dat
    Use Meteodata,           only : gph_dat
    Use Dims,                only : im,jm,lm

    ! --- in/out --------------------------------

    Integer, Intent(In)   :: region
    Integer, Intent(out)  :: status

    ! --- var   ------------------------------

    Integer :: isl,imod,i_stat
    Real :: z,rad

    ! For converting mass to mass/volume
    Real :: volume

    ! --- const ------------------------------

    Integer, Parameter :: ipm_so4 = 1
    Integer, Parameter :: ipm_bc  = 2
    Integer, Parameter :: ipm_pom = 3
    Integer, Parameter :: ipm_ss  = 4
    Integer, Parameter :: ipm_du  = 5
    Integer, Parameter :: ipm_no3 = 6
    Integer, Parameter :: ipm_nh4 = 7
    Integer, Parameter :: ipm_h2o = 8

    Real, Dimension(nmod) :: lnsigma = log(sigma)

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

    ! Fraction belonging to the PM class for each mode, modfrac
    Do i_stat = 1,n_stat+n_stat2
       If (stat(i_stat)%region /= region) Cycle ! This method is evaluated per region, so stations not in the right region can go cycling.
       stations_count(i_stat) = stations_count(i_stat) + 1
       Do imod=1,nmod
          ! Calculate the median radius of the volume weighted distribution.

          rad = rw_mode(region,imod)%d3(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls)*cmedr2mmedr(imod)

          If (rad .eq. 0.0) Then
             modfrac(imod,:) = 1.0 ! Should not matter, because if the radius is zero, there are no aerosols. But in principle, it would mean that all aerosols are infinitely small, so they would fit in any PM-class.
          Else
             Do isl=1,nsizelimits
                ! Uses the normal distribution formula. We know the Z-value of the size limit.
                z = (log(sizelimits(isl)) - log(rad)) / lnsigma(imod)
                modfrac(imod,isl) = 0.5 + 0.5 * ERF(z*hr2)
             End Do
          End If
       End Do
       
       ! And the volume. This is the nicest and cleanest line of code of this whole module.
       volume = (gph_dat(region)%data(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls+1) - &
                 gph_dat(region)%data(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls  )   ) * region_dat(region)%dxyp(stat(i_stat)%js)

       ! First each processor gathers their own contributions, then we allreduce all output and write it down.
       ! Without MPI, we strip the If-statement and always resolve the 'tracer' way.

       ! Parallel over tracers: we need the rm_t, which is over i,j, l from 1 to lm and t from 1 to ntracetloc. We check whether the tracer is active on this processor.
       stations_pm(:,ipm_so4,i_stat) = stations_pm(:,ipm_so4,i_stat) + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,iso4nus-offsetn) * modfrac(1,:) / volume 
       stations_pm(:,ipm_so4,i_stat) = stations_pm(:,ipm_so4,i_stat) + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,iso4ais-offsetn) * modfrac(2,:) / volume 
       stations_pm(:,ipm_so4,i_stat) = stations_pm(:,ipm_so4,i_stat) + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,iso4acs-offsetn) * modfrac(3,:) / volume 
       stations_pm(:,ipm_so4,i_stat) = stations_pm(:,ipm_so4,i_stat) + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,iso4cos-offsetn) * modfrac(4,:) / volume 
       stations_pm(:,ipm_bc,i_stat)  = stations_pm(:,ipm_bc,i_stat)  + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,ibcais-offsetn)  * modfrac(2,:) / volume 
       stations_pm(:,ipm_bc,i_stat)  = stations_pm(:,ipm_bc,i_stat)  + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,ibcacs-offsetn)  * modfrac(3,:) / volume 
       stations_pm(:,ipm_bc,i_stat)  = stations_pm(:,ipm_bc,i_stat)  + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,ibccos-offsetn)  * modfrac(4,:) / volume 
       stations_pm(:,ipm_bc,i_stat)  = stations_pm(:,ipm_bc,i_stat)  + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,ibcaii-offsetn)  * modfrac(5,:) / volume 
       stations_pm(:,ipm_pom,i_stat) = stations_pm(:,ipm_pom,i_stat) + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,ipomais-offsetn) * modfrac(2,:) / volume 
       stations_pm(:,ipm_pom,i_stat) = stations_pm(:,ipm_pom,i_stat) + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,ipomacs-offsetn) * modfrac(3,:) / volume 
       stations_pm(:,ipm_pom,i_stat) = stations_pm(:,ipm_pom,i_stat) + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,ipomcos-offsetn) * modfrac(4,:) / volume 
       stations_pm(:,ipm_pom,i_stat) = stations_pm(:,ipm_pom,i_stat) + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,ipomaii-offsetn) * modfrac(5,:) / volume 
       stations_pm(:,ipm_ss,i_stat)  = stations_pm(:,ipm_ss,i_stat)  + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,issacs-offsetn)  * modfrac(3,:) / volume 
       stations_pm(:,ipm_ss,i_stat)  = stations_pm(:,ipm_ss,i_stat)  + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,isscos-offsetn)  * modfrac(4,:) / volume 
       stations_pm(:,ipm_du,i_stat)  = stations_pm(:,ipm_du,i_stat)  + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,iduacs-offsetn)  * modfrac(3,:) / volume 
       stations_pm(:,ipm_du,i_stat)  = stations_pm(:,ipm_du,i_stat)  + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,iducos-offsetn)  * modfrac(4,:) / volume 
       stations_pm(:,ipm_du,i_stat)  = stations_pm(:,ipm_du,i_stat)  + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,iduaci-offsetn)  * modfrac(6,:) / volume 
       stations_pm(:,ipm_du,i_stat)  = stations_pm(:,ipm_du,i_stat)  + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,iducoi-offsetn)  * modfrac(7,:) / volume 
       stations_pm(:,ipm_no3,i_stat) = stations_pm(:,ipm_no3,i_stat) + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,ino3_a-offsetn)  * modfrac(3,:) / volume 
       stations_pm(:,ipm_nh4,i_stat) = stations_pm(:,ipm_nh4,i_stat) + mass_dat(region)%rm_t(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls,inh4-offsetn)    * modfrac(3,:) / volume 
       stations_pm(:,ipm_h2o,i_stat) = stations_pm(:,ipm_h2o,i_stat) + h2o_mode(region,1)%d3(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls)                 * modfrac(1,:) / volume 
       stations_pm(:,ipm_h2o,i_stat) = stations_pm(:,ipm_h2o,i_stat) + h2o_mode(region,2)%d3(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls)                 * modfrac(2,:) / volume 
       stations_pm(:,ipm_h2o,i_stat) = stations_pm(:,ipm_h2o,i_stat) + h2o_mode(region,3)%d3(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls)                 * modfrac(3,:) / volume 
       stations_pm(:,ipm_h2o,i_stat) = stations_pm(:,ipm_h2o,i_stat) + h2o_mode(region,4)%d3(stat(i_stat)%is,stat(i_stat)%js,stat(i_stat)%ls)                 * modfrac(4,:) / volume 
    End Do

    ! ok
    status = 0

  End Subroutine PM_Step
  ! ================================================================ }}}
  Subroutine PM_Write(status) ! {{{

    use file_hdf,            only : WriteData
    Use user_output_station, only : sds_station_pm,stat,n_stat,n_stat2
    Use Partools,            only : myid, root, tracer_active,offsetn,offsetl,lmloc

    ! --- in/out --------------------------------
    
    Integer, Intent(out)  :: status

    ! --- var   ------------------------------

    Integer :: i_stat ! Just because I do not have Rebin. I want to divide a 3d object by a 1d object. It is possible to repeat dividing a 2d object by a scalar.

    ! --- const ------------------------------

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

    stations_pm(:,:,:) = 1.e9*stations_pm(:,:,:)

    Do i_stat=1,n_stat+n_stat2
      stations_pm(:,:,i_stat) = stations_pm(:,:,i_stat) / stations_count(i_stat) 
    End Do

    stations_pm_tobewritten(:,:,:,1) = stations_pm(:,:,:)

    ! Now the root continues. The rest only resets its values.
    If (myid .eq. root) then
      Call Writedata(sds_station_pm,stations_pm_tobewritten,status,start=(/0,0,0,io_record/))
      IF_NOTOK_RETURN(status=1)

      io_record = io_record + 1
    End If ! myid .eq. root

    ! Reset all output fields.
    stations_count = 0.0
    stations_pm = 0.0

    status = 0

  End Subroutine PM_Write
  ! ================================================================ }}}
end module PM