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

  ! Generic type for electromagnetic boundary conditions
  TYPE, PUBLIC, ABSTRACT:: boundaryEMGeneric
    INTEGER:: nNodes
    TYPE(meshNodePointer), ALLOCATABLE:: nodes(:)

    CONTAINS
      PROCEDURE(applyEM_interface), DEFERRED, PASS:: apply

  END TYPE boundaryEMGeneric

  ABSTRACT INTERFACE
    ! Apply boundary condition to the load vector for the Poission equation
    SUBROUTINE applyEM_interface(self, vectorF)
      IMPORT boundaryEMGeneric
      CLASS(boundaryEMGeneric), INTENT(in):: self
      REAL(8), INTENT(inout):: vectorF(:)

    END SUBROUTINE applyEM_interface

  END INTERFACE

  TYPE, EXTENDS(boundaryEMGeneric):: boundaryEMDirichlet
    REAL(8):: potential

    CONTAINS
      ! boundaryEMGeneric DEFERRED PROCEDURES
      PROCEDURE, PASS:: apply => applyDirichlet

  END TYPE boundaryEMDirichlet

  TYPE, EXTENDS(boundaryEMGeneric):: boundaryEMDirichletTime
    REAL(8):: potential
    TYPE(table1D):: temporalProfile

    CONTAINS
      ! boundaryEMGeneric DEFERRED PROCEDURES
      PROCEDURE, PASS:: apply => applyDirichletTime

  END TYPE boundaryEMDirichletTime

  ! Container for boundary conditions
  TYPE:: boundaryEMCont
    CLASS(boundaryEMGeneric), ALLOCATABLE:: obj

  END TYPE boundaryEMCont

  INTEGER:: nBoundaryEM
  TYPE(boundaryEMCont), ALLOCATABLE:: boundaryEM(:)

  !Information of charge and reference parameters for rho vector
  REAL(8), ALLOCATABLE:: qSpecies(:)

  CONTAINS
    SUBROUTINE findNodes(self, physicalSurface)
      USE moduleMesh
      IMPLICIT NONE

      CLASS(boundaryEMGeneric), INTENT(inout):: self
      INTEGER, INTENT(in):: physicalSurface
      CLASS(meshEdge), POINTER:: edge
      INTEGER, ALLOCATABLE:: nodes(:), nodesEdge(:)
      INTEGER:: nNodes, nodesNew
      INTEGER:: e, n

      !Temporal array to hold nodes
      ALLOCATE(nodes(0))

      ! Loop thorugh the edges and identify those that are part of the boundary
      DO e = 1, mesh%numEdges
        edge => mesh%edges(e)%obj
        IF (edge%physicalSurface == physicalSurface) THEN
          ! Edge is of the right boundary index
          ! Get nodes in the edge
          nNodes = edge%nNodes
          nodesEdge = edge%getNodes(nNodes)
          ! Collect all nodes that are not already in the temporal array
          DO n = 1, nNodes
            IF (ANY(nodes == nodesEdge(n))) THEN
              ! Node already in array, skip
              CYCLE

            ELSE
              ! If not, add element to array of nodes
              nodes = [nodes, nodesEdge(n)]

            END IF

          END DO

        END IF

      END DO

      ! Point boundary to nodes
      nNodes = SIZE(nodes)
      ALLOCATE(self%nodes(nNodes))
      self%nNodes = nNodes
      DO n = 1, nNodes
        self%nodes(n)%obj => mesh%nodes(nodes(n))%obj

      END DO

    END SUBROUTINE findNodes

    ! Initialize Dirichlet boundary condition
    SUBROUTINE initDirichlet(self, physicalSurface, potential)
      USE moduleRefParam, ONLY: Volt_ref
      IMPLICIT NONE

      CLASS(boundaryEMGeneric), ALLOCATABLE, INTENT(out):: self
      INTEGER, INTENT(in):: physicalSurface
      REAL(8), INTENT(in):: potential

      ! Allocate boundary edge
      ALLOCATE(boundaryEMDirichlet:: self)

      SELECT TYPE(self)
      TYPE IS(boundaryEMDirichlet)
        self%potential = potential / Volt_ref

        CALL findNodes(self, physicalSurface)

      END SELECT

    END SUBROUTINE initDirichlet

    ! Initialize Dirichlet boundary condition
    SUBROUTINE initDirichletTime(self, physicalSurface, potential, temporalProfile)
      USE moduleRefParam, ONLY: Volt_ref, ti_ref
      IMPLICIT NONE

      CLASS(boundaryEMGeneric), ALLOCATABLE, INTENT(out):: self
      INTEGER, INTENT(in):: physicalSurface
      REAL(8), INTENT(in):: potential
      CHARACTER(:), ALLOCATABLE, INTENT(in):: temporalProfile

      ! Allocate boundary edge
      ALLOCATE(boundaryEMDirichletTime:: self)

      SELECT TYPE(self)
      TYPE IS(boundaryEMDirichletTime)
        self%potential = potential / Volt_ref

        CALL findNodes(self, physicalSurface)

        CALL self%temporalProfile%init(temporalProfile)

        CALL self%temporalProfile%convert(1.D0/ti_ref, 1.D0)

      END SELECT

    END SUBROUTINE initDirichletTime

    !Apply Dirichlet boundary condition to the poisson equation
    SUBROUTINE applyDirichlet(self, vectorF)
      USE moduleMesh
      IMPLICIT NONE

      CLASS(boundaryEMDirichlet), INTENT(in):: self
      REAL(8), INTENT(inout):: vectorF(:)
      INTEGER:: n, ni

      DO n = 1, self%nNodes
        self%nodes(n)%obj%emData%phi  = self%potential
        vectorF(self%nodes(n)%obj%n)  = self%nodes(n)%obj%emData%phi

      END DO

    END SUBROUTINE applyDirichlet

    !Apply Dirichlet boundary condition with time temporal profile
    SUBROUTINE applyDirichletTime(self, vectorF)
      USE moduleMesh
      USE moduleCaseParam, ONLY: timeStep, tauMin
      IMPLICIT NONE

      CLASS(boundaryEMDirichletTime), INTENT(in):: self
      REAL(8), INTENT(inout):: vectorF(:)
      REAL(8):: timeFactor
      INTEGER:: n, ni

      timeFactor = self%temporalProfile%get(DBLE(timeStep)*tauMin)

      DO n = 1, self%nNodes
        self%nodes(n)%obj%emData%phi  = self%potential * timeFactor
        vectorF(self%nodes(n)%obj%n)  = self%nodes(n)%obj%emData%phi

      END DO

    END SUBROUTINE applyDirichletTime

    !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, nBoundaryEM
        call boundaryEM(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
      INTEGER:: i

      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