!injection of particles
MODULE moduleInject
  USE moduleSpecies
  use velocityDistribution
  use moduleMesh, only: meshEdgePointer

  !Generic injection of particles
  TYPE, abstract:: injectGeneric
    INTEGER:: id
    CHARACTER(:), ALLOCATABLE:: name
    REAL(8):: vMod !Velocity (module)
    REAL(8):: temperature(1:3) !Temperature
    REAL(8):: n(1:3) !Direction of injection
    LOGICAL:: fixDirection !The injection of particles has a fix direction defined by n
    INTEGER:: nParticles !Number of particles to introduce each time step
    CLASS(speciesGeneric), POINTER:: species !Species of injection
    INTEGER:: nEdges
    type(meshEdgePointer), allocatable:: edges(:)
    INTEGER, ALLOCATABLE:: particlesPerEdge(:) ! Particles per edge
    REAL(8), ALLOCATABLE:: weightPerEdge(:) ! Weight per edge
    REAL(8):: surface ! Total surface of injection
    TYPE(velDistCont):: v(1:3) !Velocity distribution function in each direction
    procedure(updateInject_interface), pointer, pass:: update => null()
    procedure(printInject_interface),  pointer, pass:: print  => null()
    CONTAINS
      PROCEDURE, PASS:: addParticles

  END TYPE injectGeneric

  abstract interface
    ! Update the values of the particle boundary model
    subroutine updateInject_interface(self)
      import injectGeneric

      class(injectGeneric), intent(inout):: self

    end subroutine updateInject_interface

    ! Write the values of the particle boundary model
    subroutine printInject_interface(self, fileID)
      import injectGeneric

      class(injectGeneric), intent(inout):: self
      integer, intent(in):: fileID

    end subroutine printInject_interface

  end interface

  ! Default type, constant inject of particles
  type, extends(injectGeneric):: injectConstant

  end type injectconstant

  type, extends(injectGeneric):: injectQuasiNeutral

  end type injectQuasiNeutral

  ! Wrapper for injects
  type:: injectWrapper
    class(injectGeneric), allocatable:: obj

  end type injectWrapper

  INTEGER:: nInject
  TYPE(injectWrapper), ALLOCATABLE, target:: injects(:)

  CONTAINS
    !Initialize an injection of particles
    SUBROUTINE initInject(self, i, object, config)
      USE moduleMesh
      USE moduleRefParam
      USE moduleConstParam
      USE moduleSpecies
      USE moduleSolver
      USE moduleErrors
      use json_module
      IMPLICIT NONE

      CLASS(injectGeneric), allocatable, INTENT(inout):: self
      INTEGER, INTENT(in):: i
      character(:), allocatable, intent(in):: object
      type(json_file), intent(inout):: config
      character(:), allocatable:: type
      logical:: found
      REAL(8):: v
      real(8), allocatable:: n(:), temperature(:)
      INTEGER:: sp, ps, particlesPerEdge
      character(:), allocatable:: speciesName
      integer:: ps_index
      REAL(8):: tauInject
      REAL(8):: flow
      CHARACTER(:), ALLOCATABLE:: units
      INTEGER:: e
      REAL(8):: fluxPerStep = 0.D0

      !Type of injection
      call config%get(object // '.type',            type,        found)
      if (.not. found) then
        ! If no type is found, assume constant injection of particles
        type = 'constant'
      end if

      ! Assign type to inject and allocates it
      select case(type)
      case ('constant')
        allocate(injectConstant:: self)

      case ('quasiNeutral')
        allocate(injectQuasiNeutral:: self)

        self%update => updateQuasiNeutral
        self%print  => printQuasiNeutral

      case default
        call criticalError('No injection type ' // type // ' defined', 'initInject')
      end select

      self%id =  i
      call config%get(object // '.name',             self%name,        found)

      !Find species
      CALL config%get(object // '.species', speciesName, found)
      sp = speciesName2Index(speciesName)

      ! Velocity and direction of injection
      call config%get(object // '.v',                v,                found)
      call config%get(object // '.n',                n,                found)
      if (.not. found) then
        allocate(n(1:3))
        n = 0.D0
      end if
      ! Temperature of injection
      call config%get(object // '.T',                temperature,      found)
      ! Flow
      call config%get(object // '.flow',             flow,             found)
      call config%get(object // '.units',            units,            found)
      CALL config%get(object // '.physicalSurface',  ps,               found)
      particlesPerEdge = 0
      CALL config%get(object // '.particlesPerEdge', particlesPerEdge, found)

      self%vMod        =  v / v_ref
      self%n           =  n / NORM2(n)
      self%temperature =  temperature / T_ref

      !Get array index of physical surface for inject
      ps_index = physicalSurface_to_index(ps)

      ! Copy array of edges
      self%edges = physicalSurfaces(ps_index)%edges
      self%nEdges = size(self%edges)

      allocate(self%particlesPerEdge(1:self%nEdges))
      allocate(self%weightPerEdge(1:self%nEdges))

      !Calculates total area
      self%surface = 0.D0
      DO e = 1, self%nEdges
        self%surface = self%surface + self%edges(e)%obj%surface

      END DO

      ! Information about species and flux
      self%species => species(sp)%obj
      tauInject    = tau(self%species%n)
      ! Convert units
      SELECT CASE(units)
      CASE ("sccm")
        !Standard cubic centimeter per minute
        fluxPerStep = flow*sccm2atomPerS

      CASE ("A")
        !Current in Ampers
        SELECT TYPE(sp => self%species)
        CLASS IS(speciesCharged)
          fluxPerStep = flow/(qe*abs(sp%q))

        CLASS DEFAULT
          call criticalError('Attempted to assign a flux in "A" to a species without charge.', 'initInject')

        END SELECT

      CASE ("Am2")
        !Input current in Ampers per square meter
        SELECT TYPE(sp => self%species)
        CLASS IS(speciesCharged)
          fluxPerStep = flow*self%surface*L_ref**2/(qe*abs(sp%q))

        CLASS DEFAULT
          call criticalError('Attempted to assign a flux in "Am2" to a species without charge.', 'initInject')

        END SELECT


      CASE ("part/s")
        !Input current in Ampers
        fluxPerStep = flow

      CASE DEFAULT
        CALL criticalError("No support for units: " // units, 'initInject')

      END SELECT
      fluxPerStep = fluxPerStep * tauInject * ti_ref / self%surface

      !Assign particles per edge
      IF (particlesPerEdge > 0) THEN
        ! Particles per edge defined by the user
        self%particlesPerEdge = particlesPerEdge
        DO e = 1, self%nEdges
          self%weightPerEdge(e) = fluxPerStep * self%edges(e)%obj%surface / real(particlesPerEdge)

        END DO

        self%nParticles = SUM(self%particlesPerEdge)

      ELSE
        ! No particles assigned per edge, use the species weight
        self%weightPerEdge = self%species%weight
        DO e = 1, self%nEdges
          self%particlesPerEdge(e) = max(1,FLOOR(fluxPerStep*self%edges(e)%obj%surface / self%species%weight))
        END DO

        self%nParticles = SUM(self%particlesPerEdge)

        !Rescale weight to match flux
        self%weightPerEdge = fluxPerStep * self%surface / (real(self%nParticles))

      END IF

      !Scale particles for different species steps
      IF (self%nParticles == 0) CALL criticalError("The number of particles for inject is 0.", 'initInject')

    END SUBROUTINE initInject

    ! Update the value of injects
    subroutine updateInjects()
      implicit none
      
      integer:: i

      do i = 1, nInject
        if (associated(injects(i)%obj%update)) then
          call injects(i)%obj%update()

        end if

      end do

    end subroutine updateInjects

    ! Updates the flow in the injection to maintain quasineutrality
    subroutine updateQuasiNeutral(self)
      use moduleMesh, only: meshEdge, meshNode, mesh, qSpecies
      implicit none

      class(injectGeneric), intent(inout):: self
      integer:: e, s, n
      integer, allocatable:: nodes(:)
      class(meshEdge), pointer:: edge
      real(8), allocatable:: density_nodes(:)
      class(meshNode), pointer:: node
      real(8):: den_center
      real(8):: density_incident, density_rest
      real(8):: alpha

      select type(self)
      type is (injectQuasiNeutral)
        do e = 1, self%nEdges
          edge => self%edges(e)%obj

          density_incident = 0.d0
          density_rest     = 0.d0

          nodes = edge%getNodes(edge%nNodes)
          allocate(density_nodes(1:edge%nNodes))
          do s = 1, nSpecies
            do n = 1, edge%nNodes
              node => mesh%nodes(nodes(n))%obj

              density_nodes(n) = node%output(s)%den

            end do

            ! Gather the density at the edge center and multiply by the species charge
            den_center = qSpecies(s)*edge%gatherF(edge%centerXi(), edge%nNodes, density_nodes)

            if (s == self%species%n) then
              density_incident = den_center

            else
              density_rest = density_rest + den_center

            end if

          end do

          ! Correction for this time step
          if (density_rest > 1.0d-10) then
            ! If there is a rest population, correct
            alpha = 1.d0 + density_incident/density_rest

          else
            ! If not, alpha is assumed unchaged
            alpha = 0.d0

          end if

          ! Adjust the weight of particles to match the new current
          self%weightPerEdge(e) = self%weightPerEdge(e) + 1.0d-1 * alpha

        end do

      end select

    end subroutine updateQuasiNeutral

    !Injection of particles
    SUBROUTINE doInjects()
      USE moduleSpecies
      USE moduleSolver
      IMPLICIT NONE

      INTEGER:: i
      
      !$OMP SINGLE
      nPartInj = 0
      DO i = 1, nInject
        IF (solver%pusher(injects(i)%obj%species%n)%pushSpecies) THEN 
          nPartInj = nPartInj + injects(i)%obj%nParticles

        END IF

      END DO

      IF (ALLOCATED(partInj)) DEALLOCATE(partInj)
      ALLOCATE(partInj(1:nPartInj))
      !$OMP END SINGLE

      DO i=1, nInject
        IF (solver%pusher(injects(i)%obj%species%n)%pushSpecies) THEN 
          CALL injects(i)%obj%addParticles()

        END IF
      END DO

    END SUBROUTINE doInjects

    !Add particles for the injection
    SUBROUTINE addParticles(self)
      USE moduleSpecies
      USE moduleSolver
      USE moduleMesh
      USE moduleRandom
      USE moduleErrors
      IMPLICIT NONE

      CLASS(injectGeneric), INTENT(in):: self
      INTEGER, SAVE:: nMin
      INTEGER:: i, e
      INTEGER:: n, sp
      CLASS(meshEdge), POINTER:: randomEdge
      REAL(8):: direction(1:3)

      !Insert particles
      !$OMP SINGLE
      nMin = 0
      DO i = 1, self%id -1
        IF (solver%pusher(injects(i)%obj%species%n)%pushSpecies) THEN 
          nMin = nMin + injects(i)%obj%nParticles

        END IF

      END DO
      nMin = nMin + 1
      !$OMP END SINGLE

      !$OMP DO
      DO e = 1, self%nEdges
        ! Select edge for injection
        randomEdge => self%edges(e)%obj
        ! Inject particles in edge
        DO i = 1, self%particlesPerEdge(e)
          ! Index in the global partInj array
          n = nMin - 1 + SUM(self%particlesPerEdge(1:e-1)) + i
          !Particle is considered to be outside the domain
          partInj(n)%n_in = .FALSE.
          !Random position in edge
          partInj(n)%r = randomEdge%randPos()
          !Assign weight to particle.
          partInj(n)%weight = self%weightPerEdge(e)
          !Volume associated to the edge:
          IF (ASSOCIATED(randomEdge%e1)) THEN
            partInj(n)%cell = randomEdge%e1%n

          ELSEIF (ASSOCIATED(randomEdge%e2)) THEN
            partInj(n)%cell = randomEdge%e2%n

          ELSE
            CALL criticalError("No Volume associated to edge", 'addParticles')

          END IF
          partInj(n)%cellColl = randomEdge%eColl%n
          sp = self%species%n

          !Assign particle type
          partInj(n)%species => self%species

          if (all(self%n == 0.D0)) then
            direction =  randomEdge%normal

          else
            direction = self%n

          end if

          partInj(n)%v = 0.D0

          do while(dot_product(partInj(n)%v, direction) <= 0.d0)
            partInj(n)%v = self%vMod*direction +  (/ self%v(1)%obj%randomVel(), &
                                                     self%v(2)%obj%randomVel(), &
                                                     self%v(3)%obj%randomVel() /)
            !If injecting a no-drift distribution and velocity is negative, reflect
            if ((self%vMod == 0.D0) .and. &
                (dot_product(partInj(n)%v, direction) <= 0.D0)) then
              partInj(n)%v = - partInj(n)%v

            end if

          end do

          !Obtain natural coordinates of particle in cell
          partInj(n)%Xi = mesh%cells(partInj(n)%cell)%obj%phy2log(partInj(n)%r)
          !Push new particle with the minimum time step
          CALL solver%pusher(sp)%pushParticle(partInj(n), tau(sp))
          !Assign cell to new particle
          CALL solver%updateParticleCell(partInj(n))

        END DO

      END DO
      !$OMP END DO

    END SUBROUTINE addParticles

    ! Writes the value of injects
    subroutine writeInjects()
      use moduleCaseparam, only: timeStep
      use moduleOutput, only: fileID_inject, &
                              formatFileName, &
                              informFileCreation,&
                              injectOutput,&
                              generateFilePath, &
                              prefix
      implicit none
      
      integer:: i
      character(:), allocatable:: fileName

      if (injectOutput) then
        fileName = formatFileName(prefix, 'injects', 'csv', timeStep)
        call informFileCreation(fileName)
        open(fileID_inject, file = generateFilePath(fileName))

        do i = 1, nInject
          if (associated(injects(i)%obj%print)) then
            call injects(i)%obj%print(fileID_inject)

          end if

        end do

        close(fileID_inject)

      end if

    end subroutine writeInjects

    ! Write information about quasiNeutral inject
    subroutine printQuasiNeutral(self, fileID)
      use moduleOutput, only: fmtColInt, fmtColReal
      implicit none

      class(injectGeneric), intent(inout):: self
      integer, intent(in):: fileID
      integer:: e

      write(fileID, '(A)') self%name
      write(fileID, '(A,",",A)') '"Edge id"', '"weight of particles"'
      do e = 1, self%nEdges
        write(fileID, '('//fmtColInt//','//fmtColReal//')') self%edges(e)%obj%n, self%weightPerEdge(e)

      end do

    end subroutine printQuasiNeutral

    ! Inits for different velocity distribution functions
    SUBROUTINE initVelDistMaxwellian(velDist, temperature, m)
      IMPLICIT NONE

      CLASS(velDistGeneric), ALLOCATABLE, INTENT(out):: velDist
      REAL(8), INTENT(in):: temperature, m

      velDist = velDistMaxwellian(vTh = DSQRT(2.d0*temperature/m))

    END SUBROUTINE initVelDistMaxwellian

    SUBROUTINE initVelDistHalfMaxwellian(velDist, temperature, m)
      IMPLICIT NONE

      CLASS(velDistGeneric), ALLOCATABLE, INTENT(out):: velDist
      REAL(8), INTENT(in):: temperature, m

      velDist = velDistHalfMaxwellian(vTh = DSQRT(2.d0*temperature/m))

    END SUBROUTINE initVelDistHalfMaxwellian

    SUBROUTINE initVelDistDelta(velDist)
      IMPLICIT NONE

      CLASS(velDistGeneric), ALLOCATABLE, INTENT(out):: velDist

      velDist = velDistDelta()

    END SUBROUTINE initVelDistDelta

END MODULE moduleInject
