! The MIT License (MIT)
!
! Copyright (c) 2016 odlomax
!
! Permission is hereby granted, free of charge, to any person obtaining a copy
! of this software and associated documentation files (the "Software"), to deal
! in the Software without restriction, including without limitation the rights
! to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
! copies of the Software, and to permit persons to whom the Software is
! furnished to do so, subject to the following conditions:
!
! The above copyright notice and this permission notice shall be included in all
! copies or substantial portions of the Software.
!
! THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
! IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
! FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
! AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
! LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
! OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
! SOFTWARE.

! mcrt simulation module
module m_simulation

   use m_kind_parameters
   use m_constants_parameters
   use m_maths
   use m_atomic_update
   use m_dust
   use m_dust_d03
   use m_source
   use m_source_point_bb
   use m_source_point_ms_star
   use m_source_point_custom
   use m_source_external_bb
   use m_source_external_ps05
   use m_particle
   use m_binary_tree
   use m_kernel
   use m_kernel_m4
   use m_kernel_m6
   use m_ray
   use m_options
   use m_datacube
   use m_io
   use m_string
   use omp_lib

   implicit none

   private

   public :: simulation

   ! define mcrt simulation class
   type :: simulation

      type(particle),pointer,contiguous :: particle_array(:)      ! array of sph particles
      type(binary_tree),pointer :: sph_tree                       ! tree of particles
      type(ray),pointer :: lum_packet_array(:)                    ! array of luminosity packets
      type(options),pointer :: sim_params                         ! simulation parameters
      type(datacube),pointer :: intensity_cube                    ! datacube
      class(kernel),pointer :: sph_kernel                         ! smoothing kernel
      class(dust),pointer :: dust_prop                            ! interstellar dust properties
      class(source),pointer :: point_source_array(:)              ! array of point sources (stars)
      class(source),pointer :: isrf_prop                          ! interstellar radiation field properties

      contains

      procedure,non_overridable :: initialise
      procedure,non_overridable :: destroy
      procedure,non_overridable :: perform_iteration
      procedure,non_overridable :: source_extract
      procedure,non_overridable :: make_datacube
      procedure,non_overridable :: write_particles_bin
      procedure,non_overridable :: read_particles_bin

   end type

   contains

   ! initialise simulation object
   subroutine initialise(self,sim_params)

      ! argument declarations
      class(simulation),intent(inout) :: self                     ! simulation object
      type(options),intent(in),target :: sim_params               ! simulation parameters

      ! variable declarations
      integer(kind=int_kind) :: i,j                               ! counter
      integer(kind=int_kind) :: string_len                        ! length of string
      logical(kind=log_kind) :: h_present                         ! is smoothing length present in cloud file
      real(kind=rel_kind),allocatable :: position(:,:)            ! particle/star positions
      real(kind=rel_kind),allocatable :: mass(:)                  ! particle/star mass
      real(kind=rel_kind),allocatable :: temperature(:)           ! particle/star temperature
      real(kind=rel_kind),allocatable :: luminosity(:)            ! star luminosity
      real(kind=rel_kind),allocatable :: age(:)                   ! star age
      real(kind=rel_kind),allocatable :: metallicity(:)           ! star metallicity
      character(kind=chr_kind,len=string_length) :: file_name     ! name of file

      ! associate parameters
      self%sim_params=>sim_params

      write(*,"(A)") "initialise kernel"
      select case (trim(self%sim_params%sph_kernel))
         case ("m4")
            allocate(kernel_m4::self%sph_kernel)
         case ("m6")
            allocate(kernel_m6::self%sph_kernel)
      end select
      call self%sph_kernel%initialise()

      open(1,file="kernel.dat")
      do i=1,size(self%sph_kernel%r_array)
         write(1,*) self%sph_kernel%r_array(i),self%sph_kernel%w(self%sph_kernel%r_array(i)),&
            &self%sph_kernel%dw_dr(self%sph_kernel%r_array(i)),self%sph_kernel%d2w_dr2(self%sph_kernel%r_array(i))
      end do
      close(1)

      write(*,"(A)") "initialise dust"
      allocate(dust_d03::self%dust_prop)
      call self%dust_prop%initialise(self%sim_params%dust_r_v,self%sim_params%dust_t_min,self%sim_params%dust_t_max,&
         &self%sim_params%dust_n_t,iso_scatter=self%sim_params%dust_iso_scatter)

      write(*,"(A)") "initialise modified random walk"
      call self%dust_prop%mrw_initialise(self%sim_params%sim_n_mrw)

      write(*,"(A)") "read in particles"
      ! check if cloud input is ascii (.dat) or binary (.bin)
      string_len=len_trim(self%sim_params%sim_cloud_file)
      select case (self%sim_params%sim_cloud_file(string_len-2:string_len))
         case ("bin")
            call self%read_particles_bin()
            h_present=.true.

         case ("dat")
            call read_in_sph_particles_3d(position,mass,temperature,self%sim_params%sim_cloud_file)
            allocate(self%particle_array(size(position,2)))
            write(*,"(A)") "initialise particles"
            do i=1,size(self%particle_array)
               if (self%sim_params%sph_scattered_light) then
                  call self%particle_array(i)%initialise(position(:,i),(/(0._rel_kind,j=1,n_dim)/),mass(i),&
                     &4._rel_kind*pi*self%dust_prop%bol_mass_emissivity(temperature(i)),&
                     &self%sim_params%sph_lambda_min,self%sim_params%sph_lambda_max,self%sim_params%sph_n_lambda)
               else
                  call self%particle_array(i)%initialise(position(:,i),(/(0._rel_kind,j=1,n_dim)/),mass(i),&
                     &4._rel_kind*pi*self%dust_prop%bol_mass_emissivity(temperature(i)))
               end if
            end do
            h_present=.false.

      end select

      write(*,"(A)") "initialise tree"
      allocate(self%sph_tree)
      call self%sph_tree%initialise(self%particle_array,self%sph_kernel,self%sim_params%sph_eta,h_present)

      if (self%sim_params%point_sources) then

         write(*,"(A)") "read in point sources"
         select case (trim(self%sim_params%point_type))

            case ("bb")

               call read_in_point_sources_3d(position,luminosity=luminosity,temperature=temperature,&
                  &file_name=self%sim_params%sim_star_file)
               allocate(source_point_bb::self%point_source_array(size(position,2)))

               write(*,"(A)") "initialise point sources"
               do i=1,size(self%point_source_array)
                  call self%point_source_array(i)%initialise(position(:,i),(/(0._rel_kind,j=1,n_dim)/),temperature=temperature(i),&
                     &luminosity=luminosity(i),beta=self%sim_params%point_beta,wavelength_min=self%sim_params%point_lambda_min,&
                     &wavelength_max=self%sim_params%point_lambda_max,n_wavelength=self%sim_params%point_n_lambda)
               end do

            case ("ms_star")

               call read_in_point_sources_3d(position,mass=mass,age=age,metallicity=metallicity,&
                  &file_name=self%sim_params%sim_star_file)
               allocate(source_point_ms_star::self%point_source_array(size(position,2)))

               write(*,"(A)") "initialise point sources"
               do i=1,size(self%point_source_array)
                  call self%point_source_array(i)%initialise(position(:,i),(/(0._rel_kind,j=1,n_dim)/),mass=mass(i),&
                     &age=age(i),metallicity=metallicity(i))
               end do

            case ("custom")

               call read_in_point_sources_3d(position,file_name=self%sim_params%sim_star_file)
               allocate(source_point_custom::self%point_source_array(size(position,2)))

               write(*,"(A)") "initialise point sources"
               open(2,file=trim(self%sim_params%point_file_list))
               read(2,*) file_name
               do i=1,size(self%point_source_array)
                  call self%point_source_array(i)%initialise(position(:,i),(/(0._rel_kind,j=1,n_dim)/),luminosity_file=file_name)
               end do
               close(2)

         end select

         ! set sublimation radii for point sources
         call self%point_source_array%set_sublimation_radius(self%dust_prop,self%sim_params%dust_sub_t,&
            &self%sim_params%point_min_sub_radius)

      else
         self%point_source_array=>null()
      end if

      write(*,"(A)") "initialise radiation field"
      if (self%sim_params%ext_rf) then
         select case (trim(self%sim_params%ext_rf_type))

            case ("ps05")

               allocate(source_external_ps05::self%isrf_prop)
               call self%isrf_prop%initialise(self%sph_tree%com,(/(0._rel_kind,i=1,n_dim)/),&
                  &radius=0.5_rel_kind*self%sph_tree%max_length,gal_r=self%sim_params%ext_rf_gal_r,&
                  &gal_z=self%sim_params%ext_rf_gal_r,add_cmb=self%sim_params%ext_rf_cmb)

            case ("bb")

               allocate(source_external_bb::self%isrf_prop)
               call self%isrf_prop%initialise(self%sph_tree%com,(/(0._rel_kind,i=1,n_dim)/),&
                  &radius=0.5_rel_kind*self%sph_tree%max_length,wavelength_min=self%sim_params%ext_rf_lambda_min,&
                  &wavelength_max=self%sim_params%ext_rf_lambda_max,n_wavelength=self%sim_params%ext_rf_n_lambda,&
                  &temperature=self%sim_params%ext_rf_t,dilution=self%sim_params%ext_rf_d,&
                  &add_cmb=self%sim_params%ext_rf_cmb)

            case default

               write(*,"(A)") "Unknown radiation field type:",self%sim_params%ext_rf_type
               stop

         end select

      else
         self%isrf_prop=>null()
      end if

      return

   end subroutine

   pure subroutine destroy(self)

      ! argument declarations
      class(simulation),intent(inout) :: self                     ! simulation object

      call self%particle_array%destroy()
      call self%sph_tree%destroy()
      call self%dust_prop%destroy()
      if (associated(self%point_source_array)) call self%point_source_array%destroy()
      if (associated(self%isrf_prop)) call self%isrf_prop%destroy()
      if (associated(self%intensity_cube)) call self%intensity_cube%destroy()

      deallocate(self%particle_array)
      deallocate(self%sph_tree)
      deallocate(self%sim_params)
      deallocate(self%dust_prop)
      if (associated(self%point_source_array)) deallocate(self%point_source_array)
      if (associated(self%isrf_prop)) deallocate(self%isrf_prop)
      if (associated(self%intensity_cube)) deallocate(self%intensity_cube)

      return

   end subroutine

   ! perform temperature calculation
   subroutine perform_iteration(self,iteration)

      ! argument declarations
      class(simulation),intent(inout) :: self                     ! simulation object
      integer(kind=int_kind),intent(in) :: iteration              ! iteration number

      ! variable declarations
      integer(kind=int_kind) :: i,j,k                             ! counter
      integer(kind=int_kind) :: n_threads                         ! number of available threads
      integer(kind=int_kind) :: n_packets                         ! number of packets for this iteration
      integer(kind=int_kind) :: n_packets_source                  ! number of packets for individual source
      real(kind=rel_kind) :: total_luminosity                     ! total luminosity of all sources
      real(kind=rel_kind) :: luminosity_chunk                     ! luminosity chunk for each source
      real(kind=rel_kind) :: position(n_dim)                      ! position vector
      real(kind=rel_kind) :: direction(n_dim)                     ! direction vector
      real(kind=rel_kind) :: wavelength                           ! wavelength
      real(kind=rel_kind) :: tau                                  ! optical depth travelled by luminosity packet
      real(kind=rel_kind) :: bg_mean_tau                          ! mean optical depth travelled by background packets
      real(kind=rel_kind) :: source_luminosity                    ! total luminosity of sources
      real(kind=rel_kind) :: packet_luminosity                    ! packet energy deposition rate
      real(kind=rel_kind) :: particle_luminosity                  ! particle energy emission rate
      real(kind=rel_kind),allocatable :: position_array(:,:)      ! array of positions
      real(kind=rel_kind),allocatable :: ps_mean_tau(:)           ! mean optical depth travelled by point source packets
      character(kind=chr_kind,len=string_length) :: output_file   ! output file

      ! get number of threads
      n_threads=omp_get_num_procs()

      ! reset scattered light bins
      call self%particle_array%reset_a_dot()

      ! initialise luminosity packets (one per thread)
      write(*,"(A)") "initialise luminosity packets"
      allocate(self%lum_packet_array(n_threads))
      do i=1,n_threads
         call self%lum_packet_array(i)%initialise(self%sph_kernel,self%sph_tree,self%dust_prop)
      end do

      write(*,"(A)") "set sublimation fraction"
      ! set sublimation fraction
      ! by default, set to 1
      self%particle_array%f_sub=1._rel_kind
      select case (trim(self%sim_params%dust_sub_type))

         case ("variable")

            do i=1,size(self%particle_array)
               self%particle_array(i)%f_sub=self%sph_tree%root_node%sph_gather_f_sub_variable(self%sph_kernel,&
                  &self%particle_array(i),4._rel_kind*pi*self%dust_prop%bol_mass_emissivity(self%sim_params%dust_sub_t))
            end do

         case ("fixed")

            if (associated(self%point_source_array)) then
               do i=1,size(self%particle_array)
                  self%particle_array(i)%f_sub=self%sph_tree%root_node%sph_gather_f_sub_fixed(self%sph_kernel,&
                     &self%particle_array(i),self%point_source_array)
               end do
            end if

      end select

      if (associated(self%point_source_array)) then

         allocate(ps_mean_tau(size(self%point_source_array)))
         ps_mean_tau=0._rel_kind
         total_luminosity=sum(self%point_source_array%luminosity)

         write(*,"(A)") "follow packets from point sources"

         do i=1,size(self%point_source_array)

            if (self%sim_params%sim_n_it>1) then
               n_packets=int(exp(log(self%sim_params%sim_n_packet_multiplier*self%sim_params%sim_n_packet_point)+&
                  &(real(iteration-1,rel_kind)/real(self%sim_params%sim_n_it-1,rel_kind))*&
                  &log(1._rel_kind/self%sim_params%sim_n_packet_multiplier)))
            else
               n_packets=self%sim_params%sim_n_packet_point
            end if

            if (self%sim_params%sim_equal_packets_per_point) then
               n_packets_source=n_packets/size(self%point_source_array)
            else
               n_packets_source=&
                  &int(real(n_packets,rel_kind)*self%point_source_array(i)%luminosity/total_luminosity)
            end if

            n_packets_source=max(1,n_packets_source)

            luminosity_chunk=self%point_source_array(i)%luminosity/real(n_packets_source,rel_kind)
            k=0

            !$omp parallel do num_threads(n_threads) default(shared) private(j,tau)
               do j=1,n_packets_source

                  tau=0._rel_kind
                  call self%lum_packet_array(omp_get_thread_num()+1)%follow(self%point_source_array(i)%position,&
                     &self%point_source_array(i)%random_direction(),self%point_source_array(i)%velocity,&
                     &self%point_source_array(i)%random_wavelength(),luminosity_chunk,&
                     &self%sim_params%sim_mrw,self%sim_params%sim_mrw_gamma,tau)

                  call atomic_integer_add(k,1)
                  call atomic_real_add(ps_mean_tau(i),tau)

                  if (mod(k,min(max(n_packets/100,1),n_packets_source))==0) &
                     &write(*,"(A,I0,A,I0,A,I0,A,I0,A,I0,A,I0)") &
                     &"iteration ",iteration," of ",self%sim_params%sim_n_it,&
                     &", star ",i," of ",size(self%point_source_array),&
                     &", packet ",k," of ",n_packets_source
               end do
            !$omp end parallel do
            ps_mean_tau(i)=ps_mean_tau(i)/real(n_packets_source,rel_kind)

         end do

      end if

      if (associated(self%isrf_prop)) then

         bg_mean_tau=0._rel_kind

         write(*,"(A)") "follow packets from external radiation field"
         if (self%sim_params%sim_n_it>1) then
            n_packets=int(exp(log(self%sim_params%sim_n_packet_multiplier*self%sim_params%sim_n_packet_external)+&
               &(real(iteration-1,rel_kind)/real(self%sim_params%sim_n_it-1,rel_kind))*&
               &log(1._rel_kind/self%sim_params%sim_n_packet_multiplier)))
         else
            n_packets=self%sim_params%sim_n_packet_external
         end if
         luminosity_chunk=self%isrf_prop%luminosity/real(n_packets,rel_kind)
         k=0

         !$omp parallel do num_threads(n_threads) default(shared) private(j,position,direction,wavelength,tau)

            do j=1,n_packets
               position=self%isrf_prop%random_position()
               direction=self%isrf_prop%random_direction(position)
               wavelength=self%isrf_prop%random_wavelength()
               tau=0._rel_kind
               call self%lum_packet_array(omp_get_thread_num()+1)%follow(position,&
                  &direction,self%isrf_prop%velocity,&
                  &wavelength,luminosity_chunk,&
                  &self%sim_params%sim_mrw,self%sim_params%sim_mrw_gamma,tau)

               call atomic_integer_add(k,1)
               call atomic_real_add(bg_mean_tau,tau)

               if (mod(k,min(max(n_packets/100,1),n_packets))==0) &
                  &write(*,"(A,I0,A,I0,A,I0,A,I0)") &
                     &"iteration ",iteration," of ",self%sim_params%sim_n_it,&
                     &", external rad field, packet ",k," of ",n_packets

            end do

         !$omp end parallel do
         bg_mean_tau=bg_mean_tau/real(n_packets,rel_kind)

      end if

      ! normalise absorption rate
      call self%particle_array%normalise_a()

      ! write out packet travel stats
      if (associated(self%point_source_array)) then
         do i=1,size(ps_mean_tau)
            write(*,"(A,I0,A,E25.17)") "star ",i," mean escape tau: ",ps_mean_tau(i)
         end do
      end if

      if (associated(self%isrf_prop)) write(*,"(A,E25.17)") "external rad field escape tau:   ",bg_mean_tau

      ! calculate packet energy deposition rate
      source_luminosity=0._rel_kind
      packet_luminosity=0._rel_kind
      if (associated(self%point_source_array)) then
         packet_luminosity=packet_luminosity+sum(self%point_source_array%luminosity*ps_mean_tau)
         source_luminosity=source_luminosity+sum(self%point_source_array%luminosity)
      end if

      if (associated(self%isrf_prop)) then
         packet_luminosity=packet_luminosity+self%isrf_prop%luminosity*bg_mean_tau
         source_luminosity=source_luminosity+self%isrf_prop%luminosity
      end if

      ! calculate particle luminosity
      particle_luminosity=0._rel_kind
      do i=1,size(self%particle_array)
         particle_luminosity=particle_luminosity+&
            &self%particle_array(i)%m*self%particle_array(i)%f_sub*self%particle_array(i)%a_dot
            if (associated(self%particle_array(i)%lambda_array)) particle_luminosity=particle_luminosity+&
               &self%particle_array(i)%m*self%particle_array(i)%f_sub*&
               &sum(self%particle_array(i)%a_dot_scatter_array*(self%particle_array(i)%lambda_array(2:)-&
               &self%particle_array(i)%lambda_array(:size(self%particle_array(i)%lambda_array)-1)))
      end do

      write(*,"(A,E25.17)") "source luminosity:   ",source_luminosity
      write(*,"(A,E25.17)") "packet luminosity:   ",packet_luminosity
      write(*,"(A,E25.17)") "particle luminosity: ",particle_luminosity

      ! write out particles
      write(*,"(A)") "write out particles"
      write(output_file,"(A,I2.2,A)") trim(self%sim_params%sim_id)//"/iteration_",iteration,".dat"
      allocate(position_array(n_dim,size(self%particle_array)))
      do i=1,size(self%particle_array)
         position_array(:,i)=self%particle_array(i)%r
      end do
      call write_out_sph_particles_3d(position_array,self%particle_array%m,&
         &self%dust_prop%dust_temperature(self%particle_array%a_dot),self%particle_array%rho,self%particle_array%h,output_file)

      call self%write_particles_bin(iteration)

      ! deallocate luminosity packets
      call self%lum_packet_array%destroy()
      deallocate(self%lum_packet_array)

      return

   end subroutine

   subroutine make_datacube(self)

      ! argument declarations
      class(simulation),intent(inout) :: self                     ! simulation object

      ! allocate datacube
      allocate(self%intensity_cube)

      ! make datacube
      write(*,"(A)") "making datacube"

      call self%intensity_cube%initialise(self%sim_params,self%isrf_prop,self%point_source_array,&
         &self%sph_kernel,self%sph_tree,self%dust_prop)

      if (self%sim_params%datacube_convolve) then
         write(*,"(A)") "convolving datacube with psf"
         call self%intensity_cube%psf_convolve()
      end if

      write(*,"(A)") "write out datacube"

      call self%intensity_cube%write_out()

      return

   end subroutine

   ! write out particles in binary format
   subroutine write_particles_bin(self,iteration)

      ! character declarations
      class(simulation),intent(in) :: self                     ! simulation object
      integer(kind=int_kind),intent(in) :: iteration           ! iteration number

      ! variable delcarations
      integer(kind=int_kind) :: i                              ! counter
      character(kind=chr_kind,len=string_length) :: file_name  ! file name

      ! open file
      write(file_name,"(A,I2.2,A)") trim(self%sim_params%sim_id)//"/iteration_",iteration,".bin"
      open(1,file=trim(file_name),form="unformatted")

      ! write number of particles to file
      write(1) size(self%particle_array,kind=int_kind)

      do i=1,size(self%particle_array)

         ! write out particle properties

         write(1) self%particle_array(i)%r
         write(1) self%particle_array(i)%v
         write(1) self%particle_array(i)%m
         write(1) self%particle_array(i)%h
         write(1) self%particle_array(i)%rho
         write(1) self%particle_array(i)%a_dot
         write(1) self%particle_array(i)%f_sub
         if (associated(self%particle_array(i)%lambda_array)) then
            write(1) size(self%particle_array(i)%lambda_array,kind=int_kind)
            write(1) self%particle_array(i)%lambda_array
            write(1) self%particle_array(i)%a_dot_scatter_array
         else
            write(1) 0_int_kind
         end if

      end do

      close(1)

      return

   end subroutine

   ! write out particles in binary format
   subroutine read_particles_bin(self)

      ! character declarations
      class(simulation),intent(inout) :: self                  ! simulation object

      ! variable delcarations
      integer(kind=int_kind) :: i                              ! counter
      integer(kind=int_kind) :: n_particles                    ! number of particles
      integer(kind=int_kind) :: n_array                        ! arbitrary array length

      ! open file
      open(1,file=trim(self%sim_params%sim_cloud_file),form="unformatted")

      ! write number of particles to file
      read(1) n_particles
      allocate(self%particle_array(n_particles))

      do i=1,size(self%particle_array)

         ! write out particle properties

         read(1) self%particle_array(i)%r
         read(1) self%particle_array(i)%v
         read(1) self%particle_array(i)%m
         read(1) self%particle_array(i)%h
         read(1) self%particle_array(i)%rho
         read(1) self%particle_array(i)%a_dot
         read(1) self%particle_array(i)%f_sub
         read(1) n_array
         if (n_array>0) then
            allocate(self%particle_array(i)%lambda_array(n_array))
            allocate(self%particle_array(i)%a_dot_scatter_array(n_array-1))
            read(1) self%particle_array(i)%lambda_array
            read(1) self%particle_array(i)%a_dot_scatter_array
         else
            self%particle_array(i)%lambda_array=>null()
            self%particle_array(i)%a_dot_scatter_array=>null()
         end if

         ! calculate some derived quantities
         self%particle_array(i)%inv_h=1._rel_kind/self%particle_array(i)%h
         self%particle_array(i)%inv_rho=1._rel_kind/self%particle_array(i)%rho
         self%particle_array(i)%a_dot_new=0._rel_kind

      end do

      close(1)

      return

   end subroutine

   ! cull particle ensemble to sphere
   subroutine source_extract(self,centre,radius)

      ! argument declarations
      class(simulation),intent(inout) :: self            ! simulation object
      real(kind=rel_kind),intent(in) :: centre(n_dim)    ! source centre
      real(kind=rel_kind),intent(in) :: radius           ! radius

      ! variable declarations
      type(particle),pointer :: temp_particle_array(:)   ! temporary particle array
      integer(kind=int_kind) :: i                        ! counter
      logical(kind=log_kind) :: array_mask(size(self%particle_array))   ! true if particles are within sphere

      ! set mask
      do i=1,size(self%particle_array)
         array_mask(i)=sum((self%particle_array(i)%r-centre)**2)<=radius**2
         self%particle_array(i)%r=self%particle_array(i)%r-centre
      end do

      ! copy live particles
      allocate(temp_particle_array(count(array_mask)))

      temp_particle_array=pack(self%particle_array,array_mask)

      self%particle_array=>temp_particle_array

      ! reset particle tree
      call self%sph_tree%destroy()
      call self%sph_tree%initialise(&
         &self%particle_array,self%sph_kernel,self%sim_params%sph_eta,.true._log_kind)

      ! set luminosities of sources outside sphere to zero
      do i=1,size(self%point_source_array)
         array_mask(i)=sum((self%point_source_array(i)%position-centre)**2)<=radius**2
         self%point_source_array(i)%position=self%point_source_array(i)%position-centre
         if (.not.array_mask(i)) self%point_source_array(i)%luminosity=0._rel_kind
      end do

      return

   end subroutine

