!------------------------------------------------------------------------------
! EPFL/Swiss Plasma Center
!------------------------------------------------------------------------------
!
! MODULE: beam
!
!> @author
!> Patryk Kaminski   EPFL/SPC
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> Module responsible for initializing the magnetic field and solve the Poisson equation
!------------------------------------------------------------------------------
MODULE fields
USE constants
USE basic, ONLY: nr, nz, zgrid, rgrid, Br, Bz, Er, Ez, femorder, ngauss, nlppform, pot, Athet, splrz, nlperiod, phinorm, nlPhis, nrank, mpirank, mpisize, step, it2d
USE beam, ONLY: parts
USE bsplines
USE mumps_bsplines
USE mpi
Use omp_lib
Use mpihelper, ONLY: db_type
IMPLICIT NONE

REAL(kind=db), allocatable, SAVE :: matcoef(:,:), phi_spline(:), vec1(:), vec2(:)
REAL(kind=db), allocatable, SAVE :: loc_moments(:,:)
INTEGER, SAVE:: loc_zspan
TYPE(mumps_mat), SAVE :: femat                  ! Finite Element Method matrix


CONTAINS

!---------------------------------------------------------------------------  
!> @author
!> Patryk kaminski   EPFL/SPC 
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> 
!> @brief
!> Set-up the necessary variables for solving Poisson and computes the magnetic field on the grid
! 
!---------------------------------------------------------------------------
SUBROUTINE fields_init
    USE basic, ONLY: pot, nlperiod, nrank, rhs, moments, Zbounds, volume
    USE bsplines
    USE mumps_bsplines
    USE mpihelper
    INTEGER :: nrz(2)

    ! Set up 2d spline splrz used in the FEM
    CALL set_spline(femorder, ngauss, zgrid, rgrid, splrz, nlppform=nlppform, period=nlperiod)

    ! Calculate dimension of splines
      nrz(1)=nz
      nrz(2)=nr 
    CALL get_dim(splrz, nrank, nrz, femorder)

    ALLOCATE(matcoef(nrank(1),nrank(2)))
    ALLOCATE(pot((nr+1)*(nz+1)))
    ALLOCATE(rhs(nrank(1)*nrank(2)))
    IF (mpirank .eq. 0) THEN
      ALLOCATE(moments(10,nrank(1)*nrank(2)))
    ELSE
      ALLOCATE(moments(0,0))
    END IF
    ALLOCATE(phi_spline(nrank(1)*nrank(2))) 
    ALLOCATE(volume(nrank(1)*nrank(2))) 

    ALLOCATE(Br((nr+1)*(nz+1)),Bz((nr+1)*(nz+1)))
    ALLOCATE(Athet((nr+1)*(nz+1)))
    ALLOCATE(Er((nr+1)*(nz+1)),Ez((nr+1)*(nz+1)))
    ! Calculate magnetic field mirror components in grid points (Davidson analytical formula employed)
    CALL magnet
  
  ! Calculate and factorise FEM matrix (depends only on mesh)
    CALL fematrix
    CALL factor(femat)
    CALL comp_volume
    loc_zspan=Zbounds(mpirank+1)-Zbounds(mpirank)+femorder(1)
    ALLOCATE(loc_moments(10,loc_zspan*nrank(2)))
    IF(mpisize .gt. 1) THEN
      CALL init_overlaps(nrank, femorder, Zbounds(mpirank), Zbounds(mpirank+1))
    END IF

END SUBROUTINE fields_init

!---------------------------------------------------------------------------  
!> @author
!> Patryk kaminski   EPFL/SPC 
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> 
!> @brief
!> Construct the right hand side vector used in the FEM Poisson solver
! 
!---------------------------------------------------------------------------
SUBROUTINE rhscon(p)
    USE bsplines
    USE mumps_bsplines
    USE mpi
    USE basic, ONLY: rnorm, potinn, potout, nrank, rhs, moments, nlend, step, it2d
    USE beam, ONLY: particles
    USE mpihelper
    type(particles), INTENT(INOUT):: p
    REAL(kind=db) :: chargecoeff
    
    loc_moments=0 ! Reset the moments matrix
    rhs=0     ! reset rhs
    moments=0 ! reset moments
    !chargecoeff=0.5/pi*omegap/omegac*V/real(nplasma) ! Normalized charge density simulated by each macro particle
    chargecoeff=p%weight*p%q/(2*pi*eps_0*phinorm*rnorm) ! Normalized charge density simulated by each macro particle

