subroutine local_pssolver
use precision
use global_parameters
use particle_array
use grids_fields
use grids_setup
use grids_eqdata
use read_in_parameters, only: local_solver_model, qion, qelectron, aion, n_modes, min_r,&
                              ifullk,qfki,kinetice, ifluxtube
use simulation_domain, only: sfcore, sfsol,couplefrc
use cell_volume, only: volume
implicit none

integer :: i, j, ij, ij0, ierror, icount, ipade, spectral_method


real(lk) :: kperp2(mgrid), rdum, rdum_inv, rhos2, rhoi2, dielectric_constant, &

            coe_temp, ne_ni_ratio, phi_temp(mgrid), dentmp(mgrid), &

            dentmpr(mgrid), dentmpi(mgrid), &

            phitmpr(mgrid), phitmpi(mgrid), &

            ffilter(mtoroidal/2+1), xz(mtoroidal), phifluxlz(mgrid,mtoroidal), &

            eachzeta(mgrid), total_v, eachzphifluxlz_core(mygrid_core,mtoroidal),&

            eachzphifluxlz_sol(mygrid_sol,mtoroidal), eachtphifluxlz_core(mgrid_core),&

            eachtphifluxlz_sol(mgrid_sol), data1dtmp(mpsi,2), total_v1d(mpsi)


complex(lk) :: yz(mtoroidal/2+1)

ipade=1

if(local_solver_model==1)then
    spectral_method=0
else
    spectral_method=1
endif

if(local_solver_model==2)then
    ipade=0
elseif(local_solver_model==3)then
    ipade=1
endif 


!========================Calculation of kperp^2 ===========================
if(sfcore==1)then
!$omp parallel do private(i,j,ij,rdum,rdum_inv)
    do i=1,mpsi_core
        do j=1,mpara_core(i)
            ij=igrid_core(i)+j
            rdum=max(min_r,abs(Rcore(ij)))
            rdum_inv=1.0_lk/rdum
            kperp2(ij)=real(nmode)*real(nmode)*rdum_inv*rdum_inv*real(n_modes(1))*real(n_modes(1))
        enddo
    enddo
endif
  
if(sfsol==1)then
!$omp parallel do private(i,j,ij,ij0,rdum,rdum_inv)
    do i=1,mpsi_sol
        do j=1,mpara_sol(i)
            ij=mgrid_core+igrid_sol(i)+j
            ij0=igrid_sol(i)+j
            rdum=max(min_r,Rsol(ij0))
            rdum_inv=1.0_lk/rdum
            kperp2(ij)=real(nmode)*real(nmode)*rdum_inv*rdum_inv*real(n_modes(1))*real(n_modes(1))
        enddo
    enddo
endif
!==========================================================================

eachzeta=0.0_lk
icount=mgrid

if(sfcore==1)then
    if(ifullk==0)then
!$omp parallel do private(i, j, ij, ne_ni_ratio, coe_temp) 
    do i=1,mpsi_core
        do j=1,mpara_core(i)
            ij=igrid_core(i)+j
            if(kinetice==0)then
            ne_ni_ratio = meshne(ij)/meshni(ij)
            coe_temp = 1.0_lk + real(ipade)*meshti(ij)/meshte(ij)/qion/qion*ne_ni_ratio
            dentmp(ij) = qion*meshni(ij)*densityi(1,ij)/coe_temp
            else
                dentmp(ij) = qion*meshni(ij)*densityi(1,ij) + qelectron*meshne(ij)*densitye(1,ij)
            endif
            eachzeta(ij)=dentmp(ij)
        enddo
    enddo
    else
!$omp parallel do private(i, j, ij, ne_ni_ratio, coe_temp) 
    do i=1,mpsi_core
        do j=1,mpara_core(i)
            ij=igrid_core(i)+j
            dentmp(ij) = qfki*meshni(ij)*densityfki(1,ij)
            eachzeta(ij)=dentmp(ij)
        enddo
    enddo
    endif
endif
if(sfsol==1)then
    if(ifullk==0)then
