!Module to solve the electromagnetic field
MODULE moduleEM
  USE moduleMesh
  USE moduleTable
  IMPLICIT NONE

  contains
    !Assemble the source vector based on the charge density to solve Poisson's equation
    SUBROUTINE assembleSourceVector(vectorF, n_e)
      USE moduleMesh
      USE moduleRefParam
      IMPLICIT NONE

      REAL(8), INTENT(out):: vectorF(1:mesh%numNodes)
      REAL(8), ALLOCATABLE:: localF(:)
      INTEGER, ALLOCATABLE:: nodes(:)
      REAL(8), ALLOCATABLE:: rho(:)
      REAL(8), INTENT(in), OPTIONAL:: n_e(1:mesh%numNodes)
      INTEGER:: nNodes
      INTEGER:: e, i, ni, b
      CLASS(meshNode), POINTER:: node

      !$OMP SINGLE
      vectorF = 0.D0
      !$OMP END SINGLE

      !$OMP DO REDUCTION(+:vectorF)
      DO e = 1, mesh%numCells
        nNodes = mesh%cells(e)%obj%nNodes
        nodes  = mesh%cells(e)%obj%getNodes(nNodes)
        !Calculates charge density (rho) in element nodes
        ALLOCATE(rho(1:nNodes))
        rho = 0.D0
        DO i = 1, nNodes
          ni = nodes(i)
          node => mesh%nodes(ni)%obj
          rho(i) = DOT_PRODUCT(qSpecies(:), node%output(:)%den/(vol_ref*node%v*n_ref))
          IF (PRESENT(n_e)) THEN
            rho(i) = rho(i) - n_e(i)

          END IF

        END DO

        !Calculates local F vector
        localF = mesh%cells(e)%obj%elemF(nNodes, rho)

        !Assign local F to global F
        DO i = 1, nNodes
          ni = nodes(i)
          vectorF(ni) = vectorF(ni) + localF(i)

        END DO

        DEALLOCATE(localF)
        DEALLOCATE(nodes, rho)

      END DO
      !$OMP END DO

      !Apply boundary conditions
      !$OMP SINGLE
      do b = 1, nBoundariesEM
        call boundariesEM(b)%obj%apply(vectorF)

      end do
      !$OMP END SINGLE

    END SUBROUTINE assembleSourceVector

    !Solving the Poisson equation for electrostatic potential
    SUBROUTINE solveElecField()
      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

    FUNCTION BoltzmannElectron(phi, n) RESULT(n_e)
      USE moduleRefParam
      USE moduleConstParam
      IMPLICIT NONE

      INTEGER, INTENT(in):: n
      REAL(8), INTENT(in):: phi(1:n)
      REAL(8):: n_e(1:n)
      REAL(8):: n_e0 = 1.0D16, phi_0 = -500.0D0, T_e = 11604.0

      n_e = n_e0 / n_ref * exp(qe * (phi*Volt_ref - phi_0) / (kb * T_e))

      RETURN

    END FUNCTION BoltzmannElectron

    SUBROUTINE solveElecFieldBoltzmann()
      USE moduleMesh
      USE moduleErrors
      IMPLICIT NONE

      INTEGER, SAVE:: INFO
      INTEGER:: n
      REAL(8), ALLOCATABLE, SAVE:: tempF(:)
      REAL(8), ALLOCATABLE, SAVE:: n_e(:), phi_old(:), phi(:)
      INTEGER:: k
      EXTERNAL:: dgetrs

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

      !$OMP DO
      DO n = 1, mesh%numNodes
        phi_old(n) = mesh%nodes(n)%obj%emData%phi

      END DO
      !$OMP END DO

      !$OMP SINGLE
      DO k = 1, 100
        n_e = BoltzmannElectron(phi_old, mesh%numNodes)
        CALL assembleSourceVector(tempF, n_e)

        CALL dgetrs('N', mesh%numNodes, 1, mesh%K, mesh%numNodes, &
                                mesh%IPIV, tempF,  mesh%numNodes, info)
        phi = tempF

        PRINT *, MAXVAL(n_e), MINVAL(n_e)
        PRINT *, MAXVAL(phi), MINVAL(phi)
        PRINT*, k, "diff = ", MAXVAL(ABS(phi - phi_old))
        phi_old = phi

      END DO
      !$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 = phi_old(n)

        END DO
        !$OMP END DO

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

      END IF

      !$OMP SINGLE
      DEALLOCATE(tempF, n_e, phi_old, phi)
      !$OMP END SINGLE

    END SUBROUTINE solveElecFieldBoltzmann

END MODULE moduleEM