! Assemble rhs
    IF (p%Nptot .ne. 0) THEN
         IF(modulo(step,it2d).eq. 0 .or. nlend) THEN
            CALL deposit_moments(p, loc_moments)
         ELSE
            CALL deposit_charge(p, loc_moments)
         END IF
    END IF

  ! If we are using MPI parallelism, reduce the rhs on the root process
    IF(mpisize .gt. 1) THEN
      CALL fields_comm
    ELSE
      moments=loc_moments
    END IF
    
    IF(mpirank.eq.0 )THEN
      IF(nlphis) THEN
        ! We consider self-consistent interactions
         rhs=moments(1,:)*chargecoeff ! Normalized charge density simulated by each macro particle
      ELSE
         rhs=0
      END IF
    END IF

  ! Apply the Dirichlet boundary conditions on the coaxial insert (if present)
    IF(rgrid(0).ne.0.0_db) rhs(1:nrank(1))=potinn
  ! Apply the Dirichlet boundary conditions on the cylinder walls
    rhs((nrank(2)-1)*nrank(1)+1:nrank(2)*nrank(1))=potout
    
  END SUBROUTINE rhscon

!---------------------------------------------------------------------------  
!> @author
!> Patryk kaminski   EPFL/SPC 
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> 
!> @brief
!> Deposit the particles moments (q,v,v^2) from p on the grid
! 
!---------------------------------------------------------------------------

  SUBROUTINE deposit_moments(p, p_loc_moments)
    USE bsplines
    USE mumps_bsplines
    USE mpi
    USE basic, ONLY: Zbounds
    USE beam, ONLY: particles
    USE mpihelper
    TYPE(particles), INTENT(IN):: p
    REAL(kind=db), DIMENSION(:,:), INTENT(INOUT):: p_loc_moments
    INTEGER ::irow, jcol, it, jw, mu, i, k, iend, nbunch
    INTEGER,DIMENSION(:),ALLOCATABLE::zleft, rleft
    REAL(kind=db) :: vr, vthet, vz, coeff
    REAL(kind=db), ALLOCATABLE :: fun(:,:,:), fun2(:,:,:)
    INTEGER:: num_threads

    num_threads=omp_get_max_threads()
    nbunch=p%Nptot/num_threads ! Particle bunch size used when calling basfun
    nbunch=max(nbunch,1) ! Particle bunch size used when calling basfun
    nbunch=min(nbunch,64) ! Particle bunch size used when calling basfun

    ALLOCATE(zleft(nbunch),rleft(nbunch))
    ALLOCATE(fun(1:femorder(1)+1,0:0,nbunch),fun2(1:femorder(2)+1,0:0,nbunch)) ! Arrays keeping values of b-splines at gauss node

! Assemble rhs
    IF (p%Nptot .ne. 0) THEN