!$omp parallel do private(i, j, ij, ne_ni_ratio, coe_temp) 
    do i=1,mpsi_sol
        do j=1,mpara_sol(i)
            ij=mgrid_core+igrid_sol(i)+j
            if(kinetice==0)then
            ne_ni_ratio = meshne(ij)/meshni(ij)
            coe_temp = 1.0_lk + real(ipade)*meshti(ij)/meshte(ij)/qion/qion*ne_ni_ratio
            dentmp(ij) = qion*meshni(ij)*densityi(1,ij)/coe_temp
            else
                dentmp(ij) = qion*meshni(ij)*densityi(1,ij) + qelectron*meshne(ij)*densitye(1,ij)
            endif
            eachzeta(ij)=dentmp(ij)
        enddo
    enddo
    else
!$omp parallel do private(i, j, ij, ne_ni_ratio, coe_temp) 
    do i=1,mpsi_sol
        do j=1,mpara_sol(i)
            ij=mgrid_core+igrid_sol(i)+j
            dentmp(ij) = qfki*meshni(ij)*densityfki(1,ij)
            eachzeta(ij)=dentmp(ij)
        enddo
    enddo
    endif
endif

call MPI_ALLGATHER(eachzeta,icount,mpi_Rsize,phifluxlz,icount,&
                   mpi_Rsize,toroidal_comm,ierror)
  
if(sfcore==1)then
    yz=0.0_lk
    ffilter=0.0_lk
    ffilter(n_modes+1)=1.0_lk
!$omp parallel do private(ij,xz,yz)
    do ij=ij_min_tor_core, ij_max_tor_core
            xz=phifluxlz(ij,:)
            call fftr1d(1,mtoroidal,1.0_lk,xz,yz,2)
! linear run only keep a few modes
            yz=ffilter*yz
            dentmpr(ij)=real(real(yz(n_modes(1)+1)))
            dentmpi(ij)=real(aimag(yz(n_modes(1)+1)))
    enddo
endif

if(sfsol==1)then
    yz=0.0_lk
    ffilter=0.0_lk
    ffilter(n_modes+1)=1.0_lk
!$omp parallel do private(ij,xz,yz)
    do ij=ij_min_tor_sol, ij_max_tor_sol
            xz=phifluxlz(ij,:)
            call fftr1d(1,mtoroidal,1.0_lk,xz,yz,2)
! linear run only keep a few modes
            yz=ffilter*yz
            dentmpr(ij)=real(real(yz(n_modes(1)+1)))
            dentmpi(ij)=real(aimag(yz(n_modes(1)+1)))
    enddo
endif


!====================== Zero boundary condition =======================
if(ifluxtube==0)then
    if(sfcore == 1)then
!$omp parallel do private(j, ij)
        do j = 1, mpara_core(1)
            ij = igrid_core(1) + j
            dentmpr(ij) = 0.0_lk
            dentmpi(ij) = 0.0_lk
        enddo

        if(couplefrc == 0)then
!$omp parallel do private(j, ij)
            do j = 1,mpara_core(mpsi_core)
                ij = igrid_core(mpsi_core) + j
                dentmpr(ij) = 0.0_lk
                dentmpi(ij) = 0.0_lk
            enddo
        endif
    endif

    if(sfsol == 1)then
!$omp parallel do private(j, ij)
        do j = 1, mpara_sol(mpsi_sol)
            ij = mgrid_core + igrid_sol(mpsi_sol) + j
            dentmpr(ij) = 0.0_lk
            dentmpi(ij) = 0.0_lk
        enddo

        if(couplefrc == 0)then
!$omp parallel do private(j, ij)
            do j = 1, mpara_sol(1)
                ij = mgrid_core + igrid_sol(1) + j
                dentmpr(ij) = 0.0_lk
                dentmpi(ij) = 0.0_lk
            enddo
        endif
    endif
endif


!==========================================================================
!============================= Local Solver ===============================

phi=0.0_lk
phitmpr=0.0_lk
phitmpi=0.0_lk

if(local_solver_model==1)then
    if(kinetice==1) write(*,*) 'kinetic electron with local_solver_model==1 not supported.'
    stop
    if(sfcore==1)then
    if(ifullk==0)then
