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, volOld, volNew)
      USE moduleSpecies
      USE moduleMesh
      IMPLICIT NONE

      TYPE(particle), INTENT(inout):: part
      CLASS(meshVol), POINTER, INTENT(in):: volOld
      CLASS(meshVol), POINTER, INTENT(inout):: volNew

    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
       CALL doProbes(partArray(nStart+n))
       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)
          CALL doProbes(partOld(nn))

        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)
          CALL doProbes(partOld(nn))

        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
      DO s = 1, nSpecies
        DO e = 1, mesh%numVols
          IF (solver%pusher(s)%pushSpecies) THEN 
            CALL mesh%vols(e)%obj%listPart_in(s)%erase()
            mesh%vols(e)%obj%totalWeight(s) = 0.D0

          END IF

        END DO

      END DO

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

          END IF

        END DO

      END DO

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

      !Loops over the particles to scatter them
      !$OMP DO
      DO n = 1, nPartOld
        CALL mesh%vols(partOld(n)%vol)%obj%scatter(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, volOld, volNew)
      USE moduleSpecies
      USE moduleMesh
      USE moduleRandom
      IMPLICIT NONE

      TYPE(particle), INTENT(inout):: part
      CLASS(meshVol), POINTER, INTENT(in):: volOld
      CLASS(meshVol), POINTER, INTENT(inout):: volNew
      REAL(8):: fractionVolume, pSplit

      !If particle changes volume to smaller cell
      IF (volOld%volume > volNew%volume .AND. &
          part%weight >= part%species%weight*1.0D-1) THEN
        fractionVolume = volOld%volume/volNew%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, volNew)

        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, vol)
      USE moduleSpecies
      USE moduleList
      USE moduleMesh
      USE OMP_LIB
      IMPLICIT NONE

      TYPE(particle), INTENT(inout):: part
      INTEGER, INTENT(in):: nSplit
      CLASS(meshVol), INTENT(inout):: vol
      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 OMP_SET_LOCK(lockWScheme)
        CALL partWScheme%add(newPart)
        CALL OMP_UNSET_LOCK(lockWScheme)
        !Add particle to cell list
        CALL OMP_SET_lock(vol%lock)
        sp = part%species%n
        CALL vol%listPart_in(sp)%add(newPart)
        CALL OMP_UNSET_lock(vol%lock)

      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(meshVol), POINTER:: volOld, volNew

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

      volOld => mesh%vols(part%vol)%obj
      CALL volOld%findCell(part)
      CALL findCellColl(part)
      volNew => mesh%vols(part%vol)%obj
      !Call the NA shcme
      IF (ASSOCIATED(self%weightingScheme)) THEN
        CALL self%weightingScheme(part, volOld, volNew)

      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
      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)
        IF (ASSOCIATED(meshForMCC)) CALL meshForMCC%printColl(t)
        CALL mesh%printEM(t)
        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()

      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