!$OMP PARALLEL DO DEFAULT(SHARED), PRIVATE(zleft, rleft, jw, it, iend, irow, jcol, mu, k, vr, vz, vthet, coeff, fun, fun2)
      DO i=1,p%Nptot,nbunch
        ! Avoid segmentation fault by accessing non relevant data
         iend=min(i+nbunch-1,p%Nptot)
         k=iend-i+1
        ! Localize the particle
         !CALL locintv(splrz%sp2, p%R(i:iend), rleft(1:k))
         !CALL locintv(splrz%sp1, p%Z(i:iend), zleft(1:k))
          rleft(1:k)=p%rindex(i:iend)
          zleft(1:k)=p%zindex(i:iend)
        ! Compute the value of the splines at the particles positions 
         CALL basfun(p%Z(i:iend), splrz%sp1, fun(:,:,1:k), zleft(1:k)+1)
         CALL basfun(p%R(i:iend), splrz%sp2, fun2(:,:,1:k), rleft(1:k)+1)
         DO k=1,(iend-i+1)
          DO jw=1,(femorder(2)+1)
           DO it=1,(femorder(1)+1)
               irow=zleft(k)+it-Zbounds(mpirank)
               jcol=rleft(k)+jw   
               mu=irow+(jcol-1)*(loc_zspan)
               coeff=fun(it,0,k)*fun2(jw,0,k)
              ! Add contribution of particle nbunch to rhs grid point mu 
               vr=parts%UR(i+k-1)/parts%Gamma(i+k-1)
               vz=parts%UZ(i+k-1)/parts%Gamma(i+k-1)
               vthet=parts%UTHET(i+k-1)/parts%Gamma(i+k-1)
               !$OMP ATOMIC UPDATE 
               p_loc_moments(1, mu) = p_loc_moments(1, mu) + coeff
               !$OMP END ATOMIC 
               !$OMP ATOMIC UPDATE
               p_loc_moments(2, mu) = p_loc_moments(2, mu) + coeff*vr
               !$OMP END ATOMIC 
               !$OMP ATOMIC UPDATE
               p_loc_moments(3, mu) = p_loc_moments(3, mu) + coeff*vthet
               !$OMP END ATOMIC 
               !$OMP ATOMIC UPDATE
               p_loc_moments(4, mu) = p_loc_moments(4, mu) + coeff*vz
               !$OMP END ATOMIC 
               !$OMP ATOMIC UPDATE
               p_loc_moments(5, mu) = p_loc_moments(5, mu) + coeff*vr*vr
               !$OMP END ATOMIC 
               !$OMP ATOMIC UPDATE
               p_loc_moments(6, mu) = p_loc_moments(6, mu) + coeff*vr*vthet
               !$OMP END ATOMIC 
               !$OMP ATOMIC UPDATE
               p_loc_moments(7, mu) = p_loc_moments(7, mu) + coeff*vr*vz
               !$OMP END ATOMIC 
               !$OMP ATOMIC UPDATE
               p_loc_moments(8, mu) = p_loc_moments(8, mu) + coeff*vthet*vthet
               !$OMP END ATOMIC 
               !$OMP ATOMIC UPDATE
               p_loc_moments(9, mu) = p_loc_moments(9, mu) + coeff*vthet*vz
               !$OMP END ATOMIC 
               !$OMP ATOMIC UPDATE
               p_loc_moments(10,mu) = p_loc_moments(10,mu) + coeff*vz*vz
               !$OMP END ATOMIC 
            END DO
          END DO
         END DO
      END DO
!$OMP END PARALLEL DO
    END IF
    DEALLOCATE(fun,fun2, zleft, rleft)

  END subroutine deposit_moments

!---------------------------------------------------------------------------  
!> @author
!> Patryk kaminski   EPFL/SPC 
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> 
!> @brief
!> Deposit the particles charges (q) from p on the grid
! 
!---------------------------------------------------------------------------


  SUBROUTINE deposit_charge(p, p_loc_moments)
    USE bsplines
    USE mumps_bsplines
    USE mpi
    USE basic, ONLY: Zbounds
    USE beam, ONLY: particles
    USE mpihelper
    TYPE(particles), INTENT(IN):: p
    REAL(kind=db), DIMENSION(:,:), INTENT(INOUT):: p_loc_moments
    INTEGER ::irow, jcol, it, jw, mu, i, k, iend, nbunch
    INTEGER,DIMENSION(:),ALLOCATABLE::zleft, rleft
    REAL(kind=db), ALLOCATABLE :: fun(:,:,:), fun2(:,:,:)
    INTEGER:: num_threads

    num_threads=omp_get_max_threads()
    nbunch=p%Nptot/num_threads ! Particle bunch size used when calling basfun
    nbunch=max(nbunch,1) ! Particle bunch size used when calling basfun
    nbunch=min(nbunch,64) ! Particle bunch size used when calling basfun
    ALLOCATE(zleft(nbunch),rleft(nbunch))
    ALLOCATE(fun(1:femorder(1)+1,0:0,nbunch),fun2(1:femorder(2)+1,0:0,nbunch)) ! Arrays keeping values of b-splines at gauss node

! Assemble rhs
    IF (p%Nptot .ne. 0) THEN