!$omp parallel do private(ij,dielectric_constant)
        do ij = ij_min_core, ij_max_core
                dielectric_constant = meshne(ij)/meshte(ij)
                phi(1,ij) = qion*meshni(ij)*densityi(1,ij)/dielectric_constant
        enddo
    else
!$omp parallel do private(ij,dielectric_constant)
        do ij = ij_min_core, ij_max_core
                dielectric_constant = meshne(ij)/meshte(ij)
                phi(1,ij) = qfki*meshni(ij)*densityfki(1,ij)/dielectric_constant
        enddo
    endif
    endif
    
    if(sfsol==1)then
    if(ifullk==0)then
!$omp parallel do private(ij,dielectric_constant)
        do ij = ij_min_sol, ij_max_sol
                dielectric_constant = meshne(ij)/meshte(ij)
                phi(1,ij) = qion*meshni(ij)*densityi(1,ij)/dielectric_constant
        enddo
    else
!$omp parallel do private(ij,dielectric_constant)
        do ij = ij_min_sol, ij_max_sol
                dielectric_constant = meshne(ij)/meshte(ij)
                phi(1,ij) = qfki*meshni(ij)*densityfki(1,ij)/dielectric_constant
        enddo
    endif
    endif
    call MPI_ALLREDUCE(phi(1,:),phi_temp,mgrid,mpi_Rsize,MPI_SUM,partd_comm,ierror)
    phi(1,:) = phi_temp/npartdom

elseif(local_solver_model==2)then

    if(sfcore==1)then
    if(ifullk==0)then
!$omp parallel do private(ij, rhos2, ne_ni_ratio, dielectric_constant)
        do ij=ij_min_tor_core, ij_max_tor_core
                rhos2=aion*meshte(ij)/bmesh(ij)/bmesh(ij)/qion/qion
                ne_ni_ratio = meshne(ij)/meshni(ij)
                if(kinetice==0)then
                dielectric_constant=(ne_ni_ratio + qion*qion*kperp2(ij)*rhos2)*meshni(ij)/meshte(ij)
                else
                    dielectric_constant = qion*qion*kperp2(ij)*rhos2*meshni(ij)/meshte(ij)
                endif
                phitmpr(ij) = dentmpr(ij)/dielectric_constant
                phitmpi(ij) = dentmpi(ij)/dielectric_constant
        enddo
    else
!$omp parallel do private(ij, rhos2, ne_ni_ratio, dielectric_constant)
        do ij=ij_min_tor_core, ij_max_tor_core
                dielectric_constant=meshne(ij)/meshte(ij)
                phitmpr(ij) = dentmpr(ij)/dielectric_constant
                phitmpi(ij) = dentmpi(ij)/dielectric_constant
        enddo
    endif
    endif

    if(sfsol==1)then
    if(ifullk==0)then
!$omp parallel do private(ij, rhos2, ne_ni_ratio, dielectric_constant)
        do ij=ij_min_tor_sol, ij_max_tor_sol
                rhos2=aion*meshte(ij)/bmesh(ij)/bmesh(ij)/qion/qion
                ne_ni_ratio = meshne(ij)/meshni(ij)
                if (kinetice==0)then
                dielectric_constant=(ne_ni_ratio + qion*qion*kperp2(ij)*rhos2)*meshni(ij)/meshte(ij)
                else
                    dielectric_constant = qion*qion*kperp2(ij)*rhos2*meshni(ij)/meshte(ij)
                endif
                phitmpr(ij) = dentmpr(ij)/dielectric_constant
                phitmpi(ij) = dentmpi(ij)/dielectric_constant
        enddo
    else
!$omp parallel do private(ij, rhos2, ne_ni_ratio, dielectric_constant)
        do ij=ij_min_tor_sol, ij_max_tor_sol
                dielectric_constant=meshne(ij)/meshte(ij)
                phitmpr(ij) = dentmpr(ij)/dielectric_constant
                phitmpi(ij) = dentmpi(ij)/dielectric_constant
        enddo
    endif
    endif

elseif(local_solver_model==3)then

    if(sfcore==1)then
    if(ifullk==0)then
