!moduleMeshCyl: 2D axial symmetric extension of generic mesh from GMSH format.
!               x == z
!               y == r
!               z == theta (unused)
MODULE moduleMeshCyl
  USE moduleMesh
  IMPLICIT NONE

  !TODO: Move this, this is horrible
  REAL(8), PARAMETER:: w(1:3) = (/ 5.D0/9.D0, 8.D0/9.D0, 5.D0/9.D0 /)
  REAL(8), PARAMETER:: cor(1:3) = (/ -DSQRT(3.D0/5.D0), 0.D0, DSQRT(3.D0/5.D0) /)

  TYPE, PUBLIC, EXTENDS(meshNode):: meshNodeCyl
    !Element coordinates
    REAL(8):: r = 0.D0, z = 0.D0
    CONTAINS
      PROCEDURE, PASS:: init => initNodeCyl
      PROCEDURE, PASS:: getCoordinates => getCoordCyl

  END TYPE meshNodeCyl

  TYPE, PUBLIC, ABSTRACT, EXTENDS(meshEdge):: meshEdgeCyl
    !Element coordinates
    REAL(8):: r(1:2) = 0.D0, z(1:2) = 0.D0
    !Connectivity to nodes
    CLASS(meshNode), POINTER:: n1 => NULL(), n2 => NULL()
    CONTAINS
      PROCEDURE, PASS::                               init      => initEdgeCyl
      PROCEDURE(boundary_interface), PASS, DEFERRED:: fBoundary

  END TYPE meshEdgeCyl

  ABSTRACT INTERFACE
    SUBROUTINE boundary_interface(self, part)
      USE moduleSpecies

      IMPORT:: meshEdgeCyl
      CLASS (meshEdgeCyl), INTENT(inout):: self
      CLASS (particle), INTENT(inout):: part

    END SUBROUTINE

  END INTERFACE

  TYPE, PUBLIC, ABSTRACT, EXTENDS(meshVol):: meshVolCyl
    CONTAINS
      PROCEDURE, PASS:: collision => collision2DCyl

  END TYPE meshVolCyl

  TYPE, PUBLIC, EXTENDS(meshVolCyl):: meshVolCylQuad
    !Element coordinates
    REAL(8):: r(1:4) = 0.D0, z(1:4) = 0.D0
    !Connectivity to nodes
    CLASS(meshNode), POINTER:: n1 => NULL(), n2 => NULL(), n3 => NULL(), n4 => NULL()
    !Connectivity to adjacent elements
    CLASS(*), POINTER:: e1 => NULL(), e2 => NULL(), e3 => NULL(), e4 => NULL()
    !Maximum collision rate
    !TODO: Move to other more appropiate section
    REAL(8):: phi(1:4) = 0.D0
    REAL(8):: Ke(1:4,1:4) = 0.D0
    REAL(8):: arNodes(1:4) = 0.D0

    CONTAINS
      PROCEDURE, PASS::   init        => initVolQuadCyl
      PROCEDURE, PASS::   locKe       => localKeQuad
      PROCEDURE, PASS::   detJac      => detJQuad
      PROCEDURE, PASS::   invJ        => invJQuad
      PROCEDURE, PASS::   area        => areaQuad
      PROCEDURE, NOPASS:: weight      => weightQuad
      PROCEDURE, NOPASS:: inside      => insideQuad
      PROCEDURE, PASS::   dVal        => dValQuad
      PROCEDURE, PASS::   scatter     => scatterQuad
      PROCEDURE, PASS::   phy2log     => phy2logQuad
      PROCEDURE, PASS::   findCell    => findCellCylQuad
      PROCEDURE, PASS::   resetOutput => resetOutputQuad

  END TYPE meshVolCylQuad

  CONTAINS

    !Inits node element
    SUBROUTINE initNodeCyl(self, n, r)
      USE moduleSpecies
      USE moduleRefParam
      IMPLICIT NONE

      CLASS(meshNodeCyl), INTENT(out):: self
      INTEGER, INTENT(in):: n
      REAL(8), INTENT(in):: r(1:3)

      self%n = n
      self%r = r(1)/L_ref
      self%z = r(2)/L_ref
      !Node volume, to be determined in mesh
      self%v = 0.D0

      !Allocates output:
      ALLOCATE(self%output(1:nSpecies))

    END SUBROUTINE initNodeCyl

    FUNCTION getCoordCyl(self) RESULT(r)
      IMPLICIT NONE

      CLASS(meshNodeCyl):: self
      REAL(8):: r(1:3)

      r = (/self%z, self%r, 0.D0/)

    END FUNCTION getCoordCyl

    !Edge functions
    SUBROUTINE initEdgeCyl(self, n, p, bt, physicalSurface)
      IMPLICIT NONE

      CLASS(meshEdgeCyl), INTENT(out):: self
      INTEGER, INTENT(in):: n
      INTEGER, INTENT(in):: p(:)
      INTEGER, INTENT(in):: bt
      INTEGER, INTENT(in):: physicalSurface
      REAL(8):: r1(1:3), r2(1:3)

      self%n  = n
      self%n1 => mesh%nodes(p(1))%obj
      self%n2 => mesh%nodes(p(2))%obj
      !Get element coordinates
      r1 = self%n1%getCoordinates()
      r2 = self%n2%getCoordinates()
      self%z  = (/r1(1), r2(1)/)
      self%r  = (/r1(2), r2(2)/)
      !Normal vector
      self%normal = (/ self%r(2)-self%r(1), &
                       self%z(2)-self%z(1), &
                       0.D0 /)
      !Boundary index
      self%bt = bt
      !Phyiscal Surface
      self%physicalSurface = physicalSurface

    END SUBROUTINE initEdgeCyl

    !Vol functions
    !Quad Volume
    SUBROUTINE initVolQuadCyl(self, n, p)
      USE moduleRefParam
      IMPLICIT NONE

      CLASS(meshVolCylQuad), INTENT(out):: self
      INTEGER, INTENT(in):: n
      INTEGER, INTENT(in):: p(:)
      REAL(8):: r1(1:3), r2(1:3), r3(1:3), r4(1:3)

      self%n  = n
      self%n1 => mesh%nodes(p(1))%obj
      self%n2 => mesh%nodes(p(2))%obj
      self%n3 => mesh%nodes(p(3))%obj
      self%n4 => mesh%nodes(p(4))%obj
      !Get element coordinates
      r1 = self%n1%getCoordinates()
      r2 = self%n2%getCoordinates()
      r3 = self%n3%getCoordinates()
      r4 = self%n4%getCoordinates()
      self%z  = (/r1(1), r2(1), r3(1), r4(1)/)
      self%r  = (/r1(2), r2(2), r3(2), r4(2)/)

      !Assign node volume
      CALL self%area()
      self%n1%v = self%n1%v + self%arNodes(1)
      self%n2%v = self%n2%v + self%arNodes(2)
      self%n3%v = self%n3%v + self%arNodes(3)
      self%n4%v = self%n4%v + self%arNodes(4)

      CALL self%locKe()

      self%sigmaVrelMax = sigma_ref/L_ref**2

    END SUBROUTINE initVolQuadCyl

    FUNCTION fpsi(xi1,xi2) RESULT(psi) 
      IMPLICIT NONE

      REAL(8),INTENT(in):: xi1,xi2
      REAL(8):: psi(1:4)
      psi(1) = (1.D0-xi1)*(1.D0-xi2)
      psi(2) = (1.D0+xi1)*(1.D0-xi2)
      psi(3) = (1.D0+xi1)*(1.D0+xi2)
      psi(4) = (1.D0-xi1)*(1.D0+xi2)
      psi = psi*0.25D0

    END FUNCTION fpsi

    !Derivative element function (xi1)
    FUNCTION dpsiXi1(xi2) RESULT(psiXi1) 
      IMPLICIT NONE

      REAL(8),INTENT(in):: xi2
      REAL(8):: psiXi1(1:4)
      psiXi1(1) = -(1.D0-xi2)
      psiXi1(2) =  (1.D0-xi2)
      psiXi1(3) =  (1.D0+xi2)
      psiXi1(4) = -(1.D0+xi2)
      psiXi1 = psiXi1*0.25D0
      
    END FUNCTION dpsiXi1

    !Derivative element function (xi2)
    FUNCTION dpsiXi2(xi1) RESULT(psiXi2) 
      IMPLICIT NONE

      REAL(8),INTENT(in):: xi1
      REAL(8):: psiXi2(1:4)
      psiXi2(1) = -(1.D0-xi1)
      psiXi2(2) = -(1.D0+xi1)
      psiXi2(3) =  (1.D0+xi1)
      psiXi2(4) =  (1.D0-xi1)
      psiXi2 = psiXi2*0.25D0

    END FUNCTION dpsiXi2

    !Transforms physical coordinates to element coordinates
    FUNCTION phy2logQuad(this,r) RESULT(xN) 
      IMPLICIT NONE

      CLASS(meshVolCylQuad):: this
      REAL(8):: r(1:2)
      REAL(8):: xN(1:2), xO(1:2), dPsi(1:2,1:4), detJ, j(1:2,1:2), psi(1:4), f(1:2)
      REAL(8):: conv

      !Iterative newton method to transform coordinates
      conv=1.D0
      xO=0.D0

      DO WHILE(conv>1.D-4)
        dPsi(1,:) = dpsiXi1(xO(2))
        dPsi(2,:) = dpsiXi2(xO(1))
        j    = this%invJ(xO(1),xO(2),dPsi)
        psi = fpsi(xO(1), xO(2))
        f(1) = DOT_PRODUCT(psi,this%z)-r(1)
        f(2) = DOT_PRODUCT(psi,this%r)-r(2)
        detJ = this%detJac(xO(1),xO(2),dPsi)
        xN=xO - MATMUL(j, f)/detJ
        conv=MAXVAL(DABS(xN-xO),1)
        xO=xN

      END DO

    END FUNCTION phy2logQuad

    !Computes element local stiffness matrix
    SUBROUTINE localKeQuad(this) 
      USE moduleConstParam
      IMPLICIT NONE

      CLASS(meshVolCylQuad):: this
      REAL(8):: r, xi1, xi2, dPsi(1:2,1:4)
      REAL(8):: j(1:2,1:2), detJ
      INTEGER:: l, m

      this%Ke=0.D0
      !Start 2D Gauss Quad Integral
      DO l=1, 3
        DO m = 1, 3
          xi1 = cor(l)
          xi2 = cor(m)
          dPsi(1,:) = dpsiXi1(xi2)
          dPsi(2,:) = dpsiXi2(xi1)
          detJ = this%detJac(xi1,xi2,dPsi)
          j    = this%invJ(xi1,xi2,dPsi)
          r    = DOT_PRODUCT(fpsi(xi1,xi2),this%r)
          this%Ke = this%Ke + MATMUL(TRANSPOSE(MATMUL(j,dPsi)),MATMUL(j,dPsi))*r*w(l)*w(m)/detJ

        END DO
      END DO
      this%Ke = this%Ke*2.D0*PI

    END SUBROUTINE localKeQuad

    !Computes element Jacobian determinant
    FUNCTION detJQuad(this,xi1,xi2,dPsi_in) RESULT(dJ) 
      IMPLICIT NONE

      REAL(8),OPTIONAL:: dPsi_in(1:2,1:4)
      REAL(8):: dPsi(1:2,1:4)
      REAL(8):: dz(1:2), dr(1:2)
      REAL(8):: xi1,xi2, dJ
      CLASS(meshVolCylQuad):: this

      IF(PRESENT(dPsi_in)) THEN
        dPsi=dPsi_in

      ELSE
        dPsi(1,:) = dpsiXi1(xi2)
        dPsi(2,:) = dpsiXi2(xi1)

      END IF

      dz(1) = DOT_PRODUCT(dPsi(1,:),this%z)
      dz(2) = DOT_PRODUCT(dPsi(2,:),this%z)
      dr(1) = DOT_PRODUCT(dPsi(1,:),this%r)
      dr(2) = DOT_PRODUCT(dPsi(2,:),this%r)
      dJ = dz(1)*dr(2)-dz(2)*dr(1)

    END FUNCTION detJQuad

    FUNCTION invJQuad(this,xi1,xi2,dPsi_in) RESULT(j) !Computes element Jacobian inverse matrix (without determinant)
      IMPLICIT NONE

      REAL(8),OPTIONAL:: dpsi_in(1:2,1:4)
      REAL(8):: dPsi(1:2,1:4)
      REAL(8):: dz(1:2), dr(1:2)
      REAL(8):: xi1,xi2, j(1:2,1:2)
      CLASS(meshVolCylQuad):: this

      IF(PRESENT(dPsi_in)) THEN
        dPsi=dPsi_in

      ELSE
        dPsi(1,:) = dpsiXi1(xi2)
        dPsi(2,:) = dpsiXi2(xi1)

      END IF

      dz(1) = DOT_PRODUCT(dPsi(1,:),this%z)
      dz(2) = DOT_PRODUCT(dPsi(2,:),this%z)
      dr(1) = DOT_PRODUCT(dPsi(1,:),this%r)
      dr(2) = DOT_PRODUCT(dPsi(2,:),this%r)
      j(1,:) = (/  dr(2), -dz(2) /)
      j(2,:) = (/ -dr(1),  dz(1) /)

    END FUNCTION invJQuad

    SUBROUTINE areaQuad(this)!Computes element area
      USE moduleConstParam
      IMPLICIT NONE

      CLASS(meshVolCylQuad):: this
      REAL(8):: r, xi1, xi2
      REAL(8):: detJ, fpsi0(1:4)

      this%volume  = 0.D0
      this%arNodes = 0.D0
      !2D 1 point Gauss Quad Integral
      xi1 = 0.D0
      xi2 = 0.D0
      detJ = this%detJac(xi1,xi2)*8.D0*PI !4*2*pi
      fpsi0 = fpsi(xi1,xi2)
      r = DOT_PRODUCT(fpsi0,this%r)
      this%volume  =       r*detJ
      this%arNodes = fpsi0*r*detJ

    END SUBROUTINE areaQuad

    FUNCTION weightQuad(xi1_p, xi2_p) RESULT(w) !Computes weights in the four vertices. '_p' references particle position in the logical space
      IMPLICIT NONE

      REAL(8):: xi1_p, xi2_p
      REAL(8):: w(1:4)

      w = fpsi(xi1_p, xi2_p)

    END FUNCTION weightQuad

    FUNCTION insideQuad(xi1_p, xi2_p) RESULT(ins) !Checks if a particle is inside a quad element
      IMPLICIT NONE

      REAL(8):: xi1_p, xi2_p
      LOGICAL:: ins

      ins = (xi1_p >= -1.D0 .AND. xi1_p <= 1.D0) .AND. &
            (xi2_p >= -1.D0 .AND. xi2_p <= 1.D0)

    END FUNCTION insideQuad

    FUNCTION dValQuad(this,xi1,xi2) RESULT(EF)!Computes electric field in xi1,xi2
      IMPLICIT NONE

      CLASS(meshVolCylQuad):: this
      REAL(8):: xi1, xi2, dPsi(1:2,1:4)
      REAL(8):: j(1:2,1:2), detJ
      REAL(8):: EF(1:3)

      dPsi(1,:) = dpsiXi1(xi2)
      dPsi(2,:) = dpsiXi2(xi1)
      detJ = this%detJac(xi1,xi2,dPsi)
      j    = this%invJ(xi1,xi2,dPsi)
      EF(1) = -DOT_PRODUCT(dPsi(1,:), this%phi)*j(2,2)/detJ
      EF(2) = -DOT_PRODUCT(dPsi(2,:), this%phi)*j(1,1)/detJ
      EF(3) =  0.D0

    END FUNCTION dValQuad

    SUBROUTINE scatterQuad(self, part)
      USE moduleOutput
      USE moduleSpecies
      IMPLICIT NONE

      CLASS(meshVolCylQuad), INTENT(in):: self
      CLASS(particle), INTENT(in):: part
      TYPE(outputNode), POINTER:: vertex
      REAL(8):: w_p(1:4)
      REAL(8):: tensorS(1:3,1:3)

      w_p = self%weight(part%xLog(1), part%xLog(2))
      tensorS = outerProduct(part%v, part%v)

      vertex => self%n1%output(part%pt)
      vertex%den          = vertex%den          + part%weight*w_p(1)
      vertex%mom(:)       = vertex%mom(:)       + part%weight*w_p(1)*part%v(:)
      vertex%tensorS(:,:) = vertex%tensorS(:,:) + part%weight*w_p(1)*tensorS

      vertex => self%n2%output(part%pt)
      vertex%den          = vertex%den          + part%weight*w_p(2)
      vertex%mom(:)       = vertex%mom(:)       + part%weight*w_p(2)*part%v(:)
      vertex%tensorS(:,:) = vertex%tensorS(:,:) + part%weight*w_p(2)*tensorS

      vertex => self%n3%output(part%pt)
      vertex%den          = vertex%den          + part%weight*w_p(3)
      vertex%mom(:)       = vertex%mom(:)       + part%weight*w_p(3)*part%v(:)
      vertex%tensorS(:,:) = vertex%tensorS(:,:) + part%weight*w_p(3)*tensorS

      vertex => self%n4%output(part%pt)
      vertex%den          = vertex%den          + part%weight*w_p(4)
      vertex%mom(:)       = vertex%mom(:)       + part%weight*w_p(4)*part%v(:)
      vertex%tensorS(:,:) = vertex%tensorS(:,:) + part%weight*w_p(4)*tensorS


    END SUBROUTINE scatterQuad

    RECURSIVE SUBROUTINE findCellCylQuad(self, part, oldCell)
      USE moduleSpecies
      IMPLICIT NONE

      CLASS(meshVolCylQuad), INTENT(inout):: self
      CLASS(meshVol), OPTIONAL, INTENT(in):: oldCell
      CLASS(particle), INTENT(inout):: part
      REAL(8):: xLog(1:2)
      REAL(8):: xLogArray(1:4)
      CLASS(*), POINTER:: nextElement
      INTEGER:: nextInt

      xLog = self%phy2log(part%r(1:2))
      IF (self%inside(xLog(1), xLog(2))) THEN
        !Checks if particle is inside of current cell
        IF (PRESENT(oldCell)) THEN
          !If oldCell, recalculate particle weight, as particle has entered a new cell.
          part%weight = part%weight*oldCell%volume/self%volume
        END IF
        part%e_p  = self%n
        part%xLog = xLog
      ELSE
        !If not, searches for a neighbour and repeats the process.
        xLogArray = (/ -xLog(2), xLog(1), xLog(2), -xLog(1) /)
        nextInt = MAXLOC(xLogArray,1)
        !Selects the higher value of directions and searches in that direction
        NULLIFY(nextElement)
        SELECT CASE (nextInt)
        CASE (1)
          nextElement => self%e1
        CASE (2)
          nextElement => self%e2
        CASE (3)
          nextElement => self%e3
        CASE (4)
          nextElement => self%e4
        END SELECT

        !Defines the next step
        SELECT TYPE(nextElement)
        CLASS IS(meshVolCyl)
          !Particle moved to new cell, repeat find procedure
          CALL nextElement%findCell(part, self)

        CLASS IS (meshEdgeCyl)
          !Particle encountered an edge, execute boundary
          CALL nextElement%fBoundary(part)

        CLASS DEFAULT
          WRITE(*,*) "ERROR, CHECK findCellCylQuad"

        END SELECT
      END IF

    END SUBROUTINE findCellCylQuad

    PURE SUBROUTINE resetOutputQuad(self)
      USE moduleSpecies
      USE moduleOutput
      IMPLICIT NONE

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

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

        self%n2%output(k)%den     = 0.D0
        self%n2%output(k)%mom     = 0.D0
        self%n2%output(k)%tensorS = 0.D0

        self%n3%output(k)%den     = 0.D0
        self%n3%output(k)%mom     = 0.D0
        self%n3%output(k)%tensorS = 0.D0

        self%n4%output(k)%den     = 0.D0
        self%n4%output(k)%mom     = 0.D0
        self%n4%output(k)%tensorS = 0.D0

      END DO

    END SUBROUTINE resetOutputQuad


    SUBROUTINE collision2DCyl(self)
      USE moduleRefParam
      USE moduleConstParam
      USE moduleCollisions
      USE moduleSpecies
      USE moduleList
      IMPLICIT NONE

      CLASS(meshVolCyl), INTENT(inout):: self
      INTEGER:: Npart !Number of particles inside the cell
      REAL(8):: Fn !Specific weight
      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

      Fn = species(1)%obj%weight!TODO: Check how to do this for multiple species
      Npart = self%listPart_in%amount
      Pmax = Fn*self%sigmaVrelMax*tau/self%volume
      self%nColl = INT(REAL(Npart*(Npart-1))*Pmax*0.5D0)

      DO n = 1, self%NColl
        !Select random numbers
        rnd = 1 + FLOOR(Npart*RAND())
        part_i => self%listPart_in%get(rnd)
        rnd = 1 + FLOOR(Npart*RAND())
        part_j => self%listPart_in%get(rnd)
        ij = interactionIndex(part_i%pt, part_j%pt)
        sigmaVrelMaxNew = 0.D0
        DO k = 1, interactionMatrix(ij)%amount
          CALL interactionMatrix(ij)%collisions(k)%obj%collide(self%sigmaVrelMax, sigmaVrelMaxNew, part_i, part_j)

        END DO

        !Update maximum cross section times vrelative per each collision
        IF (sigmaVrelMaxNew > self%sigmaVrelMax) self%sigmaVrelMax = sigmaVrelMaxNew
          
      END DO

      !Reset output in nodes
      CALL self%resetOutput()

      !Erase the list of particles inside the cell
      CALL self%listPart_in%erase()

    END SUBROUTINE collision2DCyl

END MODULE moduleMeshCyl