!$OMP PARALLEL DO DEFAULT(SHARED), PRIVATE(zleft, rleft, jw, it, iend, irow, jcol, mu, k, fun, fun2)
      DO i=1,p%Nptot,nbunch
        ! Avoid segmentation fault by accessing non relevant data
         iend=min(i+nbunch-1,p%Nptot)
         k=iend-i+1
        ! Localize the particle
         !CALL locintv(splrz%sp2, p%R(i:iend), rleft(1:k))
         !CALL locintv(splrz%sp1, p%Z(i:iend), zleft(1:k))
         rleft(1:k)=p%rindex(i:iend)
         zleft(1:k)=p%zindex(i:iend)

        ! Compute the value of the splines at the particles positions 
         CALL basfun(p%Z(i:iend), splrz%sp1, fun, zleft(1:k)+1)
         CALL basfun(p%R(i:iend), splrz%sp2, fun2, rleft(1:k)+1)
           DO k=1,(iend-i+1)
            DO jw=1,(femorder(2)+1)
             DO it=1,(femorder(1)+1)
                irow=zleft(k)+it-Zbounds(mpirank)
                jcol=rleft(k)+jw   
                mu=irow+(jcol-1)*(loc_zspan)
                ! Add contribution of particle nbunch to rhs grid point mu 
                !$OMP ATOMIC update
                 p_loc_moments(1, mu) = p_loc_moments(1, mu) + fun(it,0,k)*fun2(jw,0,k)
                !$OMP END ATOMIC
              END DO
            END DO
           END DO
      END DO
!$OMP END PARALLEL DO
    END IF
    
    DEALLOCATE(fun,fun2, zleft, rleft)

  END subroutine deposit_charge

SUBROUTINE fields_comm
USE mpihelper
USE Basic, ONLY: moments, Zbounds, mpirank, nlend, it2d, step, leftproc, rightproc
INTEGER:: ierr, i, j
INTEGER:: displs(mpisize), counts(mpisize)
INTEGER:: overlap_type
INTEGER:: rcvoverlap_type


    displs=Zbounds(0:mpisize-1)
    counts=Zbounds(1:mpisize)-Zbounds(0:mpisize-1)
    counts(mpisize)=counts(mpisize)+femorder(1)

    IF(modulo(step,it2d).eq. 0 .or. nlend) THEN
        CALL start_momentscomm(mpirank, leftproc, rightproc, loc_moments, nrank, femorder, loc_zspan-femorder(1))
        IF(mpirank .gt. 0) THEN
          !$OMP PARALLEL DO SIMD DEFAULT(SHARED)
          DO j=1,femorder(1)
            DO i=1,nrank(2)
              loc_moments(1:10,(i-1)*loc_zspan+j) = loc_moments(1:10,(i-1)*loc_zspan+j)&
              & + momentsoverlap_buffer(10*(nrank(2)*(j-1)+i-1)+1:10*(nrank(2)*(j-1)+i))
            END DO
          END DO
          !$OMP END PARALLEL DO SIMD
        END IF
        ! Set communication vector type
        overlap_type=momentsoverlap_type
        rcvoverlap_type=rcvmomentsoverlap_type
    ELSE
        CALL start_rhscomm(mpirank, leftproc, rightproc, loc_moments, nrank, femorder, loc_zspan-femorder(1))
        IF(mpirank .gt. 0) THEN
        !$OMP PARALLEL DO SIMD DEFAULT(SHARED)
          DO j=1,femorder(1)
            DO i=1,nrank(2)
              loc_moments(1, (i-1)*loc_zspan+j) = loc_moments(1,(i-1)*loc_zspan+j)&
              & + rhsoverlap_buffer(nrank(2)*(j-1)+i)
            END DO
          END DO
          !$OMP END PARALLEL DO SIMD
        END IF
        ! Set communication vector type
        overlap_type=rhsoverlap_type
        rcvoverlap_type=rcvrhsoverlap_type
    END IF
    IF(mpirank.eq.0) THEN
      moments=0
    END IF
#ifdef _DEBUG
        WRITE(*,*)"started gatherv at step: ",step
        WRITE(*,*)"counts:", counts
        WRITE(*,*)"displs:", displs
#endif
    CALL MPI_GATHERV(loc_moments(1,1), counts(mpirank+1), overlap_type, &
    &                moments(1,1), counts, displs, rcvoverlap_type, 0, MPI_COMM_WORLD, ierr)
#ifdef _DEBUG
        WRITE(*,*)"ended gatherv at step: ",step