!$omp parallel do private(ij, rhos2, ne_ni_ratio, coe_temp, dielectric_constant)
        do ij=ij_min_tor_core, ij_max_tor_core
                rhos2=aion*meshte(ij)/bmesh(ij)/bmesh(ij)/qion/qion

                ne_ni_ratio = meshne(ij)/meshni(ij)
                if(kinetice==0)then
                coe_temp = 1.0_lk + real(ipade)*meshti(ij)/meshte(ij)/qion/qion*ne_ni_ratio

                dielectric_constant = (kperp2(ij)*rhos2*qion*qion*coe_temp + ne_ni_ratio)*meshni(ij)/meshte(ij)
                else
                    dielectric_constant = qion*qion*rhos2*kperp2(ij)*meshni(ij)/meshte(ij)
                endif
                phitmpr(ij) = dentmpr(ij)/dielectric_constant
                phitmpi(ij) = dentmpi(ij)/dielectric_constant

                phitmpr(ij) = phitmpr(ij) + real(ipade)*(dentmpr(ij)/qion)*(meshti(ij)/qion/meshni(ij))
                phitmpi(ij) = phitmpi(ij) + real(ipade)*(dentmpi(ij)/qion)*(meshti(ij)/qion/meshni(ij))
        enddo
    else
!$omp parallel do private(ij, rhos2, ne_ni_ratio, coe_temp, dielectric_constant)
        do ij=ij_min_tor_core, ij_max_tor_core
                dielectric_constant = meshne(ij)/meshte(ij)
                phitmpr(ij) = dentmpr(ij)/dielectric_constant
                phitmpi(ij) = dentmpi(ij)/dielectric_constant
        enddo
    endif
    endif

    if(sfsol==1)then
    if(ifullk==0)then
!$omp parallel do private(ij, rhos2, ne_ni_ratio, coe_temp, dielectric_constant)
        do ij=ij_min_tor_sol, ij_max_tor_sol
                rhos2=aion*meshte(ij)/bmesh(ij)/bmesh(ij)/qion/qion
                ne_ni_ratio = meshne(ij)/meshni(ij)
                if(kinetice==0)then
                coe_temp = 1.0_lk + real(ipade)*meshti(ij)/meshte(ij)/qion/qion*ne_ni_ratio

                dielectric_constant = (kperp2(ij)*rhos2*qion*qion*coe_temp + ne_ni_ratio)*meshni(ij)/meshte(ij)
                else
                    dielectric_constant = qion*qion*rhos2*kperp2(ij)*meshni(ij)/meshte(ij)
                endif
                phitmpr(ij) = dentmpr(ij)/dielectric_constant
                phitmpi(ij) = dentmpi(ij)/dielectric_constant

                phitmpr(ij) = phitmpr(ij) + real(ipade)*(dentmpr(ij)/qion)*(meshti(ij)/qion/meshni(ij))
                phitmpi(ij) = phitmpi(ij) + real(ipade)*(dentmpi(ij)/qion)*(meshti(ij)/qion/meshni(ij))
        enddo
    else
!$omp parallel do private(ij, rhos2, ne_ni_ratio, coe_temp, dielectric_constant)
        do ij=ij_min_tor_sol, ij_max_tor_sol
                dielectric_constant = meshne(ij)/meshte(ij)
                phitmpr(ij) = dentmpr(ij)/dielectric_constant
                phitmpi(ij) = dentmpi(ij)/dielectric_constant
        enddo
    endif
    endif

else 

    if(mype==0)write(gtcout,*)"wrong set of local_solver_model"
    stop

endif  

!==================================================================
if(spectral_method==1)then
!================== transform back to real space ==================

eachzphifluxlz_core=0.0_lk
eachzphifluxlz_sol=0.0_lk
eachtphifluxlz_core=0.0_lk
eachtphifluxlz_sol=0.0_lk

    call mpi_allreduce(mpi_in_place, phitmpr, mgrid, MPI_RSIZE, MPI_SUM, mpi_comm_world, ierror)
    call mpi_allreduce(mpi_in_place, phitmpi, mgrid, MPI_RSIZE, MPI_SUM, mpi_comm_world, ierror)
    if(sfcore==1)then
        yz=0.0_lk
        xz=0.0_lk
