!injection of particles
MODULE moduleInject
  USE moduleSpecies

  !Generic type for velocity distribution function
  TYPE, ABSTRACT:: velDistGeneric
    CONTAINS
      !Returns random velocity from distribution function
      PROCEDURE(randomVel_interface), DEFERRED, PASS:: randomVel

  END TYPE velDistGeneric

  ABSTRACT INTERFACE
    FUNCTION randomVel_interface(self) RESULT(v)
      IMPORT velDistGeneric

      CLASS(velDistGeneric), INTENT(in):: self
      REAL(8):: v

    END FUNCTION randomVel_interface

  END INTERFACE

  !Container for velocity distributions
  TYPE:: velDistCont
    CLASS(velDistGeneric), ALLOCATABLE:: obj

  END TYPE velDistCont

  !Maxwellian distribution function
  TYPE, EXTENDS(velDistGeneric):: velDistMaxwellian
    REAL(8):: vTh !Thermal Velocity
    CONTAINS
      PROCEDURE, PASS:: randomVel => randomVelMaxwellian
  
  END TYPE velDistMaxwellian

  TYPE, EXTENDS(velDistGeneric):: velDistHalfMaxwellian
    REAL(8):: vTh !Thermal Velocity
    CONTAINS
      PROCEDURE, PASS:: randomVel => randomVelHalfMaxwellian
  
  END TYPE velDistHalfMaxwellian

  !Dirac's delta distribution function
  TYPE, EXTENDS(velDistGeneric):: velDistDelta
    CONTAINS
      PROCEDURE, PASS:: randomVel => randomVelDelta

  END TYPE velDistDelta

  !Generic injection of particles
  TYPE:: 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
    INTEGER, ALLOCATABLE:: edges(:) !Array with 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
    CONTAINS
      PROCEDURE, PASS:: init         => initInject
      PROCEDURE, PASS:: addParticles

  END TYPE injectGeneric

  INTEGER:: nInject
  TYPE(injectGeneric), ALLOCATABLE:: inject(:)

  CONTAINS
    !Initialize an injection of particles
    SUBROUTINE initInject(self, i, v, n, temperature, flow, units, sp, physicalSurface, particlesPerEdge)
      USE moduleMesh
      USE moduleRefParam
      USE moduleConstParam
      USE moduleSpecies
      USE moduleSolver
      USE moduleErrors
      IMPLICIT NONE

      CLASS(injectGeneric), INTENT(inout):: self
      INTEGER, INTENT(in):: i
      REAL(8), INTENT(in):: v, n(1:3), temperature(1:3)
      INTEGER, INTENT(in):: sp, physicalSurface, particlesPerEdge
      REAL(8):: tauInject
      REAL(8), INTENT(in):: flow
      CHARACTER(:), ALLOCATABLE, INTENT(in):: units
      INTEGER:: e, et
      INTEGER:: phSurface(1:mesh%numEdges)
      INTEGER:: nVolColl
      REAL(8):: fluxPerStep = 0.D0

      self%id          =  i
      self%vMod        =  v / v_ref
      self%n           =  n / NORM2(n)
      self%temperature =  temperature / T_ref
      !Gets the edge elements from which particles are injected
      DO e = 1, mesh%numEdges
        phSurface(e) = mesh%edges(e)%obj%physicalSurface

      END DO
      self%nEdges = COUNT(phSurface == physicalSurface)
      ALLOCATE(self%edges(1:self%nEdges))
      ALLOCATE(self%particlesPerEdge(1:self%nEdges))
      ALLOCATE(self%weightPerEdge(1:self%nEdges))
      et = 0
      DO e=1, mesh%numEdges
        IF (mesh%edges(e)%obj%physicalSurface == physicalSurface) THEN
          et = et + 1
          self%edges(et) = mesh%edges(e)%obj%n
          !Assign connectivity between injection edge and meshColl volume
          IF (doubleMesh) THEN
            nVolColl = findCellBrute(meshColl, mesh%edges(e)%obj%randPos())
            IF (nVolColl > 0) THEN
              mesh%edges(e)%obj%eColl => meshColl%cells(nVolColl)%obj

            ELSE
              CALL criticalError("No connection between edge and meshColl", "initInject")

            END IF

          ELSE
            IF (ASSOCIATED(mesh%edges(e)%obj%e1)) THEN
              mesh%edges(e)%obj%eColl => mesh%edges(e)%obj%e1

            ELSE
              mesh%edges(e)%obj%eColl => mesh%edges(e)%obj%e2

            END IF

          END IF

        END IF

      END DO

      !Calculates total area
      self%surface = 0.D0
      DO et = 1, self%nEdges
        self%surface = self%surface + mesh%edges(self%edges(et))%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 et = 1, self%nEdges
          self%weightPerEdge(et) = fluxPerStep*mesh%edges(self%edges(et))%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 et = 1, self%nEdges
          self%particlesPerEdge(et) = max(1,FLOOR(fluxPerStep*mesh%edges(self%edges(et))%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

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

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

        END IF

      END DO

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

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

        END IF
      END DO

    END SUBROUTINE doInjects

    SUBROUTINE initVelDistMaxwellian(velDist, temperature, m)
      IMPLICIT NONE

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

      velDist = velDistMaxwellian(vTh = DSQRT(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(temperature/m))

    END SUBROUTINE initVelDistHalfMaxwellian

    SUBROUTINE initVelDistDelta(velDist)
      IMPLICIT NONE

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

      velDist = velDistDelta()

    END SUBROUTINE initVelDistDelta

    !Random velocity from Maxwellian distribution
    FUNCTION randomVelMaxwellian(self) RESULT (v)
      USE moduleRandom
      IMPLICIT NONE

      CLASS(velDistMaxwellian), INTENT(in):: self
      REAL(8):: v
      v = 0.D0

      v = sqrt(2.0)*self%vTh*randomMaxwellian()

    END FUNCTION randomVelMaxwellian

    !Random velocity from Half Maxwellian distribution
    FUNCTION randomVelHalfMaxwellian(self) RESULT (v)
      USE moduleRandom
      IMPLICIT NONE

      CLASS(velDistHalfMaxwellian), INTENT(in):: self
      REAL(8):: v
      v = 0.D0

      v = sqrt(2.0)*self%vTh*randomHalfMaxwellian()

    END FUNCTION randomVelHalfMaxwellian

    !Random velocity from Dirac's delta distribution
    PURE FUNCTION randomVelDelta(self) RESULT(v)
      IMPLICIT NONE

      CLASS(velDistDelta), INTENT(in):: self
      REAL(8):: v

      v = 0.D0

    END FUNCTION randomVelDelta

    !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(inject(i)%species%n)%pushSpecies) THEN 
          nMin = nMin + inject(i)%nParticles

        END IF

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

      !$OMP DO
      DO e = 1, self%nEdges
        ! Select edge for injection
        randomEdge => mesh%edges(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(direction, partInj(n)%v) < 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

END MODULE moduleInject