#endif
END SUBROUTINE fields_comm
!---------------------------------------------------------------------------  
!> @author
!> Patryk kaminski   EPFL/SPC 
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> 
!> @brief
!> Solves Poisson equation using FEM. Distributes the result on all MPI workers and interpolate the electric forces
!> for each particle.
! 
!---------------------------------------------------------------------------
  SUBROUTINE poisson
    USE basic, ONLY: rhs, nrank, pot 
    USE bsplines
    USE mumps_bsplines
    USE futils
    INTEGER:: ierr, jder(2)
    jder(1)=0
    jder(2)=0

  ! Only the root process solves Poisson
    IF(mpirank.eq.0) CALL bsolve(femat,rhs,phi_spline)
    
  ! Broadcast the solution to all MPI workers
    CALL MPI_Bcast(phi_spline, nrank(1)*nrank(2), db_type, 0, MPI_COMM_WORLD, ierr)
    matcoef= reshape(phi_spline, (/nrank(1),nrank(2)/))

  ! Evaluate the electric potential on the grid points
    CALL gridval(splrz, vec1, vec2, pot, jder, matcoef)
   
    IF( mpirank .eq. 0 .and. modulo(step,it2d) .eq. 0) THEN
    ! On the root process, compute the electric field for diagnostic purposes
      CALL gridval(splrz, vec1, vec2, Ez, (/1,0/))
      CALL gridval(splrz, vec1, vec2, Er, (/0,1/))
      Ez=-Ez
      Er=-Er
    END IF

  END SUBROUTINE poisson



  SUBROUTINE EForcescomp(p)
    Use beam, ONLY: particles
    TYPE(particles), INTENT(INOUT):: p
    INTEGER:: i, iend
    INTEGER:: nbunch=64
    INTEGER:: num_threads

    num_threads=omp_get_max_threads()
    nbunch=p%Nptot/num_threads ! Particle bunch size used when calling basfun
    nbunch=max(nbunch,1) ! Particle bunch size used when calling basfun
    nbunch=min(nbunch,64) ! Particle bunch size used when calling basfun

    ! Evaluate the electric potential and field at the particles position
  !$OMP PARALLEL DO DEFAULT(SHARED), PRIVATE(iend)
    DO i=1,p%Nptot,nbunch
      ! Avoid segmentation fault by accessing non relevant data
        iend=min(i+nbunch-1,p%Nptot)
        CALL getgrad(splrz, p%Z(i:iend), p%R(i:iend), p%pot(i:iend), p%EZ(i:iend), p%ER(i:iend))
        p%EZ(i:iend)=-p%Ez(i:iend)
        p%ER(i:iend)=-p%ER(i:iend)
      END DO
   !$OMP END PARALLEL DO

  END SUBROUTINE EForcescomp

!---------------------------------------------------------------------------  
!> @author
!> Patryk kaminski   EPFL/SPC 
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> 
!> @brief
!> Constucts the FEM matrix using bsplines initialized in fields_init
!---------------------------------------------------------------------------
SUBROUTINE fematrix

    USE bsplines
    USE mumps_bsplines

    REAL(kind=db), ALLOCATABLE :: xgauss(:,:), wgauss(:), zg(:), rg(:), wzg(:), wrg(:)
    REAL(kind=db), ALLOCATABLE :: f(:,:), aux(:)
    REAL(kind=db), ALLOCATABLE :: coefs(:)
    REAL(kind=db), ALLOCATABLE :: fun(:,:), fun2(:,:)
    REAL(kind=db)  :: contrib
    INTEGER, ALLOCATABLE :: idert(:), iderw(:) 
    INTEGER :: i, j, k, jt, iw, irow, jcol, mu, igauss, iterm, irow2, jcol2, mu2, kterms
    kterms=2 !Number of terms in weak form

    ALLOCATE(fun(1:femorder(1)+1,0:1),fun2(1:femorder(2)+1,0:1))!Arrays keeping values of b-splines at gauss node
    ALLOCATE(xgauss(ngauss(1)*ngauss(2),2), wgauss(ngauss(1)*ngauss(2)),zg(ngauss(1)),rg(ngauss(2)), wzg(ngauss(1)), wrg(ngauss(2)))   !Gaussian nodes and weights arrays 
    ALLOCATE(f((femorder(1)+1)*(femorder(2)+1),2),aux(femorder(1)+1))                                                                  !Auxiliary arrays ordering bsplines
    ALLOCATE(idert(kterms), iderw(kterms), coefs(kterms))                                                                              !Pointers on the order of derivatives 