!$omp parallel do private(i,j,ij,xz,yz)
        do ij=ij_min_tor_core, ij_max_tor_core
                yz(n_modes(1)+1)=cmplx(phitmpr(ij),phitmpi(ij))
                yz=ffilter*yz
                call fftr1d(-1,mtoroidal,1.0_lk,xz,yz,2)
                eachzphifluxlz_core(ij-ij_min_core+1,:)=xz
        enddo

        CALL MPI_ALLREDUCE(MPI_IN_PLACE,eachzphifluxlz_core,mygrid_core*mtoroidal,MPI_RSIZE, MPI_SUM, toroidal_comm,ierror)
        eachtphifluxlz_core(ij_min_core:ij_max_core)=eachzphifluxlz_core(:,myrank_toroidal+1)
        CALL MPI_ALLREDUCE(MPI_IN_PLACE,eachtphifluxlz_core,mgrid_core,MPI_RSIZE, MPI_SUM, partd_comm,ierror)
 
!$omp parallel do private(i,j,ij) 
        do i=1,mpsi_core
            do j=1,mpara_core(i)
                ij=igrid_core(i)+j
                phi(1,ij)=eachtphifluxlz_core(ij)
            enddo
        enddo
    endif
    
    if(sfsol==1)then
        if(ifluxtube==1)then
        phi_real=0.0_lk
        phi_imag=0.0_lk
        coe_temp=1.0_lk/mpsi_sol
        do i=1,mpsi_sol
           do j=1,mpara_sol(i)
              ij = mgrid_core+igrid_sol(i)+j
              phi_real(mgrid_core+igrid_sol(1)+j)=phi_real(mgrid_core+igrid_sol(1)+j)+phitmpr(ij)*coe_temp
              phi_imag(mgrid_core+igrid_sol(1)+j)=phi_imag(mgrid_core+igrid_sol(1)+j)+phitmpi(ij)*coe_temp
           enddo
        enddo
        do i=1,mpsi_sol
           do j=1,mpara_sol(i)
              ij = mgrid_core+igrid_sol(i)+j
              phitmpr(ij)=phi_real(mgrid_core+igrid_sol(1)+j)
              phitmpi(ij)=phi_imag(mgrid_core+igrid_sol(1)+j)
           enddo
        enddo
        endif
     
        yz=0.0_lk
        xz=0.0_lk
!$omp parallel do private(ij,xz,yz)
        do ij=ij_min_tor_sol, ij_max_tor_sol
                yz(n_modes(1)+1)=cmplx(phitmpr(ij),phitmpi(ij))
                yz=ffilter*yz
                call fftr1d(-1,mtoroidal,1.0_lk,xz,yz,2)
                eachzphifluxlz_sol(ij-ij_min_sol+1,:)=xz
        enddo

        CALL MPI_ALLREDUCE(MPI_IN_PLACE,eachzphifluxlz_sol,mygrid_sol*mtoroidal,MPI_RSIZE, MPI_SUM, toroidal_comm,ierror)
        eachtphifluxlz_sol(ij_min_sol:ij_max_sol)=eachzphifluxlz_sol(:,myrank_toroidal+1)
        CALL MPI_ALLREDUCE(MPI_IN_PLACE,eachtphifluxlz_sol,mgrid_sol,MPI_RSIZE, MPI_SUM, partd_comm,ierror)
    
!$omp parallel do private(i,j,ij) 
        do i=1,mpsi_sol
            do j=1,mpara_sol(i)
                ij=mgrid_core+igrid_sol(i)+j
                phi(1,ij)=eachtphifluxlz_sol(ij)
            enddo
        enddo
    endif
    phi_real=phitmpr
    phi_imag=phitmpi

endif
!==================================================================


!====================== Zero boundary condition for phi =======================
    if(ifluxtube==0)then
    if(sfcore == 1)then
!$omp parallel do private(j, ij)
        do j = 1, mpara_core(1)
            ij = igrid_core(1) + j
            phi(1, ij) = 0.0_lk
            phi_real(ij) = 0.0_lk
            phi_imag(ij) = 0.0_lk
        enddo

        if(couplefrc == 0)then
