!moduleMesh: General module for Finite Element mesh
MODULE moduleMesh
  USE moduleList
  USE moduleOutput
  USE moduleBoundary
  IMPLICIT NONE

  !Generic mesh element
  TYPE, PUBLIC, ABSTRACT:: meshElement
    !Index
    INTEGER:: n = 0
    CONTAINS

  END TYPE meshElement

  !Parent of Node element
  TYPE, PUBLIC, ABSTRACT, EXTENDS(meshElement):: meshNode
    !Node volume
    REAL(8):: v = 0.D0
    !Output values
    TYPE(outputNode), ALLOCATABLE:: output(:)
    TYPE(emNode):: emData
    !Lock indicator for scattering
    INTEGER(KIND=OMP_LOCK_KIND):: lock
    CONTAINS
      PROCEDURE(initNode_interface),    DEFERRED, PASS:: init
      PROCEDURE(getCoord_interface),    DEFERRED, PASS:: getCoordinates
      PROCEDURE,                                  PASS:: resetOutput

  END TYPE meshNode

  ABSTRACT INTERFACE
    !Interface of init a node (3D generic coordinates)
    SUBROUTINE initNode_interface(self, n, r)
      IMPORT:: meshNode
      CLASS(meshNode), INTENT(out):: self
      INTEGER, INTENT(in):: n
      REAL(8), INTENT(in):: r(1:3)

    END SUBROUTINE initNode_interface

    !Interface to get coordinates from node
    PURE FUNCTION getCoord_interface(self) RESULT(r)
      IMPORT:: meshNode
      CLASS(meshNode), INTENT(in):: self
      REAL(8):: r(1:3)

    END FUNCTION getCoord_interface

  END INTERFACE

  !Containers for nodes in the mesh
  TYPE:: meshNodeCont
    CLASS(meshNode), ALLOCATABLE:: obj
    CONTAINS

  END TYPE meshNodeCont

  !Type for array of boundary functions (one per species)
  TYPE, PUBLIC:: fBoundaryGeneric
    PROCEDURE(boundary_interface), POINTER, NOPASS:: apply => NULL()
    CONTAINS

  END TYPE

  !Parent of Edge element
  TYPE, PUBLIC, ABSTRACT, EXTENDS(meshElement):: meshEdge
    !Connectivity to vols
    CLASS(meshVol), POINTER:: e1 => NULL(), e2 => NULL()
    !Connectivity to vols in meshColl
    CLASS(meshVol), POINTER:: eColl => NULL()
    !Normal vector
    REAL(8):: normal(1:3)
    !Weight for random injection of particles
    REAL(8):: weight = 1.D0
    !Pointer to boundary type
    TYPE(boundaryCont), POINTER:: boundary
    !Array of functions for boundary conditions
    TYPE(fBoundaryGeneric), ALLOCATABLE:: fBoundary(:)
    !Physical surface for the edge
    INTEGER:: physicalSurface
    CONTAINS
      PROCEDURE(initEdge_interface),         DEFERRED, PASS:: init
      PROCEDURE(getNodesEdge_interface),     DEFERRED, PASS:: getNodes
      PROCEDURE(intersectionEdge_interface), DEFERRED, PASS:: intersection
      PROCEDURE(randPosEdge_interface),      DEFERRED, PASS:: randPos

  END TYPE meshEdge

  ABSTRACT INTERFACE
    !Inits the edge parameters
    SUBROUTINE initEdge_interface(self, n, p, bt, physicalSurface)
      IMPORT:: meshEdge

      CLASS(meshEdge), INTENT(out):: self
      INTEGER, INTENT(in):: n
      INTEGER, INTENT(in):: p(:)
      INTEGER, INTENT(in):: bt
      INTEGER, INTENT(in):: physicalSurface

    END SUBROUTINE initEdge_interface

    !Get nodes index from node
    PURE FUNCTION getNodesEdge_interface(self) RESULT(n)
      IMPORT:: meshEdge
      CLASS(meshEdge), INTENT(in):: self
      INTEGER, ALLOCATABLE:: n(:)

    END FUNCTION getNodesEdge_interface

    !Returns the intersecction between an edge and a line defined by point r0
    PURE FUNCTION intersectionEdge_interface(self, r0) RESULT(r)
      IMPORT:: meshEdge
      CLASS(meshEdge), INTENT(in):: self
      REAL(8), INTENT(in), DIMENSION(1:3):: r0
      REAL(8):: r(1:3)

    END FUNCTION intersectionEdge_interface

    !Returns a random position in the edge
    FUNCTION randPosEdge_interface(self) RESULT(r)
      IMPORT:: meshEdge
      CLASS(meshEdge), INTENT(in):: self
      REAL(8):: r(1:3)

    END FUNCTION randPosEdge_interface

  END INTERFACE

  INTERFACE
    SUBROUTINE boundary_interface(edge, part)
      USE moduleSpecies

      IMPORT:: meshEdge
      CLASS (meshEdge), INTENT(inout):: edge
      CLASS (particle), INTENT(inout):: part

    END SUBROUTINE

  END INTERFACE

  !Containers for edges in the mesh
  TYPE:: meshEdgeCont
    CLASS(meshEdge), ALLOCATABLE:: obj

  END TYPE meshEdgeCont

  !Parent of Volume element
  TYPE, PUBLIC, ABSTRACT, EXTENDS(meshElement):: meshVol
    !Maximum collision rate
    REAL(8):: sigmaVrelMax = 0.D0
    !Volume
    REAL(8):: volume = 0.D0
    !List of particles inside the volume
    TYPE(listNode):: listPart_in
    !Lock indicator for listPart_in
    INTEGER(KIND=OMP_LOCK_KIND):: lock
    !Number of collisions per volume
    INTEGER:: nColl = 0
    !Total weight of particles inside cell
    REAL(8):: totalWeight = 0.D0
    CONTAINS
      PROCEDURE(initVol_interface),     DEFERRED, PASS::   init
      PROCEDURE(getNodesVol_interface), DEFERRED, PASS::   getNodes
      PROCEDURE(randPosVol_interface),  DEFERRED, PASS::   randPos
      PROCEDURE(fPsi_interface),        DEFERRED, NOPASS:: fPsi
      PROCEDURE(scatter_interface),     DEFERRED, PASS::   scatter
      PROCEDURE(gatherEF_interface),    DEFERRED, PASS::   gatherEF
      PROCEDURE(elemK_interface),       DEFERRED, PASS::   elemK
      PROCEDURE(elemF_interface),       DEFERRED, PASS::   elemF
      PROCEDURE,                                  PASS::   findCell
      PROCEDURE(phy2log_interface),     DEFERRED, PASS::   phy2log
      PROCEDURE(inside_interface),      DEFERRED, NOPASS:: inside
      PROCEDURE(nextElement_interface), DEFERRED, PASS::   nextElement

  END TYPE meshVol

  ABSTRACT INTERFACE
    SUBROUTINE initVol_interface(self, n, p, nodes)
      IMPORT:: meshVol
      IMPORT meshNodeCont
      CLASS(meshVol), INTENT(out):: self
      INTEGER, INTENT(in):: n
      INTEGER, INTENT(in):: p(:)
      TYPE(meshNodeCont), INTENT(in), TARGET:: nodes(:)

    END SUBROUTINE initVol_interface

    SUBROUTINE scatter_interface(self, part)
      USE moduleSpecies
      
      IMPORT:: meshVol
      CLASS(meshVol), INTENT(in):: self
      CLASS(particle), INTENT(in):: part

    END SUBROUTINE scatter_interface

    PURE FUNCTION gatherEF_interface(self, xi) RESULT(EF)
      IMPORT:: meshVol
      CLASS(meshVol), INTENT(in):: self
      REAL(8), INTENT(in):: xi(1:3)
      REAL(8):: EF(1:3)

    END FUNCTION gatherEF_interface

    PURE FUNCTION getNodesVol_interface(self) RESULT(n)
      IMPORT:: meshVol
      CLASS(meshVol), INTENT(in):: self
      INTEGER, ALLOCATABLE:: n(:)

    END FUNCTION getNodesVol_interface

    PURE FUNCTION fPsi_interface(xi) RESULT(fPsi)
      REAL(8), INTENT(in):: xi(1:3)
      REAL(8), ALLOCATABLE:: fPsi(:)

    END FUNCTION fPsi_interface

    PURE FUNCTION elemK_interface(self) RESULT(localK)
      IMPORT:: meshVol
      CLASS(meshVol), INTENT(in):: self
      REAL(8), ALLOCATABLE:: localK(:,:)

    END FUNCTION elemK_interface

    PURE FUNCTION elemF_interface(self, source) RESULT(localF)
      IMPORT:: meshVol
      CLASS(meshVol), INTENT(in):: self
      REAL(8), INTENT(in):: source(1:)
      REAL(8), ALLOCATABLE:: localF(:)

    END FUNCTION elemF_interface

    SUBROUTINE nextElement_interface(self, xi, nextElement)
      IMPORT:: meshVol, meshElement
      CLASS(meshVol), INTENT(in):: self
      REAL(8), INTENT(in):: xi(1:3)
      CLASS(meshElement), POINTER, INTENT(out):: nextElement

    END SUBROUTINE nextElement_interface

    PURE FUNCTION phy2log_interface(self,r) RESULT(xN) 
      IMPORT:: meshVol
      CLASS(meshVol), INTENT(in):: self
      REAL(8), INTENT(in):: r(1:3)
      REAL(8):: xN(1:3)

    END FUNCTION phy2log_interface

    PURE FUNCTION inside_interface(xi) RESULT(ins)
      IMPORT:: meshVol
      REAL(8), INTENT(in):: xi(1:3)
      LOGICAL:: ins

    END FUNCTION inside_interface

    FUNCTION randPosVol_interface(self) RESULT(r)
      IMPORT:: meshVol
      CLASS(meshVol), INTENT(in):: self
      REAL(8):: r(1:3)

    END FUNCTION randPosVol_interface

  END INTERFACE

  !Containers for volumes in the mesh
  TYPE:: meshVolCont
    CLASS(meshVol), ALLOCATABLE:: obj

  END TYPE meshVolCont

  !Generic mesh type
  TYPE, ABSTRACT:: meshGeneric
    !Geometry of the mesh
    CHARACTER(:), ALLOCATABLE:: geometry
    !Number of elements
    INTEGER:: numNodes, numVols
    !Array of nodes
    TYPE(meshNodeCont), ALLOCATABLE:: nodes(:)
    !Array of volume elements
    TYPE(meshVolCont),  ALLOCATABLE:: vols(:)
    PROCEDURE(readMesh_interface),     POINTER, PASS::   readMesh     => NULL()
    PROCEDURE(readInitial_interface),  POINTER, NOPASS:: readInitial  => NULL()
    PROCEDURE(connectMesh_interface),  POINTER, PASS::   connectMesh  => NULL()
    PROCEDURE(printColl_interface),    POINTER, PASS::   printColl    => NULL()
    CONTAINS
      PROCEDURE, PASS:: doCollisions

  END TYPE

  ABSTRACT INTERFACE
    !Reads the mesh from a file
    SUBROUTINE readMesh_interface(self, filename)
      IMPORT meshGeneric

      CLASS(meshGeneric), INTENT(inout):: self
      CHARACTER(:), ALLOCATABLE, INTENT(in):: filename

    END SUBROUTINE readMesh_interface

    SUBROUTINE readInitial_interface(sp, filename, density, velocity, temperature)
      INTEGER, INTENT(in):: sp
      CHARACTER(:), ALLOCATABLE, INTENT(in):: filename
      REAL(8), ALLOCATABLE, INTENT(out), DIMENSION(:):: density
      REAL(8), ALLOCATABLE, INTENT(out), DIMENSION(:,:):: velocity
      REAL(8), ALLOCATABLE, INTENT(out), DIMENSION(:):: temperature

    END SUBROUTINE readInitial_interface

    !Connects volume and edges to the mesh
    SUBROUTINE connectMesh_interface(self)
      IMPORT meshGeneric

      CLASS(meshGeneric), INTENT(inout):: self

    END SUBROUTINE connectMesh_interface

    !Prints number of collisions in each volume
    SUBROUTINE printColl_interface(self, t)
      IMPORT meshGeneric

      CLASS(meshGeneric), INTENT(in):: self
      INTEGER, INTENT(in):: t

    END SUBROUTINE printColl_interface

  END INTERFACE

  !Particle mesh
  TYPE, EXTENDS(meshGeneric), PUBLIC:: meshParticles
    INTEGER:: numEdges
    !Array of boundary elements
    TYPE(meshEdgeCont), ALLOCATABLE:: edges(:)
    !Global stiffness matrix
    REAL(8), ALLOCATABLE, DIMENSION(:,:):: K 
    !Permutation matrix for P L U factorization
    INTEGER, ALLOCATABLE, DIMENSION(:,:):: IPIV
    PROCEDURE(printOutput_interface), POINTER, PASS:: printOutput => NULL()
    PROCEDURE(printEM_interface),     POINTER, PASS:: printEM     => NULL()
    PROCEDURE(doCoulomb_interface),   POINTER, PASS:: doCoulomb   => NULL()
    CONTAINS
      PROCEDURE, PASS:: constructGlobalK

  END TYPE meshParticles

  ABSTRACT INTERFACE
    !Perform Coulomb Scattering
    SUBROUTINE doCoulomb_interface(self)
      IMPORT meshParticles

      CLASS(meshParticles), INTENT(inout):: self

    END SUBROUTINE doCoulomb_interface
    
    !Prints Species data
    SUBROUTINE printOutput_interface(self, t)
      IMPORT meshParticles

      CLASS(meshParticles), INTENT(in):: self
      INTEGER, INTENT(in):: t

    END SUBROUTINE printOutput_interface

    !Prints EM info
    SUBROUTINE printEM_interface(self, t)
      IMPORT meshParticles

      CLASS(meshParticles), INTENT(in):: self
      INTEGER, INTENT(in):: t

    END SUBROUTINE printEM_interface


  END INTERFACE

  TYPE(meshParticles), TARGET:: mesh

  !Collision (MCC) mesh
  TYPE, EXTENDS(meshGeneric):: meshCollisions
    CONTAINS

  END TYPE meshCollisions

  TYPE(meshCollisions), TARGET:: meshColl

  ABSTRACT INTERFACE
    SUBROUTINE readMeshColl_interface(self, filename)
      IMPORT meshCollisions

      CLASS(meshCollisions), INTENT(inout):: self
      CHARACTER(:), ALLOCATABLE, INTENT(in):: filename

    END SUBROUTINE readMeshColl_interface

    SUBROUTINE connectMeshColl_interface(self)
      IMPORT meshParticles

      CLASS(meshParticles), INTENT(inout):: self

    END SUBROUTINE connectMeshColl_interface

  END INTERFACE

  !Pointer to mesh used for MC collisions
  CLASS(meshGeneric), POINTER:: meshForMCC => NULL()

  !Procedure to find a volume for a particle in meshColl
  PROCEDURE(findCellColl_interface), POINTER:: findCellColl => NULL()

  ABSTRACT INTERFACE
    SUBROUTINE findCellColl_interface(part)
      USE moduleSpecies

      TYPE(particle), INTENT(inout):: part

    END SUBROUTINE findCellColl_interface

  END INTERFACE

  !Logical to indicate if an specific mesh for MC Collisions is used
  LOGICAL:: doubleMesh
  !Complete path for the two meshes
  CHARACTER(:), ALLOCATABLE:: pathMeshColl, pathMeshParticle

  CONTAINS
    !Constructs the global K matrix
    SUBROUTINE constructGlobalK(self)
      IMPLICIT NONE

      CLASS(meshParticles), INTENT(inout):: self
      INTEGER:: e
      INTEGER, ALLOCATABLE:: n(:)
      REAL(8), ALLOCATABLE:: localK(:,:)
      INTEGER:: nNodes, i, j

      DO e = 1, self%numVols
        n      = self%vols(e)%obj%getNodes()
        localK = self%vols(e)%obj%elemK()
        nNodes = SIZE(n)

        DO i = 1, nNodes
          DO j = 1, nNodes
            self%K(n(i), n(j)) = self%K(n(i), n(j)) + localK(i, j)

          END DO

        END DO

      END DO

    END SUBROUTINE constructGlobalK

    !Reset the output of node
    PURE SUBROUTINE resetOutput(self)
      USE moduleSpecies
      USE moduleOutput
      IMPLICIT NONE

      CLASS(meshNode), INTENT(inout):: self
      INTEGER:: k

      DO k = 1, nSpecies
        self%output(k)%den     = 0.D0
        self%output(k)%mom     = 0.D0
        self%output(k)%tensorS = 0.D0

      END DO

    END SUBROUTINE resetOutput

    !Find next cell for particle
    RECURSIVE SUBROUTINE findCell(self, part, oldCell)
      USE moduleSpecies
      USE moduleErrors
      USE OMP_LIB
      IMPLICIT NONE

      CLASS(meshVol), INTENT(inout):: self
      CLASS(particle), INTENT(inout), TARGET:: part
      CLASS(meshVol), OPTIONAL, INTENT(in):: oldCell
      REAL(8):: xi(1:3)
      CLASS(meshElement), POINTER:: nextElement

      xi = self%phy2log(part%r)
      !Checks if particle is inside 'self' cell
      IF (self%inside(xi)) THEN
        part%vol  = self%n
        part%xi   = xi
        part%n_in = .TRUE.
        !Assign particle to listPart_in
        CALL OMP_SET_LOCK(self%lock)
        CALL self%listPart_in%add(part)
        self%totalWeight = self%totalWeight + part%weight
        CALL OMP_UNSET_LOCK(self%lock)

      ELSE
        !If not, searches for a neighbour and repeats the process.
        CALL self%nextElement(xi, nextElement)
        !Defines the next step
        SELECT TYPE(nextElement)
        CLASS IS(meshVol)
          !Particle moved to new cell, repeat find procedure
          CALL nextElement%findCell(part, self)

        CLASS IS (meshEdge)
          !Particle encountered a surface, apply boundary
          CALL nextElement%fBoundary(part%species%n)%apply(nextElement,part)

          !If particle is still inside the domain, call findCell
          IF (part%n_in) THEN 
            IF(PRESENT(oldCell)) THEN
              CALL self%findCell(part, oldCell)

            ELSE
              CALL self%findCell(part)

            END IF
          END IF

        CLASS DEFAULT
          WRITE (*, "(A, I6)") "Element = ", self%n
          CALL criticalError("No connectivity found for element", "findCell")

        END SELECT

      END IF

    END SUBROUTINE findCell

    !If Coll and Particle are the same, simply copy the part%vol into part%volColl
    SUBROUTINE findCellSameMesh(part)
      USE moduleSpecies
      IMPLICIT NONE

      TYPE(particle), INTENT(inout):: part

      part%volColl = part%vol

    END SUBROUTINE findCellSameMesh

    !TODO: try to combine this with the findCell for a regular mesh
    !Find the volume in which particle reside in the mesh for collisions
    SUBROUTINE findCellCollMesh(part)
      USE moduleSpecies
      IMPLICIT NONE

      TYPE(particle), INTENT(inout):: part
      LOGICAL:: found
      CLASS(meshVol), POINTER:: vol
      REAL(8), DIMENSION(1:3):: xii
      CLASS(meshElement), POINTER:: nextElement

      found = .FALSE.
      
      vol => meshColl%vols(part%volColl)%obj
      DO WHILE(.NOT. found)
        xii = vol%phy2log(part%r)
        IF (vol%inside(xii)) THEN
          part%volColl = vol%n
          CALL OMP_SET_LOCK(vol%lock)
          CALL vol%listPart_in%add(part)
          vol%totalWeight = vol%totalWeight + part%weight
          CALL OMP_UNSET_LOCK(vol%lock)
          found = .TRUE.

        ELSE
          CALL vol%nextElement(xii, nextElement)
          SELECT TYPE(nextElement)
          CLASS IS(meshVol)
            !Try next element
            vol => nextElement

          CLASS DEFAULT
            !Should never happend, but just in case, stops loops
            found = .TRUE.

          END SELECT

        END IF

      END DO

    END SUBROUTINE findCellCollMesh

    !Returns index of volume associated to a position (if any)
    !If no voulme is found, returns 0
    !WARNING: This function is slow and should only be used in initialization phase
    FUNCTION findCellBrute(self, r) RESULT(nVol)
      USE moduleSpecies
      IMPLICIT NONE

      CLASS(meshGeneric), INTENT(in):: self
      REAL(8), DIMENSION(1:3), INTENT(in):: r
      INTEGER:: nVol
      INTEGER:: e
      REAL(8), DIMENSION(1:3):: xii

      !Inits RESULT
      nVol = 0

      DO e = 1, self%numVols
        xii = self%vols(e)%obj%phy2log(r)
        IF(self%vols(e)%obj%inside(xii)) THEN
          nVol = self%vols(e)%obj%n
          EXIT

        END IF

      END DO

    END FUNCTION findCellBrute

    !Computes collisions in element
    SUBROUTINE doCollisions(self)
      USE moduleCollisions
      USE moduleSpecies
      USE moduleList
      use moduleRefParam
      USE moduleRandom
      IMPLICIT NONE

      CLASS(meshGeneric), INTENT(inout), TARGET:: self
      INTEGER:: e
      CLASS(meshVol), POINTER:: vol
      INTEGER:: nPart !Number of particles inside the cell
      REAL(8):: pMax !Maximum probability of collision
      INTEGER:: rnd !random index
      TYPE(particle), POINTER:: part_i, part_j
      INTEGER:: n !collision
      INTEGER:: ij, k
      REAL(8):: sigmaVrelMaxNew
      TYPE(pointerArray), ALLOCATABLE:: partTemp(:)

      !$OMP DO SCHEDULE(DYNAMIC)
      DO e=1, self%numVols
        vol => self%vols(e)%obj
        nPart = vol%listPart_in%amount
        !Computes iterations if there is more than one particle in the cell
        IF (nPart > 1) THEN
          !Probability of collision
          pMax = vol%totalWeight*vol%sigmaVrelMax*tauMin/vol%volume

          !Number of collisions in the cell
          vol%nColl = NINT(REAL(nPart)*pMax*0.5D0)

          IF (vol%nColl > 0) THEN
            !Converts the list of particles to an array for easy access
            partTemp = vol%listPart_in%convert2Array()

          END IF

          DO n = 1, vol%nColl
            !Select random numbers
            rnd = random(1, nPart)
            part_i => partTemp(rnd)%part
            rnd = random(1, nPart)
            part_j => partTemp(rnd)%part
            ij = interactionIndex(part_i%species%n, part_j%species%n)
            sigmaVrelMaxNew = 0.D0
            DO k = 1, interactionMatrix(ij)%amount
              CALL interactionMatrix(ij)%collisions(k)%obj%collide(vol%sigmaVrelMax, sigmaVrelMaxNew, part_i, part_j)

            END DO

            !Update maximum cross section*v_rel per each collision
            IF (sigmaVrelMaxNew > vol%sigmaVrelMax) THEN
              vol%sigmaVrelMax = sigmaVrelMaxNew

            END IF

          END DO

        END IF

      END DO
      !$OMP END DO

    END SUBROUTINE doCollisions

    SUBROUTINE doCoulomb(self)
      IMPLICIT NONE

      CLASS(meshParticles), INTENT(inout):: self

    END SUBROUTINE doCoulomb

END MODULE moduleMesh