! Constuction of auxiliary array ordering bsplines in given interval
    DO i=1,(femorder(1)+1)
       aux(i)=i
    END DO
    DO i=1,(femorder(2)+1)
       f((i-1)*(femorder(1)+1)+1:i*(femorder(1)+1),1)=aux
       f((i-1)*(femorder(1)+1)+1:i*(femorder(1)+1),2)=i
    END DO

! Initialisation of FEM matrix
    CALL init(nrank(1)*nrank(2),2,femat)
   
! Assemble FEM matrix
    DO j=1,nr ! Loop on r position
       DO i=1,nz	! Loop on z position
         ! Computation of gauss weight and position in r and z direction for gaussian integration
          CALL get_gauss(splrz%sp1, ngauss(1), i, zg, wzg)
          CALL get_gauss(splrz%sp2, ngauss(2), j, rg, wrg)
          ! Construction of matrix xgauss and wgauss storing the weight and position for 2d gaussian integration
          DO k=1,ngauss(2)
             xgauss((k-1)*ngauss(1)+1:k*ngauss(1),1)=zg
             xgauss((k-1)*ngauss(1)+1:k*ngauss(1),2)=rg(k)
             wgauss((k-1)*ngauss(1)+1:k*ngauss(1))=wrg(k)*wzg	
          END DO
          DO igauss=1,ngauss(1)*ngauss(2) ! Loop on gaussian weights and positions
             CALL basfun(xgauss(igauss,1), splrz%sp1, fun, i)
             CALL basfun(xgauss(igauss,2), splrz%sp2, fun2, j)
             CALL coefeq(xgauss(igauss,:), idert, iderw, coefs)
             DO iterm=1,kterms ! Loop on the two integration dimensions
                DO jt=1,(1+femorder(1))*(femorder(2)+1) 
                   DO iw=1,(1+femorder(1))*(femorder(2)+1)
                      irow=i+f(jt,1)-1; jcol=j+f(jt,2)-1
                      irow2=i+f(iw,1)-1; jcol2=j+f(iw,2)-1
                      mu=irow+(jcol-1)*nrank(1)
                      mu2=irow2+(jcol2-1)*nrank(1)
                      contrib=fun(f(jt,1),idert(iterm))* fun(f(iw,1),idert(iterm))* &
                           &  fun2(f(jt,2),iderw(iterm))*fun2(f(iw,2),iderw(iterm))* &
                           &  wgauss(igauss)*xgauss(igauss,2)
                      CALL updtmat(femat, mu, mu2, contrib)	
                   END DO
                END DO
             END DO
          END DO
       END DO
    END DO

! Impose Dirichlet boundary conditions on FEM matrix 
    CALL fe_dirichlet 

    DEALLOCATE(xgauss, wgauss,zg,rg, wzg, wrg)   
    DEALLOCATE(f, aux)
    DEALLOCATE(idert, iderw, coefs, fun,fun2)
  END SUBROUTINE fematrix

  SUBROUTINE comp_volume

    USE bsplines
    USE mumps_bsplines
    USE basic, ONLY: Volume

    REAL(kind=db), ALLOCATABLE :: xgauss(:,:), wgauss(:), zg(:), rg(:), wzg(:), wrg(:)
    REAL(kind=db), ALLOCATABLE :: f(:,:), aux(:)
    REAL(kind=db), ALLOCATABLE :: fun(:,:), fun2(:,:)
    INTEGER :: i, j, k, jt, irow, jcol, mu, igauss

    ALLOCATE(fun(1:femorder(1)+1,0:0),fun2(1:femorder(2)+1,0:0))!Arrays keeping values of b-splines at gauss node
    ALLOCATE(xgauss(ngauss(1)*ngauss(2),2), wgauss(ngauss(1)*ngauss(2)),zg(ngauss(1)),rg(ngauss(2)), wzg(ngauss(1)), wrg(ngauss(2)))   !Gaussian nodes and weights arrays 
    ALLOCATE(f((femorder(1)+1)*(femorder(2)+1),2),aux(femorder(1)+1))                                                                  !Auxiliary arrays ordering bsplines 

! Constuction of auxiliary array ordering bsplines in given interval
    DO i=1,(femorder(1)+1)
       aux(i)=i
    END DO
    DO i=1,(femorder(2)+1)
       f((i-1)*(femorder(1)+1)+1:i*(femorder(1)+1),1)=aux
       f((i-1)*(femorder(1)+1)+1:i*(femorder(1)+1),2)=i
    END DO
   