!$omp parallel do private(j, ij)
            do j = 1,mpara_core(mpsi_core)
                ij = igrid_core(mpsi_core) + j
                phi(1, ij) = 0.0_lk
                phi_real(ij) = 0.0_lk
                phi_imag(ij) = 0.0_lk
            enddo
        endif
    endif

    if(sfsol == 1)then
!$omp parallel do private(j, ij)
        do j = 1, mpara_sol(mpsi_sol)
            ij = mgrid_core + igrid_sol(mpsi_sol) + j
            phi(1, ij) = 0.0_lk
            phi_real(ij) = 0.0_lk
            phi_imag(ij) = 0.0_lk
        enddo

        if(couplefrc == 0)then
!$omp parallel do private(j, ij)
            do j = 1, mpara_sol(1)
                ij = mgrid_core + igrid_sol(1) + j
                phi(1, ij) = 0.0_lk
                phi_real(ij) = 0.0_lk
                phi_imag(ij) = 0.0_lk
            enddo
        endif
    endif
    endif
!==================================================================


CALL SMOOTH(phi)
call toroidal_filter(phi)

phi_rms=0.0_lk
phi0_rms=0.0_lk
total_v = 0.0_lk
data1dtmp = 0.0_lk
total_v1d = 0.0_lk
if(sfcore==1)then
  do i=1,mpsi_core
     do j=1,mpara_core(i)
        ij=igrid_core(i)+j
        data1dtmp(i,1) = data1dtmp(i,1)+phi(1,ij)**2*volume(ij)
        data1dtmp(i,2) = data1dtmp(i,2)+phi0(1,ij)**2*volume(ij)
        total_v1d(i) = total_v1d(i) + volume(ij)
     enddo
     phi_rms = phi_rms + data1dtmp(i,1)
     phi0_rms = phi0_rms + data1dtmp(i,2)
     total_v = total_v + total_v1d(i)
  enddo
endif
if(sfsol==1)then
  do i=1,mpsi_sol
     do j=1,mpara_sol(i)
        ij=mgrid_core+igrid_sol(i)+j
        data1dtmp(mpsi_core+i,1) = data1dtmp(mpsi_core+i,1)+phi(1,ij)**2*volume(ij)
        data1dtmp(mpsi_core+i,2) = data1dtmp(mpsi_core+i,2)+phi0(1,ij)**2*volume(ij)
        total_v1d(mpsi_core+i) = total_v1d(mpsi_core+i) + volume(ij)
     enddo
     phi_rms = phi_rms + data1dtmp(mpsi_core+i,1)
     phi0_rms = phi0_rms + data1dtmp(mpsi_core+i,2)
     total_v = total_v + total_v1d(mpsi_core+i)
  enddo
endif

phi_rms = sqrt(phi_rms/total_v)
phi0_rms = sqrt(phi0_rms/total_v)
call MPI_ALLREDUCE(phi_rms,total_v,1,mpi_Rsize,MPI_SUM,toroidal_comm,ierror)
total_v = total_v / mtoroidal
phi_rms = total_v
call MPI_ALLREDUCE(phi0_rms,total_v,1,mpi_Rsize,MPI_SUM,toroidal_comm,ierror)
total_v = total_v / mtoroidal
phi0_rms = total_v

data1dtmp(:,1) = sqrt(data1dtmp(:,1)/total_v1d)
data1dtmp(:,2) = sqrt(data1dtmp(:,2)/total_v1d)
call MPI_ALLREDUCE(data1dtmp(:,1),total_v1d,mpsi,mpi_Rsize,MPI_SUM,toroidal_comm,ierror)
total_v1d = total_v1d / mtoroidal
data1dtmp(:,1) = total_v1d
call MPI_ALLREDUCE(data1dtmp(:,2),total_v1d,mpsi,mpi_Rsize,MPI_SUM,toroidal_comm,ierror)
total_v1d = total_v1d / mtoroidal
data1dtmp(:,2) = total_v1d

if(ifullk==0)then
data1di(:,4:5) = data1dtmp
else
data1dfki(:,4:5) = data1dtmp
endif

end subroutine local_pssolver
