MODULE moduleSolver
  USE moduleAverage

  !Generic type for pusher of particles
  TYPE, PUBLIC:: pusherGeneric
    PROCEDURE(push_interafece), POINTER, NOPASS:: pushParticle => NULL()
    !Boolean to indicate if the species is moved in the iteration
    LOGICAL:: pushSpecies
    !How many interations between advancing the species
    INTEGER:: every
    CONTAINS
      PROCEDURE, PASS:: init => initPusher

  END TYPE pusherGeneric

  !Generic type for solver
  TYPE, PUBLIC:: solverGeneric
    TYPE(pusherGeneric), ALLOCATABLE:: pusher(:)
    TYPE(averageGeneric), ALLOCATABLE:: averageScheme
    PROCEDURE(solveEM_interface), POINTER, NOPASS:: solveEM => NULL()
    PROCEDURE(weightingScheme_interface), POINTER, NOPASS:: weightingScheme => NULL()
    CONTAINS
      PROCEDURE, PASS:: initEM
      PROCEDURE, PASS:: initWS
      PROCEDURE, PASS:: updateParticleCell
      PROCEDURE, PASS:: updatePushSpecies

  END TYPE solverGeneric

  INTERFACE
    !Push a particle
    PURE SUBROUTINE push_interafece(part, tauIn)
      USE moduleSpecies

      TYPE(particle), INTENT(inout):: part
      REAL(8), INTENT(in):: tauIn

    END SUBROUTINE push_interafece

    !Solve the electromagnetic field
    SUBROUTINE solveEM_interface()
    
    END SUBROUTINE solveEM_interface

    !Apply nonAnalogue scheme to a particle
    SUBROUTINE weightingScheme_interface(part, cellOld, cellNew)
      USE moduleSpecies
      USE moduleMesh
      IMPLICIT NONE

      TYPE(particle), INTENT(inout):: part
      CLASS(meshCell), POINTER, INTENT(in):: cellOld
      CLASS(meshCell), POINTER, INTENT(inout):: cellNew

    END SUBROUTINE weightingScheme_interface

  END INTERFACE

  TYPE(solverGeneric):: solver

  CONTAINS
    !Init Pusher
    SUBROUTINE initPusher(self, pusherType, tau, tauSp)
      USE moduleErrors
      USE modulePusher
      USE moduleMesh, ONLY: mesh
      IMPLICIT NONE

      CLASS(pusherGeneric), INTENT(out):: self
      CHARACTER(:), ALLOCATABLE:: pusherType
      REAL(8):: tau, tauSp

      !TODO: Reorganize if Cart pushers are combined
      SELECT CASE(mesh%dimen)
      CASE(0)
        self%pushParticle => push0D

      CASE(1:3)
        SELECT CASE(mesh%geometry)
        CASE ('Cart')
          SELECT CASE(pusherType)
          CASE('Neutral')
            self%pushParticle => pushCartNeutral

          CASE('Electrostatic')
            self%pushParticle => pushCartElectrostatic

          CASE('Electromagnetic')
            self%pushParticle => pushCartElectromagnetic

          CASE DEFAULT
            CALL criticalError('Pusher ' // pusherType // ' not found for Cart','initPusher')

          END SELECT

        CASE('Cyl')
          SELECT CASE(pusherType)
          CASE('Neutral')
            self%pushParticle => push2DCylNeutral

          CASE('Electrostatic')
            self%pushParticle => push2DCylElectrostatic

          CASE DEFAULT
            CALL criticalError('Pusher ' // pusherType // ' not found for Cyl','initPusher')

          END SELECT

        CASE('Rad')
          SELECT CASE(pusherType)
          CASE('Neutral')
            self%pushParticle => push1DRadNeutral

          CASE('Electrostatic')
            self%pushParticle => push1DRadElectrostatic

          CASE DEFAULT
            CALL criticalError('Pusher ' // pusherType // ' not found for Rad','initPusher')

          END SELECT

        END SELECT

      END SELECT

      self%pushSpecies = .FALSE.
      self%every = INT(tauSp/tau)

    END SUBROUTINE initPusher

    SUBROUTINE initEM(self, EMType)
      USE moduleEM
      IMPLICIT NONE

      CLASS(solverGeneric), INTENT(inout):: self
      CHARACTER(:), ALLOCATABLE:: EMType

      SELECT CASE(EMType)
      CASE('Electrostatic','ConstantB')
        self%solveEM => solveElecField

      END SELECT

    END SUBROUTINE initEM

    !Initialize the non-analogue scheme
    SUBROUTINE initWS(self, NAType)
      USE moduleMesh
      IMPLICIT NONE

      CLASS(solverGeneric), INTENT(inout):: self
      CHARACTER(:), ALLOCATABLE:: NAType

      SELECT CASE(NAType)
      CASE ('Volume')
        self%weightingScheme => volumeWScheme

      END SELECT

    END SUBROUTINE initWS
    
    !Do all pushes for particles
    SUBROUTINE doPushes()
      USE moduleSpecies
      USE moduleMesh
      IMPLICIT NONE

      INTEGER:: n
      INTEGER:: sp

      !$OMP DO
      DO n = 1, nPartOld
        !Select species type
        sp = partOld(n)%species%n
        !Checks if the species sp is update this iteration
        IF (solver%pusher(sp)%pushSpecies) THEN
          !Push particle
          CALL solver%pusher(sp)%pushParticle(partOld(n), tau(sp))
          !Find cell in wich particle reside
          CALL solver%updateParticleCell(partOld(n))

        END IF

      END DO
      !$OMP END DO

    END SUBROUTINE doPushes

    !Takes the particles from a list and put them into an array
    !nStart indicates the last fill index in the array
    SUBROUTINE resetList(partList, partArray, nStart)
      USE moduleSpecies
      USE moduleList
      USE moduleProbe
      IMPLICIT NONE

      TYPE(listNode), INTENT(inout):: partList
      TYPE(particle), INTENT(inout), DIMENSION(1:):: partArray
      INTEGER, INTENT(in):: nStart
      TYPE(lNode), POINTER:: partCurr, partNext
      INTEGER:: n

      partCurr => partList%head
      DO n = 1, partList%amount
       partNext => partCurr%next
       partArray(nStart + n) = partCurr%part
       IF (nProbes > 0) THEN
         CALL doProbes(partArray(nStart+n))

       END IF
       DEALLOCATE(partCurr)
       partCurr => partNext

      END DO

      !Resest head, tail and number of nodes in the list
      IF (ASSOCIATED(partList%head)) NULLIFY(partList%head)
      IF (ASSOCIATED(partList%tail)) NULLIFY(partList%tail)
      partList%amount = 0

    END SUBROUTINE resetList

    SUBROUTINE doReset()
      USE moduleSpecies
      USE moduleMesh
      USE moduleList
      USE moduleProbe
      IMPLICIT NONE

      INTEGER:: nn, n, e
      INTEGER, SAVE:: nPartNew
      INTEGER, SAVE:: nInjIn, nOldIn, nWScheme, nCollisions, nSurfaces
      TYPE(particle), ALLOCATABLE, SAVE:: partTemp(:)
      INTEGER:: s

      !$OMP SECTIONS
      !$OMP SECTION
      nInjIn = 0
      IF (ALLOCATED(partInj)) THEN
        nInjIn = COUNT(partInj%n_in)

      END IF
      !$OMP SECTION
      nOldIn = 0
      IF (ALLOCATED(partOld)) THEN
        nOldIn = COUNT(partOld%n_in)

      END IF
      !$OMP SECTION
      nWScheme = partWScheme%amount

      !$OMP SECTION
      nCollisions = partCollisions%amount

      !$OMP SECTION
      nSurfaces = partSurfaces%amount

      !$OMP END SECTIONS

      !$OMP BARRIER

      !$OMP SINGLE
      CALL MOVE_ALLOC(partOld, partTemp)
      nPartNew = nInjIn + nOldIn + nWScheme + nCollisions + nSurfaces
      ALLOCATE(partOld(1:nPartNew))
      !$OMP END SINGLE

      !$OMP SECTIONS
      !$OMP SECTION
      !Reset particles from injection
      nn = 0
      DO n = 1, nPartInj
        IF (partInj(n)%n_in) THEN
          nn = nn + 1
          partOld(nn) = partInj(n)
          IF (nProbes > 0) THEN
            CALL doProbes(partOld(nn))

          END IF

        END IF

      END DO

      !$OMP SECTION
      !Reset particles from previous iteration
      nn = nInjIn
      DO n = 1, nPartOld
        IF (partTemp(n)%n_in) THEN
          nn = nn + 1
          partOld(nn) = partTemp(n)
          IF (nProbes > 0) THEN
            CALL doProbes(partOld(nn))

          END IF

        END IF

      END DO

      !$OMP SECTION
      !Reset particles from weighting scheme
      nn = nInjIn + nOldIn
      CALL resetList(partWScheme, partOld, nn)

      !$OMP SECTION
      !Reset particles from collisional process
      nn = nInjIn + nOldIn + nWScheme
      CALL resetList(partCollisions, partOld, nn)

      !$OMP SECTION
      !Reset particles from surface process
      nn = nInjIn + nOldIn + nWScheme + nCollisions
      CALL resetList(partSurfaces, partOld, nn)

      !$OMP SECTION
      !Reset output in nodes
      DO e = 1, mesh%numNodes
        CALL mesh%nodes(e)%obj%resetOutput()

      END DO

      !$OMP SECTION
      !Erase the list of particles inside the cell if particles have been pushed
      IF (doMCC) THEN
        DO s = 1, nSpecies
          DO e = 1, mesh%numCells
            IF (solver%pusher(s)%pushSpecies) THEN 
              CALL mesh%cells(e)%obj%listPart_in(s)%erase()
              mesh%cells(e)%obj%totalWeight(s) = 0.D0

            END IF

          END DO

        END DO

      END IF

      !$OMP SECTION
      !Erase the list of particles inside the cell in coll mesh
      IF (doubleMesh) THEN
        DO s = 1, nSpecies
          DO e = 1, meshColl%numCells
            IF (solver%pusher(s)%pushSpecies) THEN 
              CALL meshColl%cells(e)%obj%listPart_in(s)%erase()
              meshColl%cells(e)%obj%totalWeight(s) = 0.D0

            END IF

          END DO

        END DO
      
      END IF

      !$OMP END SECTIONS

      !$OMP SINGLE
      nPartOld = nPartNew
      !$OMP END SINGLE

    END SUBROUTINE doReset

    !Scatter particles in the grid
    SUBROUTINE doScatter
      USE moduleSpecies
      USE moduleMesh
      IMPLICIT NONE

      INTEGER:: n
      CLASS(meshCell), POINTER:: cell

      !Loops over the particles to scatter them
      !$OMP DO
      DO n = 1, nPartOld
        cell => mesh%cells(partOld(n)%cell)%obj
        CALL cell%scatter(cell%nNodes, partOld(n))

      END DO
      !$OMP END DO

    END SUBROUTINE doScatter

    SUBROUTINE doEMField()
      IMPLICIT NONE

      IF (ASSOCIATED(solver%solveEM)) THEN
        CALL solver%solveEM()

      END IF

    END SUBROUTINE doEMField

    !Split particles as a function of cell volume and splits particle
    SUBROUTINE volumeWScheme(part, cellOld, cellNew)
      USE moduleSpecies
      USE moduleMesh
      USE moduleRandom
      IMPLICIT NONE

      TYPE(particle), INTENT(inout):: part
      CLASS(meshCell), POINTER, INTENT(in):: cellOld
      CLASS(meshCell), POINTER, INTENT(inout):: cellNew
      REAL(8):: fractionVolume, pSplit

      !If particle changes volume to smaller cell
      IF (cellOld%volume > cellNew%volume .AND. &
          part%weight >= part%species%weight*1.0D-1) THEN
        fractionVolume = cellOld%volume/cellNew%volume

        !Calculate probability of splitting particle
        pSplit = 1.D0 - DEXP(-fractionVolume*1.0D-1)

        IF (random() < pSplit) THEN
          !Split particle in two
          CALL splitParticle(part, 2, cellNew)

        END IF

      END IF

    END SUBROUTINE volumeWScheme

    !Subroutine to split the particle 'part' into a number 'nSplit' of particles.
    !'nSplit-1' particles are added to the partNAScheme list
    SUBROUTINE splitParticle(part, nSplit, cell)
      USE moduleSpecies
      USE moduleList
      USE moduleMesh
      USE OMP_LIB
      IMPLICIT NONE

      TYPE(particle), INTENT(inout):: part
      INTEGER, INTENT(in):: nSplit
      CLASS(meshCell), INTENT(inout):: cell
      REAL(8):: newWeight
      TYPE(particle), POINTER:: newPart
      INTEGER:: p
      INTEGER:: sp

      newWeight = part%weight / nSplit

      !Assign new weight to original particle
      part%weight = newWeight
      
      !Add new particles to list of NA particles
      DO p = 2, nSplit
        !Allocate the pointer for the new particles
        ALLOCATE(newPart)
        !Copy data from original particle
        newPart = part
        !Add particle to list of new particles from weighting scheme
        CALL partWScheme%setLock()
        CALL partWScheme%add(newPart)
        CALL partWScheme%unsetLock()
        !Add particle to cell list
        sp = part%species%n
        IF (doMCC) THEN
          CALL OMP_SET_lock(cell%lock)
          CALL cell%listPart_in(sp)%add(newPart)
          CALL OMP_UNSET_lock(cell%lock)

        END IF

      END DO

    END SUBROUTINE splitParticle

    SUBROUTINE updateParticleCell(self, part)
      USE moduleSpecies
      USE moduleMesh
      IMPLICIT NONE

      CLASS(solverGeneric), INTENT(in):: self
      TYPE(particle), INTENT(inout):: part
      CLASS(meshCell), POINTER:: cellOld, cellNew

      !Assume that particle is outside the domain
      part%n_in = .FALSE.

      cellOld => mesh%cells(part%cell)%obj
      CALL cellOld%findCell(part)
      CALL findCellColl(part)
      !Call the NA shcme
      IF (ASSOCIATED(self%weightingScheme)) THEN
        cellNew => mesh%cells(part%cell)%obj
        CALL self%weightingScheme(part, cellOld, cellNew)

      END IF

    END SUBROUTINE updateParticleCell

    !Update the information about if a species needs to be moved this iteration
    SUBROUTINE updatePushSpecies(self, t)
      USE moduleSpecies
      IMPLICIT NONE

      CLASS(solverGeneric), INTENT(inout):: self
      INTEGER, INTENT(in):: t
      INTEGER:: s

      DO s=1, nSpecies
        self%pusher(s)%pushSpecies = MOD(t, self%pusher(s)%every) == 0

      END DO

    END SUBROUTINE updatePushSpecies

    !Output the different data and information
    SUBROUTINE doOutput(t)
      USE moduleMesh
      USE moduleOutput
      USE moduleSpecies
      USE moduleCompTime
      USE moduleProbe
      USE moduleMeshOutputVTU !TEMPORARY TO TEST VTU OUTPUT
      IMPLICIT NONE

      INTEGER, INTENT(in):: t

      counterOutput = counterOutput + 1
      IF (counterOutput >= triggerOutput .OR. &
          t == tFinal .OR. t == tInitial) THEN

        !Resets output counter
        counterOutput=0

        CALL outputProbes(t)
        CALL mesh%printOutput(t)
        CALL printOutputVTU(mesh, t) !TEMPORARY TO TEST VTU OUTPUT
        IF (ASSOCIATED(meshForMCC)) CALL meshForMCC%printColl(t)
        IF (ASSOCIATED(meshForMCC)) CALL printCollVTU(meshForMCC,t) !TEMPORARY TO TEST VTU OUTPUT
        CALL mesh%printEM(t)
        CALL printEMVTU(mesh, t) !TEMPORARY TO TEST VTU OUTPUT
        WRITE(*, "(5X,A21,I10,A1,I10)") "t/tFinal: ", t, "/", tFinal
        WRITE(*, "(5X,A21,I10)")        "Particles: ", nPartOld
        IF (t == 0) THEN
          WRITE(*, "(5X,A21,F8.1,A2)")   "     init time: ",    1.D3*tStep, "ms"

        ELSE
          WRITE(*, "(5X,A21,F8.1,A2)")   "iteration time: ",    1.D3*tStep, "ms"

        END IF

        IF (nPartOld > 0) THEN
          WRITE(*, "(5X,A21,F8.1,A2)") "avg t/particle: ", 1.D9*tStep/DBLE(nPartOld), "ns"

        END IF
        WRITE(*,*)

      END IF

      counterCPUTime = counterCPUTime + 1
      IF (counterCPUTime >= triggerCPUTime .OR. &
          t == tFinal .OR. t == tInitial) THEN

        !Reset CPU Time counter
        counterCPUTime = 0

        CALL printTime(t, t == 0)

      END IF

      !Output average values
      IF (useAverage .AND. t == tFinal) THEN
        CALL mesh%printAverage()
        CALL printAverageVTU(mesh) !TEMPORARY TO TEST VTU OUTPUT

      END IF

    END SUBROUTINE doOutput

    SUBROUTINE doAverage(t)
      USE moduleAverage
      USE moduleMesh
      IMPLICIT NONE

      INTEGER, INTENT(in):: t
      INTEGER:: tAverage, n


      IF (useAverage) THEN
        tAverage = t - tAverageStart

        IF (tAverage == 1) THEN
          !First iteration in which average scheme is used
          DO n = 1, mesh%numNodes
            CALL averageScheme(n)%firstAverage(mesh%nodes(n)%obj%output)

          END DO
        ELSEIF (tAverage > 1) THEN
          DO n = 1, mesh%numNodes
            CALL averageScheme(n)%updateAverage(mesh%nodes(n)%obj%output, tAverage)

          END DO

        END IF

      END IF

    END SUBROUTINE doAverage


END MODULE moduleSolver