MODULE moduleSolver

  !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(:)
    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
      IMPLICIT NONE

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

      SELECT CASE(pusherType)
      CASE('2DCylNeutral')
        self%pushParticle => pushCylNeutral

      CASE('2DCylCharged')
        self%pushParticle => pushCylCharged

      CASE('1DCartCharged')
        self%pushParticle => push1DCartCharged

      CASE('1DRadCharged')
        self%pushParticle => push1DRadCharged

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

      END SELECT

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

    END SUBROUTINE initPusher

    SUBROUTINE initEM(self, EMType)
      IMPLICIT NONE

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

      SELECT CASE(EMType)
      CASE('Electrostatic')
        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)%sp
        !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

    !Push one particle. Boris pusher for 2D Cyl Neutral particle
    PURE SUBROUTINE pushCylNeutral(part, tauIn)
      USE moduleSpecies
      IMPLICIT NONE

      TYPE(particle), INTENT(inout):: part
      REAL(8), INTENT(in):: tauIn
      TYPE(particle):: part_temp
      REAL(8):: x_new, y_new, r, sin_alpha, cos_alpha
      REAL(8):: v_p_oh_star(2:3)

      part_temp = part
      !z
      part_temp%v(1) = part%v(1)
      part_temp%r(1) = part%r(1) + part_temp%v(1)*tauIn
      !r,theta
      v_p_oh_star(2) = part%v(2)
      x_new = part%r(2) + v_p_oh_star(2)*tauIn
      v_p_oh_star(3) = part%v(3)
      y_new =             v_p_oh_star(3)*tauIn
      r = DSQRT(x_new**2+y_new**2)
      part_temp%r(2) = r
      IF (r > 0.D0) THEN
        sin_alpha = y_new/r
        cos_alpha = x_new/r
      ELSE
        sin_alpha = 0.D0
        cos_alpha = 1.D0
      END IF
      part_temp%v(2) =  cos_alpha*v_p_oh_star(2)+sin_alpha*v_p_oh_star(3)
      part_temp%v(3) = -sin_alpha*v_p_oh_star(2)+cos_alpha*v_p_oh_star(3)
      part_temp%n_in = .FALSE. !Assume particle is outside until cell is found
      !Copy temporal particle to particle
      part=part_temp

    END SUBROUTINE pushCylNeutral

    !Push one particle. Boris pusher for 2D Cyl Charged particle
    PURE SUBROUTINE pushCylCharged(part, tauIn)
      USE moduleSpecies
      USE moduleEM
      IMPLICIT NONE

      TYPE(particle), INTENT(inout):: part
      REAL(8), INTENT(in):: tauIn
      REAL(8):: v_p_oh_star(2:3)
      TYPE(particle):: part_temp
      REAL(8):: x_new, y_new, r, sin_alpha, cos_alpha
      REAL(8):: qmEFt(1:3)!charge*tauIn*EF/mass

      part_temp = part
      !Get electric field at particle position
      qmEFt = part_temp%qm*gatherElecField(part_temp)*tauIn
      !z
      part_temp%v(1) = part%v(1) + qmEFt(1)
      part_temp%r(1) = part%r(1) + part_temp%v(1)*tauIn
      !r,theta
      v_p_oh_star(2) = part%v(2) + qmEFt(2)
      x_new = part%r(2) + v_p_oh_star(2)*tauIn
      v_p_oh_star(3) = part%v(3) + qmEFt(3)
      y_new =             v_p_oh_star(3)*tauIn
      r = DSQRT(x_new**2+y_new**2)
      part_temp%r(2) = r
      IF (r > 0.D0) THEN
        sin_alpha = y_new/r
        cos_alpha = x_new/r
      ELSE
        sin_alpha = 0.D0
        cos_alpha = 1.D0
      END IF
      part_temp%v(2) =  cos_alpha*v_p_oh_star(2)+sin_alpha*v_p_oh_star(3)
      part_temp%v(3) = -sin_alpha*v_p_oh_star(2)+cos_alpha*v_p_oh_star(3)
      part_temp%n_in = .FALSE. !Assume particle is outside until cell is found
      !Copy temporal particle to particle
      part=part_temp

    END SUBROUTINE pushCylCharged

    !Push charged particles in 1D cartesian coordinates
    PURE SUBROUTINE push1DCartCharged(part, tauIn)
      USE moduleSPecies
      USE moduleEM
      IMPLICIT NONE

      TYPE(particle), INTENT(inout):: part
      REAL(8), INTENT(in):: tauIn
      TYPE(particle):: part_temp
      REAL(8):: qmEFt(1:3)

      part_temp = part
      !Get the electric field at particle position
      qmEFt = part_temp%qm*gatherElecField(part_temp)*tauIn

      !x
      part_temp%v(1) = part%v(1) + qmEFt(1)
      part_temp%r(1) = part%r(1) + part_temp%v(1)*tauIn

      part_temp%n_in = .FALSE.

      part = part_temp

    END SUBROUTINE push1DCartCharged

    !Push one particle. Boris pusher for 1D Radial Charged particle
    PURE SUBROUTINE push1DRadCharged(part, tauIn)
      USE moduleSpecies
      USE moduleEM
      IMPLICIT NONE

      TYPE(particle), INTENT(inout):: part
      REAL(8), INTENT(in):: tauIn
      REAL(8):: v_p_oh_star(1:2)
      TYPE(particle):: part_temp
      REAL(8):: x_new, y_new, r, sin_alpha, cos_alpha
      REAL(8):: qmEFt(1:3)!charge*tauIn*EF/mass

      part_temp = part
      !Time step for the species
      !Get electric field at particle position
      qmEFt = part_temp%qm*gatherElecField(part_temp)*tauMin
      !r,theta
      v_p_oh_star(1) = part%v(1) + qmEFt(1)
      x_new = part%r(1) + v_p_oh_star(1)*tauIn
      v_p_oh_star(2) = part%v(2) + qmEFt(2)
      y_new =             v_p_oh_star(2)*tauIn
      r = DSQRT(x_new**2+y_new**2)
      part_temp%r(1) = r
      IF (r > 0.D0) THEN
        sin_alpha = y_new/r
        cos_alpha = x_new/r
      ELSE
        sin_alpha = 0.D0
        cos_alpha = 1.D0
      END IF
      part_temp%v(1) =  cos_alpha*v_p_oh_star(1)+sin_alpha*v_p_oh_star(2)
      part_temp%v(2) = -sin_alpha*v_p_oh_star(1)+cos_alpha*v_p_oh_star(2)
      part_temp%n_in = .FALSE. !Assume particle is outside until cell is found
      !Copy temporal particle to particle
      part=part_temp

    END SUBROUTINE push1DRadCharged

    !Do the collisions in all the cells
    SUBROUTINE doCollisions()
      USE moduleMesh
      IMPLICIT NONE

      INTEGER:: e

      !$OMP DO SCHEDULE(DYNAMIC)
      DO e=1, mesh%numVols
        CALL mesh%vols(e)%obj%collision()
      END DO
      !$OMP END DO

    END SUBROUTINE doCollisions
    
    SUBROUTINE doReset()
      USE moduleSpecies
      USE moduleMesh
      USE moduleList
      IMPLICIT NONE

      INTEGER:: nn, n, e
      INTEGER, SAVE:: nPartNew
      INTEGER, SAVE:: nInjIn, nOldIn, nWScheme, nCollisions
      TYPE(particle), ALLOCATABLE, SAVE:: partTemp(:)
      TYPE(lNode), POINTER:: partCurr, partNext

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

      !$OMP BARRIER

      !$OMP SINGLE
      CALL MOVE_ALLOC(partOld, partTemp)
      nPartNew = nInjIn + nOldIn + nWScheme + nCollisions
      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)

        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)

        END IF

      END DO
      !$OMP SECTION
      !Reset particles from weighting scheme
      nn = nInjIn + nOldIn
      partCurr => partWScheme%head
      DO n = 1, nWScheme
        partNext => partCurr%next
        partOld(nn+n) = partCurr%part
        DEALLOCATE(partCurr)
        partCurr => partNext

      END DO
      IF (ASSOCIATED(partWScheme%head)) NULLIFY(partWScheme%head)
      IF (ASSOCIATED(partWScheme%tail)) NULLIFY(partWScheme%tail)
      partWScheme%amount = 0

      !$OMP SECTION
      !Reset particles from collisional process
      nn = nInjIn + nOldIn + nWScheme
      partCurr => partCollisions%head
      DO n = 1, nCollisions
        partNext => partCurr%next
        partOld(nn+n) = partCurr%part
        DEALLOCATE(partCurr)
        partCurr => partNext

      END DO
      IF (ASSOCIATED(partCollisions%head)) NULLIFY(partCollisions%head)
      IF (ASSOCIATED(partCollisions%tail)) NULLIFY(partCollisions%tail)
      partCollisions%amount = 0

      !$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
      DO e = 1, mesh%numVols
        mesh%vols(e)%obj%totalWeight = 0.D0
        CALL mesh%vols(e)%obj%listPart_in%erase()

      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

    !Solving the Poisson equation for electrostatic potential
    SUBROUTINE solveElecField()
      USE moduleEM
      USE moduleMesh
      USE moduleErrors
      IMPLICIT NONE

      INTEGER, SAVE:: INFO
      INTEGER:: n
      REAL(8), ALLOCATABLE, SAVE:: tempF(:)
      EXTERNAL:: dgetrs

      !$OMP SINGLE
      ALLOCATE(tempF(1:mesh%numNodes))
      !$OMP END SINGLE

      CALL assembleSourceVector(tempF)

      !$OMP SINGLE
      CALL dgetrs('N', mesh%numNodes, 1, mesh%K, mesh%numNodes, &
                              mesh%IPIV, tempF,  mesh%numNodes, info)
      !$OMP END SINGLE

      IF (info == 0) THEN
        !Suscessful resolution of Poission equation
        !$OMP DO
        DO n = 1, mesh%numNodes
          mesh%nodes(n)%obj%emData%phi = tempF(n)

        END DO
        !$OMP END DO

      ELSE
        !$OMP SINGLE
        CALL criticalError('Poisson equation failed', 'solveElecField')
        !$OMP END SINGLE

      END IF

      !$OMP SINGLE
      DEALLOCATE(tempF)
      !$OMP END SINGLE

    END SUBROUTINE solveElecField

    !Modify particle weight as a function of cell volume and splits particle
    SUBROUTINE volumeWScheme(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
      REAL(8):: fractionVolume, fractionWeight
      INTEGER:: nSplit

      !If particle has change cell, call Weighting scheme
      IF (volOld%n /= volNew%n) THEN
        fractionVolume = volOld%volume/volNew%volume

        part%weight = part%weight * fractionVolume

        fractionWeight = part%weight / species(part%sp)%obj%weight

        IF (fractionWeight >= 2.D0) THEN
          nSplit = FLOOR(fractionWeight)
          CALL splitParticle(part, nSplit, volNew)

        ELSEIF (part%weight < 1.D0) THEN
          !Particle has lost statistical meaning and will be terminated
          part%n_in = .FALSE.

        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

      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
        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)
        CALL vol%listPart_in%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

      volOld => mesh%vols(part%vol)%obj
      CALL volOld%findCell(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
      IMPLICIT NONE

      INTEGER, INTENT(in):: t

      counterOutput = counterOutput + 1
      IF (counterOutput >= triggerOutput .OR. &
          t == tmax .OR. t == 0) THEN

        !Resets output counter
        counterOutput=0

        CALL mesh%printOutput(t)
        CALL mesh%printColl(t)
        CALL mesh%printEM(t)
        WRITE(*, "(5X,A21,I10,A1,I10)") "t/tmax: ",    t,       "/", tmax
        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 == tmax .OR. t == 0) THEN

        !Reset CPU Time counter
        counterCPUTime = 0

        CALL printTime(t, t == 0)

      END IF

    END SUBROUTINE doOutput

END MODULE moduleSolver