! Assemble Volume matrix
    DO j=1,nr ! Loop on r position
       DO i=1,nz	! Loop on z position
         ! Computation of gauss weight and position in r and z direction for gaussian integration
          CALL get_gauss(splrz%sp1, ngauss(1), i, zg, wzg)
          CALL get_gauss(splrz%sp2, ngauss(2), j, rg, wrg)
          ! Construction of matrix xgauss and wgauss storing the weight and position for 2d gaussian integration
          DO k=1,ngauss(2)
             xgauss((k-1)*ngauss(1)+1:k*ngauss(1),1)=zg
             xgauss((k-1)*ngauss(1)+1:k*ngauss(1),2)=rg(k)
             wgauss((k-1)*ngauss(1)+1:k*ngauss(1))=wrg(k)*wzg	
          END DO
          DO igauss=1,ngauss(1)*ngauss(2) ! Loop on gaussian weights and positions
             CALL basfun(xgauss(igauss,1), splrz%sp1, fun, i)
             CALL basfun(xgauss(igauss,2), splrz%sp2, fun2, j)
             DO jt=1,(1+femorder(1))*(femorder(2)+1) 
                irow=i+f(jt,1)-1; jcol=j+f(jt,2)-1
                mu=irow+(jcol-1)*nrank(1)
                volume(mu)=volume(mu)+2*pi*fun(f(jt,1),0) * fun2(f(jt,2),0)* wgauss(igauss)*xgauss(igauss,2)
             END DO
          END DO
       END DO
    END DO


    DEALLOCATE(xgauss, wgauss,zg,rg, wzg, wrg)   
    DEALLOCATE(f, aux)                      
    DEALLOCATE(fun,fun2)

  END SUBROUTINE comp_volume

!---------------------------------------------------------------------------  
!> @author
!> Patryk kaminski   EPFL/SPC 
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> 
!> @brief
!> Imposes the dirichlet boundary conditions on the FEM matrix.
!---------------------------------------------------------------------------
  SUBROUTINE fe_dirichlet
    REAL(kind=db), ALLOCATABLE :: arr(:)
    INTEGER :: i
    ALLOCATE(arr(nrank(1)*nrank(2)))
    DO i=1,nrank(1)*nrank(2)
      CALL getrow(femat,i,arr)
    END DO
    DO i=1,nrank(1)
       IF(rgrid(0).ne.0.0_db) THEN
         arr=0; arr(i)=1; 
         CALL putrow(femat,i,arr)
       END IF
       arr=0; arr(nrank(1)*nrank(2)+1-i)=1 ; 
       CALL putrow(femat,nrank(1)*nrank(2)+1-i,arr)  
    END DO
    DEALLOCATE(arr)
  END SUBROUTINE fe_dirichlet
!________________________________________________________________________________
  SUBROUTINE coefeq(x, idt, idw, c)
    REAL(kind=db), INTENT(in) :: x(2)
    INTEGER, INTENT(out) :: idt(:), idw(:)
    REAL(kind=db), INTENT(out) :: c(SIZE(idt))

    c(1) = x(2)    
    idt(1) = 0     
    idw(1) = 1     
    c(2) = x(2)     
    idt(2) = 1     
    idw(2) = 0 
  END SUBROUTINE coefeq

!---------------------------------------------------------------------------  
!> @author
!> Patryk kaminski   EPFL/SPC 
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> 
!> @brief
!> Computes the magnetic field on the grid according to a magnetic mirror.
!---------------------------------------------------------------------------
SUBROUTINE magnet
    USE basic, ONLY: B0, Rcurv, rgrid, zgrid, width, rnorm, nr, nz, bnorm
    USE constants, ONLY: Pi
    REAL(kind=db) :: rg, zg,halfLz, MirrorRatio
    INTEGER :: i, rindex
    halfLz=(zgrid(nz)+zgrid(0))/2
    MirrorRatio=(Rcurv-1)/(Rcurv+1)
    DO i=1,(nr+1)*(nz+1)
       rindex=(i-1)/(nz+1)
       rg=rgrid(rindex)
       zg=zgrid(i-rindex*(nz+1)-1)-halfLz
       Br(i)=-B0*MirrorRatio*SIN(2*pi*zg/width*rnorm)*bessi1(2*pi*rg/width*rnorm)/bnorm
       Bz(i)=B0*(1-MirrorRatio*COS(2*pi*zg/width*rnorm)*bessi0(2*pi*rg/width*rnorm))/bnorm
       Athet(i)=0.5*B0*(rg*rnorm - width/pi*MirrorRatio*bessi1(2*pi*rg/width*rnorm)*COS(2*pi*zg/width*rnorm))
    END DO
  END SUBROUTINE magnet
