!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 inject 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
    real(8):: qSign_incident

  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 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

      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

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

      END IF

      ! Total number of particles to inject
      self%nParticles = SUM(self%particlesPerEdge)

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

      ! Specific parameters per type
      select type(self)
      type is(injectQuasiNeutral)
        ! Get the sign of the charged species
        self%qSign_incident = sign(1.d0, self%species%qm)

      end select

    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, meshCell
      implicit none

      class(injectGeneric), intent(inout):: self
      integer:: e
      class(meshEdge), pointer:: edge
      class(meshCell), pointer:: cell
      real(8):: Xi(1:3)
      real(8):: EF_normal
      real(8):: alpha

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

          Xi = edge%centerXi()

          if (associated(edge%e1)) then
            cell => edge%e1

          else
            cell => edge%e2

          end if

          ! Projection of EF on the edge normal vector
          EF_normal = dot_product(cell%gatherElectricField(Xi), edge%normal)

          ! Correction for this time step
          alpha = self%qSign_incident * EF_normal

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

          ! Limit the weight of a macroparticle to 1.0
          self%weightPerEdge(e) = max(self%weightPerEdge(e), 1.d0)

        end do

      end select

    end subroutine updateQuasiNeutral

    !Add particles for the injection
    SUBROUTINE addParticles(self)
      USE moduleSpecies
      USE moduleMesh
      USE moduleRandom
      USE moduleErrors
      use moduleList, only: partInj
      IMPLICIT NONE

      CLASS(injectGeneric), INTENT(in):: self
      type(particle), pointer:: part
      INTEGER:: i, e
      INTEGER:: sp
      CLASS(meshEdge), POINTER:: edge

      REAL(8):: direction(1:3)

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

          ELSEIF (ASSOCIATED(edge%e2)) THEN
            part%cell = edge%e2%n

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

          END IF
          part%cellColl = edge%eColl%n
          sp = self%species%n

          !Assign particle type
          part%species => self%species

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

          else
            direction = self%n

          end if

          part%v = 0.D0

          ! do while(dot_product(partInj(n)%v, direction) <= 0.d0)
            part%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(part%v, direction) <= 0.D0)) then
              part%v = - part%v

            end if

          ! end do

          !Obtain natural coordinates of particle in cell
          part%Xi = mesh%cells(part%cell)%obj%phy2log(part%r)

          ! Add particle to global list
          call partInj%setLock()
          call partInj%add(part)
          call partInj%unsetLock()

        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(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

END MODULE moduleInject