end module


      subroutine foo(d)
      parameter (n=1,m=2,k=3)
      character *17 xx
      integer a(1:m*n,-n+2), b(n*(m+k)*2,k+1)
      integer d(*,*)
      common /xx/ xx(101)
      character*1 str, str2*(n), str3*(*), str4*(m+n*k)
      external str
      dimension str2(k*m)
      character *10 c(len('won''t'))
      a(n,1) = a(1,m)
      c(2)(a(ichar('A'//'hello'), ichar(char(40))+3)
     $     :ichar(str("won""t",40) (1:1))) = 'x'
      end


      subroutine set_pint_contract2(flist,ansatz,
     &     idx_opsin,nopsin,
     &     op_info,orb_info)
*----------------------------------------------------------------------*
*     Routine used to form the various contractions necessary to form
*     the CABS approximation to the P-intermediate.
*     Modified formulation: August 2008
*----------------------------------------------------------------------*

      implicit none

      integer, parameter ::
     &     ntest = 00

      include 'stdunit.h'
      include 'opdim.h'
      include 'mdef_operator_info.h'
      include 'mdef_formula_info.h'
      include 'def_contraction.h'
      include 'def_formula_item.h'
      include 'def_orbinf.h'

      type(formula_item), target, intent(inout) ::
     &     flist
      integer, intent(in) ::
     &     ansatz, nopsin, idx_opsin(nopsin)
      type(operator_info), intent(inout) ::
     &     op_info
      type(orbinf), intent(in) ::
     &     orb_info

      integer ::
     &     idx, idx_shape, idx_temp, ndef, idx_prj, nj_intm
      character ::
     &     name*(form_maxlen_label*2)

      type(operator), pointer ::
     &     op
      type(formula_item), pointer ::
     &     form_pnt

      integer, external ::
     &     idx_oplist2

      if(ntest.ge.100)then
        write(lulog,*)'P-Intermediate Contraction'
        write(lulog,*)'Constituent operators: '
        do idx = 1, nopsin
          write(lulog,*)trim(op_info%op_arr(idx_opsin(idx))%op%name)
        enddo
      endif

      if(ansatz.ne.3)call quit(1,'set_pint_contract2','wrong ansatz')

      ! Get indices of input operators.
      idx_shape = idx_opsin(1)
      nj_intm = op_info%op_arr(idx_opsin(1))%op%njoined

      ! Point to the formula and move to the end of the list.
      form_pnt => flist
      do while(associated(form_pnt%next))
        form_pnt => form_pnt%next
      enddo

      ! Add the F2G part:
      call set_1contrib(form_pnt,1d0,6,
     &     idx_shape,idx_opsin,nopsin,op_info)

      ! Move to the end of the list.
      do while(associated(form_pnt%next))
        form_pnt => form_pnt%next
      enddo

      if (nj_intm.eq.1) then
      ! Add the (FG) contracted with F part.
        idx_prj = 8
        call expand_op_product2(form_pnt,idx_shape,
     &     -1d0,5,3,
     &     (/idx_shape,-idx_opsin(2),
     &       idx_opsin(4),idx_opsin(4),idx_shape/),
     &     (/        1,            2,
     &                  3,           3,        1/),
     &     -1,-1,
     &     0,0,
     &     0,0,
     &     (/2,3,2,idx_prj/),1,
     &     .false.,op_info)

      ! Move to the end of the list.
        do while(associated(form_pnt%next))
          form_pnt => form_pnt%next
        enddo

      ! Add the (FG) contracted with F part.
        idx_prj = 10
        call expand_op_product2(form_pnt,idx_shape,
     &     -1d0,5,3,
     &     (/idx_shape,-idx_opsin(2),
     &       idx_opsin(4),idx_opsin(4),idx_shape/),
     &     (/        1,            2,
     &                  3,           3,        1/),
     &     -1,-1,
     &     0,0,
     &     0,0,
     &     (/2,3,2,idx_prj/),1,
     &     .false.,op_info)

      ! Move to the end of the list.
        do while(associated(form_pnt%next))
          form_pnt => form_pnt%next
        enddo

      ! Add the V^{p"m}_{kl}.F^{ij}_{p"m}
c      idx_prj = 9
        if(ansatz.gt.1) idx_prj = 10
        call expand_op_product2(form_pnt,idx_shape,
     &     -1d0,4,3,
     &     (/idx_shape,-idx_opsin(5),idx_opsin(2),idx_shape/),
     &     (/        1,            2,           3,        1/),
     &     -1,-1,
     &     0,0,
     &     0,0,
     &     (/2,3,2,idx_prj/),1,
     &     .false.,op_info)

        ! Move to the end of the list.
        do while(associated(form_pnt%next))
          form_pnt => form_pnt%next
        enddo

      ! Add the V^{pq}_{kl}.F^{ij}_{pq}
c      idx_prj = 9
        if(ansatz.gt.1) idx_prj = 8
        call expand_op_product2(form_pnt,idx_shape,
     &     -1d0,4,3,
     &     (/idx_shape,-idx_opsin(5),idx_opsin(2),idx_shape/),
     &     (/        1,            2,           3,        1/),
     &     -1,-1,
     &     0,0,
     &     0,0,
     &     (/2,3,2,idx_prj/),1,
     &     .false.,op_info)
      else
      ! Add the (FG) contracted with F part.
        idx_prj = 8
        call expand_op_product2(form_pnt,idx_shape,
     &     -1d0,7,3,
     &     (/idx_shape,-idx_opsin(2),idx_shape,idx_shape,
     &       idx_opsin(4),idx_opsin(4),idx_shape/),
     &     (/        1,            2,        1,        1,
     &                  3,           3,        1/),
     &     -1,-1,
     &     0,0,
     &     0,0,
     &     (/2,5,2,idx_prj/),1,
     &     .false.,op_info)

      ! Move to the end of the list.
        do while(associated(form_pnt%next))
          form_pnt => form_pnt%next
        enddo

      ! Add the (FG) contracted with F part.
        idx_prj = 10
        call expand_op_product2(form_pnt,idx_shape,
     &     -1d0,7,3,
     &     (/idx_shape,-idx_opsin(2),idx_shape,idx_shape,
     &       idx_opsin(4),idx_opsin(4),idx_shape/),
     &     (/        1,            2,        1,        1,
     &                  3,           3,        1/),
     &     -1,-1,
     &     0,0,
     &     0,0,
     &     (/2,5,2,idx_prj/),1,
     &     .false.,op_info)

      ! Move to the end of the list.
        do while(associated(form_pnt%next))
          form_pnt => form_pnt%next
        enddo

      ! Add the V^{p"m}_{kl}.F^{ij}_{p"m}
c      idx_prj = 9
        if(ansatz.gt.1) idx_prj = 10
        call expand_op_product2(form_pnt,idx_shape,
     &     -1d0,7,3,
     &     (/idx_shape,-idx_opsin(5),idx_shape,idx_shape,
     &       -idx_opsin(5),idx_opsin(2),idx_shape/),
     &     (/        1,            2,        1,        1,
     &                   2,           3,        1/),
     &     -1,-1,
     &     0,0,
     &     0,0,
     &     (/5,6,2,idx_prj/),1,
     &     .false.,op_info)

        ! Move to the end of the list.
        do while(associated(form_pnt%next))
          form_pnt => form_pnt%next
        enddo

      ! Add the V^{pq}_{kl}.F^{ij}_{pq}
c      idx_prj = 9
        if(ansatz.gt.1) idx_prj = 8
        call expand_op_product2(form_pnt,idx_shape,
     &     -1d0,7,3,
     &     (/idx_shape,-idx_opsin(5),idx_shape,idx_shape,
     &       -idx_opsin(5),idx_opsin(2),idx_shape/),
     &     (/        1,            2,        1,        1,
     &                   2,           3,        1/),
     &     -1,-1,
     &     0,0,
     &     0,0,
     &     (/5,6,2,idx_prj/),1,
     &     .false.,op_info)
      end if

      if(ntest.ge.100)then
        write(lulog,*)'Final formula: P-Int.'
        call print_form_list(lulog,flist,op_info)
      endif

      return
      end


SUBROUTINE pgsblk(bit,itmx,cnvf,phi,phiold,psii,psio,oldpsi)

!-------------------------------------------------------------
!
!  Parallel Gauss-Seidel  -  Black
!
!  Check the convergence for all sub-domains in the system
!   If the system is not converged pass the black angular
!   flux data back to the red sub-domains. Update angular flux
!   at reflective edges.
!  If convergence is met, say so and exit the PGS loop
!
!-------------------------------------------------------------

USE invar
IMPLICIT NONE
INTEGER, INTENT(IN) :: bit, itmx
INTEGER, INTENT(OUT) :: cnvf
INTEGER :: i, j, k, i1, sdi, sdii, sdio, reset
INTEGER :: myztag, myytag, myxtag, znxt, zprv, ynxt, yprv, xnxt, xprv
INTEGER :: znxtag, zprtag, ynxtag, yprtag, xnxtag, xprtag, neqz, neqy, neqx, ieq, ierr
INTEGER, DIMENSION(3) :: istat
REAL*8 :: df, dfmx
REAL*8, DIMENSION(nsdp) :: dfsd
REAL*8, DIMENSION(neq,nsdp), INTENT(IN) :: phi
REAL*8, DIMENSION(neq,nsdp), INTENT(OUT) :: phiold
REAL*8, DIMENSION(bcs,8,nsdp), INTENT(INOUT) :: psii, oldpsi
REAL*8, DIMENSION(bcs,8,nsdp), INTENT(OUT) :: psio

INCLUDE 'mpif.h'

! Set up some tags
myztag = 100 + zrank
myytag = 100 + yrank
myxtag = 100 + xrank
znxt = zrank + 1
znxtag = 100 + znxt
zprv = zrank - 1
zprtag = 100 + zprv
ynxt = yrank + 1
ynxtag = 100 + ynxt
yprv = yrank - 1
yprtag = 100 + yprv
xnxt = xrank + 1
xnxtag = 100 + xnxt
xprv = xrank - 1
xprtag = 100 + xprv

! Check convergence of scalar flux
! Each process checks its own fluxes, then perform MPI Reduction
IF (MINVAL(ABS(phiold)) >= tolr) THEN
   DO i = 1, nsdp
      dfsd(i) = MAXVAL(ABS((phi(:,i) - phiold(:,i))/phiold(:,i)))
   END DO
ELSE
   DO i = 1, nsdp
      dfsd(i) = MAXVAL(ABS(phi(:,i) - phiold(:,i)))
   END DO
END IF
df = MAXVAL(dfsd)

! Get the overall max difference
CALL MPI_ALLREDUCE(df,dfmx,1,MPI_DOUBLE_PRECISION,MPI_MAX,allcomm,ierr)

! Only print iteration information if requested
IF (irank == root) THEN
   IF (itp == 1) WRITE(8,111) bit, dfmx
END IF
111 FORMAT(2X,'It ', I5,' Dfmx ', ES11.3)

! Set cnvf to 1 if converged
IF (dfmx < err) cnvf = 1

! Set the previous iterate of the flux equal to the current
phiold = phi

! Get a copy of the old psii for computing the residual for red sub-domains
IF (bit == itmx) THEN
   sdi = 0
   DO k = 1, nzsd
      DO j = 1, nysd
         i1 = MOD(j+k,2) + 1
         DO i = i1, nxsd, 2
            sdi = (k-1)*nysd*nxsd + (j-1)*nxsd + i
            oldpsi(:,:,sdi) = psii(:,:,sdi)
         END DO
      END DO
   END DO
END IF

! Send psio to neighbors. Inner blocks have 24 communications (8 octants to
! 3 neigbors each). Then set the incoming data into psii

!--------------------------------------------------------------------------------
! Perform the copies 'up z' from black to red, internally
neqz = apo*xys
DO k = 1, nzsd-1
   DO j = 1, nysd
      i1 = MOD(j+k+1,2) + 1
      DO i = i1, nxsd, 2
         sdio = (k-1)*nysd*nxsd + (j-1)*nxsd + i
         sdii = sdio + nysd*nxsd
         psii(1:neqz,1,sdii) = psio(1:neqz,1,sdio)
         psii(1:neqz,2,sdii) = psio(1:neqz,2,sdio)
         psii(1:neqz,3,sdii) = psio(1:neqz,3,sdio)
         psii(1:neqz,4,sdii) = psio(1:neqz,4,sdio)
      END DO
   END DO
END DO
! Now perform the sends 'up z' from black to red across processors
reset = (nzsd-1)*nysd*nxsd
DO j = 1, nysd
   i1 = MOD(j+1,2) + 1
   DO i = i1, nxsd, 2
      sdii = (j-1)*nxsd + i
      sdio = sdii + reset
      IF (zrank/=npz-1) CALL MPI_SEND(psio(1,1,sdio),neqz,MPI_DOUBLE_PRECISION,znxt,znxtag,zcomm,ierr)
      IF (zrank/=0) CALL MPI_RECV(psii(1,1,sdii),neqz,MPI_DOUBLE_PRECISION,zprv,myztag,zcomm,istat,ierr)
      IF (zrank/=npz-1) CALL MPI_SEND(psio(1,2,sdio),neqz,MPI_DOUBLE_PRECISION,znxt,znxtag,zcomm,ierr)
      IF (zrank/=0) CALL MPI_RECV(psii(1,2,sdii),neqz,MPI_DOUBLE_PRECISION,zprv,myztag,zcomm,istat,ierr)
      IF (zrank/=npz-1) CALL MPI_SEND(psio(1,3,sdio),neqz,MPI_DOUBLE_PRECISION,znxt,znxtag,zcomm,ierr)
      IF (zrank/=0) CALL MPI_RECV(psii(1,3,sdii),neqz,MPI_DOUBLE_PRECISION,zprv,myztag,zcomm,istat,ierr)
      IF (zrank/=npz-1) CALL MPI_SEND(psio(1,4,sdio),neqz,MPI_DOUBLE_PRECISION,znxt,znxtag,zcomm,ierr)
      IF (zrank/=0) CALL MPI_RECV(psii(1,4,sdii),neqz,MPI_DOUBLE_PRECISION,zprv,myztag,zcomm,istat,ierr)
   END DO
END DO
! Perform the copies 'down z' from black to red, internally
DO k = 2, nzsd
   DO j = 1, nysd
      i1 = MOD(j+k+1,2) + 1
      DO i = i1, nxsd, 2
         sdio = (k-1)*nysd*nxsd + (j-1)*nxsd + i
         sdii = sdio - nysd*nxsd
         psii(1:neqz,5,sdii) = psio(1:neqz,5,sdio)
         psii(1:neqz,6,sdii) = psio(1:neqz,6,sdio)
         psii(1:neqz,7,sdii) = psio(1:neqz,7,sdio)
         psii(1:neqz,8,sdii) = psio(1:neqz,8,sdio)
      END DO
   END DO
END DO
! Now perform the sends 'down z' from black to red across processors
DO j = 1, nysd
   i1 = MOD(j,2) + 1
   DO i = i1, nxsd, 2
      sdio = (j-1)*nxsd + i
      sdii = sdio + reset
      IF (zrank/=0) CALL MPI_SEND(psio(1,5,sdio),neqz,MPI_DOUBLE_PRECISION,zprv,zprtag,zcomm,ierr)
      IF (zrank/=npz-1) CALL MPI_RECV(psii(1,5,sdii),neqz,MPI_DOUBLE_PRECISION,znxt,myztag,zcomm,istat,ierr)
      IF (zrank/=0) CALL MPI_SEND(psio(1,6,sdio),neqz,MPI_DOUBLE_PRECISION,zprv,zprtag,zcomm,ierr)
      IF (zrank/=npz-1) CALL MPI_RECV(psii(1,6,sdii),neqz,MPI_DOUBLE_PRECISION,znxt,myztag,zcomm,istat,ierr)
      IF (zrank/=0) CALL MPI_SEND(psio(1,7,sdio),neqz,MPI_DOUBLE_PRECISION,zprv,zprtag,zcomm,ierr)
      IF (zrank/=npz-1) CALL MPI_RECV(psii(1,7,sdii),neqz,MPI_DOUBLE_PRECISION,znxt,myztag,zcomm,istat,ierr)
      IF (zrank/=0) CALL MPI_SEND(psio(1,8,sdio),neqz,MPI_DOUBLE_PRECISION,zprv,zprtag,zcomm,ierr)
      IF (zrank/=npz-1) CALL MPI_RECV(psii(1,8,sdii),neqz,MPI_DOUBLE_PRECISION,znxt,myztag,zcomm,istat,ierr)
   END DO
END DO

!----------------------------------------------------------------------------------------------------
! Perform the copies 'up y' from black to red, internally
ieq = neqz + 1
neqy = apo*xzs
DO k = 1, nzsd
   DO j = 1, nysd-1
      i1 = MOD(j+k+1,2) + 1
      DO i = i1, nxsd, 2
         sdio = (k-1)*nysd*nxsd + (j-1)*nxsd + i
         sdii = sdio + nxsd
         psii(ieq:neqz+neqy,1,sdii) = psio(ieq:neqz+neqy,1,sdio)
         psii(ieq:neqz+neqy,2,sdii) = psio(ieq:neqz+neqy,2,sdio)
         psii(ieq:neqz+neqy,5,sdii) = psio(ieq:neqz+neqy,5,sdio)
         psii(ieq:neqz+neqy,6,sdii) = psio(ieq:neqz+neqy,6,sdio)
      END DO
   END DO
END DO
! Now perform the sends 'up y' from black to red across processors
reset = (nysd-1)*nxsd
DO k = 1, nzsd
   i1 = MOD(k+1,2) + 1
   DO i = i1, nxsd, 2
      sdii = (k-1)*nysd*nxsd + i
      sdio = sdii + reset
      IF (yrank/=npy-1) CALL MPI_SEND(psio(ieq,1,sdio),neqy,MPI_DOUBLE_PRECISION,ynxt,ynxtag,ycomm,ierr)
      IF (yrank/=0) CALL MPI_RECV(psii(ieq,1,sdii),neqy,MPI_DOUBLE_PRECISION,yprv,myytag,ycomm,istat,ierr)
      IF (yrank/=npy-1) CALL MPI_SEND(psio(ieq,2,sdio),neqy,MPI_DOUBLE_PRECISION,ynxt,ynxtag,ycomm,ierr)
      IF (yrank/=0) CALL MPI_RECV(psii(ieq,2,sdii),neqy,MPI_DOUBLE_PRECISION,yprv,myytag,ycomm,istat,ierr)
      IF (yrank/=npy-1) CALL MPI_SEND(psio(ieq,5,sdio),neqy,MPI_DOUBLE_PRECISION,ynxt,ynxtag,ycomm,ierr)
      IF (yrank/=0) CALL MPI_RECV(psii(ieq,5,sdii),neqy,MPI_DOUBLE_PRECISION,yprv,myytag,ycomm,istat,ierr)
      IF (yrank/=npy-1) CALL MPI_SEND(psio(ieq,6,sdio),neqy,MPI_DOUBLE_PRECISION,ynxt,ynxtag,ycomm,ierr)
      IF (yrank/=0) CALL MPI_RECV(psii(ieq,6,sdii),neqy,MPI_DOUBLE_PRECISION,yprv,myytag,ycomm,istat,ierr)
   END DO
END DO
! Perform the copies 'down y' from black to red, internally
DO k = 1, nzsd
   DO j = 2, nysd
      i1 = MOD(j+k+1,2) + 1
      DO i = i1, nxsd, 2
         sdio = (k-1)*nysd*nxsd + (j-1)*nxsd + i
         sdii = sdio - nxsd
         psii(ieq:neqz+neqy,3,sdii) = psio(ieq:neqz+neqy,3,sdio)
         psii(ieq:neqz+neqy,4,sdii) = psio(ieq:neqz+neqy,4,sdio)
         psii(ieq:neqz+neqy,7,sdii) = psio(ieq:neqz+neqy,7,sdio)
         psii(ieq:neqz+neqy,8,sdii) = psio(ieq:neqz+neqy,8,sdio)
      END DO
   END DO
END DO
! Now perform the sends 'down y' from black to red across processors
reset = (nysd-1)*nxsd
DO k = 1, nzsd
   i1 = MOD(k,2) + 1
   DO i = i1, nxsd, 2
      sdio = (k-1)*nysd*nxsd + i
      sdii = sdio + reset
      IF (yrank/=0) CALL MPI_SEND(psio(ieq,3,sdio),neqy,MPI_DOUBLE_PRECISION,yprv,yprtag,ycomm,ierr)
      IF (yrank/=npy-1) CALL MPI_RECV(psii(ieq,3,sdii),neqy,MPI_DOUBLE_PRECISION,ynxt,myytag,ycomm,istat,ierr)
      IF (yrank/=0) CALL MPI_SEND(psio(ieq,4,sdio),neqy,MPI_DOUBLE_PRECISION,yprv,yprtag,ycomm,ierr)
      IF (yrank/=npy-1) CALL MPI_RECV(psii(ieq,4,sdii),neqy,MPI_DOUBLE_PRECISION,ynxt,myytag,ycomm,istat,ierr)
      IF (yrank/=0) CALL MPI_SEND(psio(ieq,7,sdio),neqy,MPI_DOUBLE_PRECISION,yprv,yprtag,ycomm,ierr)
      IF (yrank/=npy-1) CALL MPI_RECV(psii(ieq,7,sdii),neqy,MPI_DOUBLE_PRECISION,ynxt,myytag,ycomm,istat,ierr)
      IF (yrank/=0) CALL MPI_SEND(psio(ieq,8,sdio),neqy,MPI_DOUBLE_PRECISION,yprv,yprtag,ycomm,ierr)
      IF (yrank/=npy-1) CALL MPI_RECV(psii(ieq,8,sdii),neqy,MPI_DOUBLE_PRECISION,ynxt,myytag,ycomm,istat,ierr)
   END DO
END DO

!--------------------------------------------------------------------------------------------
! Perform the copies 'up x' from black to red, internally
ieq = ieq + neqy
neqx = apo*yzs
DO k = 1, nzsd
   DO j = 1, nysd
      i1 = MOD(j+k+1,2) + 1
      DO i = i1, nxsd-1, 2
         sdio = (k-1)*nysd*nxsd + (j-1)*nxsd + i
         sdii = sdio + 1
         psii(ieq:bcs,1,sdii) = psio(ieq:bcs,1,sdio)
         psii(ieq:bcs,4,sdii) = psio(ieq:bcs,4,sdio)
         psii(ieq:bcs,5,sdii) = psio(ieq:bcs,5,sdio)
         psii(ieq:bcs,8,sdii) = psio(ieq:bcs,8,sdio)
      END DO
   END DO
END DO
! Now perform the sends 'up x' from black to red across processors
reset = nxsd-1
DO k = 1, nzsd
   i1 = MOD(k+1,2) + 1
   DO j = i1, nysd, 2
      sdio = (k-1)*nysd*nxsd + j*nxsd
      sdii = sdio - reset
      IF (xrank/=npx-1) CALL MPI_SEND(psio(ieq,1,sdio),neqx,MPI_DOUBLE_PRECISION,xnxt,xnxtag,xcomm,ierr)
      IF (xrank/=0) CALL MPI_RECV(psii(ieq,1,sdii),neqx,MPI_DOUBLE_PRECISION,xprv,myxtag,xcomm,istat,ierr)
      IF (xrank/=npx-1) CALL MPI_SEND(psio(ieq,4,sdio),neqx,MPI_DOUBLE_PRECISION,xnxt,xnxtag,xcomm,ierr)
      IF (xrank/=0) CALL MPI_RECV(psii(ieq,4,sdii),neqx,MPI_DOUBLE_PRECISION,xprv,myxtag,xcomm,istat,ierr)
      IF (xrank/=npx-1) CALL MPI_SEND(psio(ieq,5,sdio),neqx,MPI_DOUBLE_PRECISION,xnxt,xnxtag,xcomm,ierr)
      IF (xrank/=0) CALL MPI_RECV(psii(ieq,5,sdii),neqx,MPI_DOUBLE_PRECISION,xprv,myxtag,xcomm,istat,ierr)
      IF (xrank/=npx-1) CALL MPI_SEND(psio(ieq,8,sdio),neqx,MPI_DOUBLE_PRECISION,xnxt,xnxtag,xcomm,ierr)
      IF (xrank/=0) CALL MPI_RECV(psii(ieq,8,sdii),neqx,MPI_DOUBLE_PRECISION,xprv,myxtag,xcomm,istat,ierr)
   END DO
END DO
! Perform the copies 'down x' from black to red, internally
DO k = 1, nzsd
   DO j = 1, nysd
      i1 = MOD(j+k,2) + 2
      DO i = i1, nxsd, 2
         sdio = (k-1)*nysd*nxsd + (j-1)*nxsd + i
         sdii = sdio - 1
         psii(ieq:bcs,2,sdii) = psio(ieq:bcs,2,sdio)
         psii(ieq:bcs,3,sdii) = psio(ieq:bcs,3,sdio)
         psii(ieq:bcs,6,sdii) = psio(ieq:bcs,6,sdio)
         psii(ieq:bcs,7,sdii) = psio(ieq:bcs,7,sdio)
      END DO
   END DO
END DO
! Now perform the sends 'down x' from black to red across processors
DO k = 1, nzsd
   i1 = MOD(k,2) + 1
   DO j = i1, nysd, 2
      sdii = (k-1)*nysd*nxsd + j*nxsd
      sdio = sdii - reset
      IF (xrank/=0) CALL MPI_SEND(psio(ieq,2,sdio),neqx,MPI_DOUBLE_PRECISION,xprv,xprtag,xcomm,ierr)
      IF (xrank/=npx-1) CALL MPI_RECV(psii(ieq,2,sdii),neqx,MPI_DOUBLE_PRECISION,xnxt,myxtag,xcomm,istat,ierr)
      IF (xrank/=0) CALL MPI_SEND(psio(ieq,3,sdio),neqx,MPI_DOUBLE_PRECISION,xprv,xprtag,xcomm,ierr)
      IF (xrank/=npx-1) CALL MPI_RECV(psii(ieq,3,sdii),neqx,MPI_DOUBLE_PRECISION,xnxt,myxtag,xcomm,istat,ierr)
      IF (xrank/=0) CALL MPI_SEND(psio(ieq,6,sdio),neqx,MPI_DOUBLE_PRECISION,xprv,xprtag,xcomm,ierr)
      IF (xrank/=npx-1) CALL MPI_RECV(psii(ieq,6,sdii),neqx,MPI_DOUBLE_PRECISION,xnxt,myxtag,xcomm,istat,ierr)
      IF (xrank/=0) CALL MPI_SEND(psio(ieq,7,sdio),neqx,MPI_DOUBLE_PRECISION,xprv,xprtag,xcomm,ierr)
      IF (xrank/=npx-1) CALL MPI_RECV(psii(ieq,7,sdii),neqx,MPI_DOUBLE_PRECISION,xnxt,myxtag,xcomm,istat,ierr)
   END DO
END DO

RETURN
END SUBROUTINE pgsblk


! Copyright (c) 2015-2021, the ELSI team.
! All rights reserved.
!
! This file is part of ELSI and is distributed under the BSD 3-clause license,
! which may be found in the LICENSE file in the ELSI root directory.

!>
!! Provide derived types used in ELSI. Most variables in this module can be set
!! by subroutines in ELSI_SET, and should be initialized with default values by
!! subroutines in ELSI_UTIL.
!!
module ELSI_DATATYPE

   use, intrinsic :: ISO_C_BINDING
   use ELSI_PRECISION, only: r8,i4
   use ELPA, only: elpa_t,elpa_autotune_t
   use F_PPEXSI_INTERFACE, only: f_ppexsi_options
   use FORTJSON, only: fjson_handle
   use NTPOLY, only: Permutation_t,Matrix_ps,SolverParameters_t,ProcessGrid_t

   implicit none

   private

   type, public :: elsi_basic_t

      ! IO
      integer(kind=i4) :: print_info
      integer(kind=i4) :: print_unit
      integer(kind=i4) :: print_json
      logical :: json_init = .false.

      ! MPI
      integer(kind=i4) :: myid
      integer(kind=i4) :: myid_all
      integer(kind=i4) :: n_procs
      integer(kind=i4) :: n_procs_all
      integer(kind=i4) :: comm
      integer(kind=i4) :: comm_all
      logical :: mpi_ready = .false.
      logical :: mpi_all_ready = .false.

      ! BLACS
      integer(kind=i4) :: blacs_ctxt
      integer(kind=i4) :: desc(9)
      integer(kind=i4) :: blk
      integer(kind=i4) :: n_prow
      integer(kind=i4) :: n_pcol
      integer(kind=i4) :: my_prow
      integer(kind=i4) :: my_pcol
      integer(kind=i4) :: n_lrow
      integer(kind=i4) :: n_lcol
      integer(kind=i4) :: nnz_l ! Local number of nonzeros
      logical :: blacs_ready = .false.

      ! Sparse matrix information (common)
      integer(kind=i4) :: nnz_g ! Global number of nonzeros
      integer(kind=i4) :: nnz_l_sp ! Local number of nonzeros
      integer(kind=i4) :: n_lcol_sp ! Local number of columns
      real(kind=r8) :: def0 ! Zero threshold

      ! Sparse matrix information (1D block)
      integer(kind=i4) :: nnz_l_sp1 ! Local number of nonzeros
      integer(kind=i4) :: n_lcol_sp1 ! Local number of columns
      logical :: pexsi_csc_ready = .false.

      ! Sparse matrix information (1D block-cyclic)
      integer(kind=i4) :: nnz_l_sp2 ! Local number of nonzeros
      integer(kind=i4) :: n_lcol_sp2 ! Local number of columns
      integer(kind=i4) :: blk_sp2
      logical :: siesta_csc_ready = .false.

      ! Sparse matrix information (generic)
      integer(kind=i4) :: nnz_l_sp3 ! Local number of nonzeros
      logical :: generic_coo_ready = .false.

   end type

   type, public :: elsi_param_t

      ! General info
      integer(kind=i4) :: solver
      integer(kind=i4) :: matrix_format
      integer(kind=i4) :: parallel_mode
      integer(kind=i4) :: n_calls ! Number of calls in this geometry step
      integer(kind=i4) :: n_calls_all ! Total number of calls

      ! Overlap
      logical :: save_ovlp
      logical :: unit_ovlp
      logical :: ill_ovlp
      logical :: ill_check
      real(kind=r8) :: ill_tol
      integer(kind=i4) :: n_good ! Number of non-ill-conditioned basis functions
      real(kind=r8) :: ovlp_ev_min
      real(kind=r8) :: ovlp_ev_max

      ! Physics
      real(kind=r8) :: n_electrons
      integer(kind=i4) :: n_basis
      integer(kind=i4) :: n_spins
      integer(kind=i4) :: n_kpts
      integer(kind=i4) :: n_states
      integer(kind=i4) :: n_states_solve
      integer(kind=i4) :: i_spin ! Local spin index
      integer(kind=i4) :: i_kpt ! Local k-point index
      real(kind=r8) :: i_wt ! Weight of local k-point
      real(kind=r8) :: spin_degen
      logical :: spin_is_set = .false.
      real(kind=r8) :: ebs ! Band structure energy
      real(kind=r8) :: energy_gap
      real(kind=r8) :: spectrum_width
      integer(kind=i4) :: dimensionality
      integer(kind=i4) :: extrapolation
      logical :: edm_ready = .false.
      logical :: eval_ready = .false.
      logical :: evec_ready = .false.
      logical :: occ_ready = .false.
      logical :: occ_non_aufbau = .false. ! occ for lowest excited state ( ...222000... => ...221100... )

      ! Chemical potential
      real(kind=r8) :: mu ! Fermi level
      real(kind=r8) :: ts ! Entropy
      integer(kind=i4) :: mu_scheme
      real(kind=r8) :: mu_width
      real(kind=r8) :: mu_tol
      integer(kind=i4) :: mu_max_steps
      integer(kind=i4) :: mu_mp_order

      ! Frozen core
      integer(kind=i4) :: fc_method
      integer(kind=i4) :: n_basis_c
      integer(kind=i4) :: n_basis_v
      integer(kind=i4) :: n_lrow_v
      integer(kind=i4) :: n_lcol_v
      integer(kind=i4) :: desc_v(9)
      logical :: fc_perm ! Permute matrices before solving

      ! Matrix redistribution
      integer(kind=i4) :: sparsity_mask
      logical :: first_blacs_to_ntpoly
      logical :: first_blacs_to_pexsi
      logical :: first_generic_to_blacs
      logical :: first_generic_to_ntpoly
      logical :: first_generic_to_pexsi
      logical :: first_siesta_to_blacs
      logical :: first_siesta_to_ntpoly
      logical :: first_siesta_to_pexsi
      logical :: first_sips_to_blacs
      logical :: first_sips_to_ntpoly

      ! ELPA
      integer(kind=i4) :: elpa_comm_row
      integer(kind=i4) :: elpa_comm_col
      integer(kind=i4) :: elpa_solver
      integer(kind=i4) :: elpa_n_single
      integer(kind=i4) :: elpa_gpu
      integer(kind=i4) :: elpa_autotune
      integer(kind=i4) :: elpa_n_lrow ! Non-ill-conditioned basis functions
      integer(kind=i4) :: elpa_n_lcol ! Non-ill-conditioned basis functions
      logical :: elpa_first
      logical :: elpa_started = .false.
      class(elpa_t), pointer :: elpa_aux => null()
      class(elpa_t), pointer :: elpa_solve => null()
      class(elpa_autotune_t), pointer :: elpa_tune => null()

      ! libOMM
      integer(kind=i4) :: omm_n_lrow
      integer(kind=i4) :: omm_n_elpa
      integer(kind=i4) :: omm_flavor
      integer(kind=i4) :: omm_desc(9)
      real(kind=r8) :: omm_tol
      logical :: omm_output
      logical :: omm_first
      logical :: omm_started = .false.

      ! PEXSI
      integer(kind=i4) :: pexsi_np_per_pole
      integer(kind=i4) :: pexsi_np_per_point
      integer(kind=i4) :: pexsi_my_prow
      integer(kind=i4) :: pexsi_my_pcol
      integer(kind=i4) :: pexsi_n_prow
      integer(kind=i4) :: pexsi_n_pcol
      integer(kind=i4) :: pexsi_my_point
      integer(kind=i4) :: pexsi_myid_point
      integer(kind=i4) :: pexsi_comm_intra_pole
      integer(kind=i4) :: pexsi_comm_inter_pole
      integer(kind=i4) :: pexsi_comm_inter_point
      real(kind=r8) :: pexsi_ne
      real(kind=r8) :: pexsi_mu_min
      real(kind=r8) :: pexsi_mu_max
      logical :: pexsi_first
      logical :: pexsi_started = .false.
      integer(kind=c_intptr_t) :: pexsi_plan
      type(f_ppexsi_options) :: pexsi_options

      ! EigenExa
      integer(kind=i4) :: exa_n_lrow
      integer(kind=i4) :: exa_n_lcol
      integer(kind=i4) :: exa_my_prow
      integer(kind=i4) :: exa_my_pcol
      integer(kind=i4) :: exa_n_prow
      integer(kind=i4) :: exa_n_pcol
      integer(kind=i4) :: exa_method
      integer(kind=i4) :: exa_blk_fwd
      integer(kind=i4) :: exa_blk_bkwd
      logical :: exa_first
      logical :: exa_started = .false.

      ! SLEPc-SIPs
      integer(kind=i4) :: sips_n_elpa
      integer(kind=i4) :: sips_n_slices
      real(kind=r8) :: sips_buffer ! Buffer for adjusting interval
      real(kind=r8) :: sips_interval(2)
      real(kind=r8) :: sips_inertia_tol
      logical :: sips_do_inertia
      logical :: sips_first
      logical :: sips_started = .false.

      ! NTPoly
      integer(kind=i4) :: nt_n_prow
      integer(kind=i4) :: nt_n_pcol
      integer(kind=i4) :: nt_method
      integer(kind=i4) :: nt_isr ! Method to find S^(-1/2)
      integer(kind=i4) :: nt_max_iter
      real(kind=r8) :: nt_tol
      real(kind=r8) :: nt_filter
      logical :: nt_output
      logical :: nt_first
      logical :: nt_started = .false.
      type(SolverParameters_t) :: nt_options
      type(Permutation_t) :: nt_perm
      type(ProcessGrid_t) :: nt_pgrid

      ! MAGMA
      integer(kind=i4) :: magma_solver
      integer(kind=i4) :: magma_n_gpus
      logical :: magma_started = .false.

      ! BSEPACK
      integer(kind=i4) :: bse_n_lrow
      integer(kind=i4) :: bse_n_lcol
      integer(kind=i4) :: bse_desc(9)

   end type

   type, public :: elsi_handle

      type(elsi_basic_t) :: bh
      type(elsi_param_t) :: ph
      type(fjson_handle) :: jh

      ! Dense
      real(kind=r8), allocatable :: ham_real_den(:,:)
      complex(kind=r8), allocatable :: ham_cmplx_den(:,:)
      real(kind=r8), allocatable :: ovlp_real_den(:,:)
      complex(kind=r8), allocatable :: ovlp_cmplx_den(:,:)
      real(kind=r8), allocatable :: eval(:)
      real(kind=r8), allocatable :: evec_real(:,:)
      complex(kind=r8), allocatable :: evec_cmplx(:,:)
      real(kind=r8), allocatable :: dm_real_den(:,:)
      complex(kind=r8), allocatable :: dm_cmplx_den(:,:)

      ! Sparse
      real(kind=r8), allocatable :: ham_real_sp(:)
      complex(kind=r8), allocatable :: ham_cmplx_sp(:)
      real(kind=r8), allocatable :: ovlp_real_sp(:)
      complex(kind=r8), allocatable :: ovlp_cmplx_sp(:)
      real(kind=r8), allocatable :: dm_real_sp(:)
      complex(kind=r8), allocatable :: dm_cmplx_sp(:)
      integer(kind=i4), allocatable :: row_ind_sp1(:)
      integer(kind=i4), allocatable :: col_ptr_sp1(:)
      integer(kind=i4), allocatable :: row_ind_sp2(:)
      integer(kind=i4), allocatable :: col_ptr_sp2(:)
      integer(kind=i4), allocatable :: row_ind_sp3(:)
      integer(kind=i4), allocatable :: col_ind_sp3(:)
      type(Matrix_ps) :: nt_ham
      type(Matrix_ps) :: nt_ovlp
      type(Matrix_ps) :: nt_dm

      ! Frozen core
      integer(kind=i4), allocatable :: perm_fc(:)
      real(kind=r8), allocatable :: ham_real_v(:,:)
      complex(kind=r8), allocatable :: ham_cmplx_v(:,:)
      real(kind=r8), allocatable :: ovlp_real_v(:,:)
      complex(kind=r8), allocatable :: ovlp_cmplx_v(:,:)
      real(kind=r8), allocatable :: evec_real_v(:,:)
      complex(kind=r8), allocatable :: evec_cmplx_v(:,:)

      ! Matrix redistribution
      integer(kind=i4), allocatable :: mask(:,:)
      integer(kind=i4), allocatable :: map_den(:,:)
      integer(kind=i4), allocatable :: map_sp1(:)
      integer(kind=i4), allocatable :: perm_sp3(:)

      ! Auxiliary
      real(kind=r8), allocatable :: ovlp_real_copy(:,:)
      complex(kind=r8), allocatable :: ovlp_cmplx_copy(:,:)
      real(kind=r8), allocatable :: dm_real_copy(:,:)
      complex(kind=r8), allocatable :: dm_cmplx_copy(:,:)
      real(kind=r8), allocatable :: occ(:,:,:)
      real(kind=r8), allocatable :: omm_c_real(:,:)
      complex(kind=r8), allocatable :: omm_c_cmplx(:,:)
      real(kind=r8), allocatable :: pexsi_ne_vec(:)
      type(Matrix_ps) :: nt_ovlp_copy
      type(Matrix_ps) :: nt_dm_copy
      type(Matrix_ps) :: nt_map

      logical :: handle_init = .false.

   end type

   type, public :: elsi_rw_handle

      type(elsi_basic_t) :: bh

      integer(kind=i4) :: rw_task
      integer(kind=i4) :: parallel_mode
      integer(kind=i4) :: matrix_format
      real(kind=r8) :: n_electrons
      integer(kind=i4) :: n_basis
      integer(kind=i4) :: header_user(8)
      logical :: handle_init = .false.

   end type

end module ELSI_DATATYPE