!________________________________________________________________________________
!Modified Bessel functions of the first kind of the zero order
FUNCTION bessi0(x)
    REAL(kind=db) :: bessi0,x
    REAL(kind=db) ::  ax
    REAL(kind=db) p1,p2,p3,p4,p5,p6,p7,q1,q2,q3,q4,q5,q6,q7,q8,q9,y
    SAVE p1,p2,p3,p4,p5,p6,p7,q1,q2,q3,q4,q5,q6,q7,q8,q9
    DATA p1,p2,p3,p4,p5,p6,p7/1.0d0,3.5156229d0,3.0899424d0,1.2067492d0,0.2659732d0,0.360768d-1,0.45813d-2/
    DATA q1,q2,q3,q4,q5,q6,q7,q8,q9 /0.39894228d0, 0.1328592d-1, 0.225319d-2, -0.157565d-2 ,0.916281d-2, &
         &                          -0.2057706d-1, 0.2635537d-1,-0.1647633d-1, 0.392377d-2 /
    if (abs(x).lt.3.75) then
       y=(x/3.75)**2
       bessi0=p1+y*(p2+y*(p3+y*(p4+y*(p5+y*(p6+y*p7)))))
    else
       ax=abs(x)
       y=3.75/ax
       bessi0=(exp(ax)/sqrt(ax))*(q1+y*(q2+y*(q3+y*(q4+y*(q5+y*(q6+y*(q7+y*(q8+y*q9))))))))
    endif
    return
  END FUNCTION bessi0
!________________________________________________________________________________
!Modified Bessel functions of the first kind of the first order
  FUNCTION bessi1(x)
    REAL(kind=db) ::  bessi1,x
    REAL(kind=db) ::  ax
    REAL(kind=db) p1,p2,p3,p4,p5,p6,p7,q1,q2,q3,q4,q5,q6,q7,q8,q9,y
    SAVE p1,p2,p3,p4,p5,p6,p7,q1,q2,q3,q4,q5,q6,q7,q8,q9
    DATA p1,p2,p3,p4,p5,p6,p7/0.5d0,0.87890594d0,0.51498869d0,0.15084934d0,0.2658733d-1,0.301532d-2,0.32411d-3/
    DATA q1,q2,q3,q4,q5,q6,q7,q8,q9 /0.39894228d0, -0.3988024d-1, -0.362018d-2,  0.163801d-2,-0.1031555d-1, &
         &                           0.2282967d-1, -0.2895312d-1,  0.1787654d-1,-0.420059d-2/
    if (abs(x).lt.3.75) then
       y=(x/3.75)**2
       bessi1=x*(p1+y*(p2+y*(p3+y*(p4+y*(p5+y*(p6+y*p7))))))
    else
       ax=abs(x)
       y=3.75/ax
       bessi1=(exp(ax)/sqrt(ax))*(q1+y*(q2+y*(q3+y*(q4+y*(q5+y*(q6+y*(q7+y*(q8+y*q9))))))))
       if(x.lt.0.)bessi1=-bessi1
    endif
    return
  END FUNCTION bessi1

!---------------------------------------------------------------------------  
!> @author
!> Patryk kaminski   EPFL/SPC 
!> Guillaume Le Bars EPFL/SPC
!
! DESCRIPTION: 
!> 
!> @brief
!> Free the memory used by the fields module
!---------------------------------------------------------------------------
  SUBROUTINE clean_fields
    Use bsplines
    USE basic, ONLY: rhs, moments
    DEALLOCATE(matcoef)
    DEALLOCATE(pot)
    DEALLOCATE(rhs)
    DEALLOCATE(phi_spline)
    DEALLOCATE(moments)
    DEALLOCATE(Br,Bz)
    DEALLOCATE(Er,Ez)
    DEALLOCATE(vec1,vec2)
    Call DESTROY_SP(splrz)

  END SUBROUTINE clean_fields

END MODULE fields