module MW_diagnostic
   use MW_kinds, only: wp, PARTICLE_NAME_LEN
   use MW_errors, only: &
         MW_errors_allocate_error => allocate_error, &
         MW_errors_deallocate_error => deallocate_error, &
         MW_errors_open_error => open_error
   use MW_ion, only: MW_ion_t
   use MW_electrode, only: MW_electrode_t
   use MW_fileunit, only: MW_fileunit_get_new_unit => get_new_unit
   implicit none
   private

   ! Public Type
   public :: MW_diagnostic_t

   ! Public subroutines
   public :: set_num_bins
   public :: define_type
   public :: void_type
   public :: compute_density

   type MW_diagnostic_t
      ! Density profile along z-axis
      integer :: num_density_samples = 0     !< number of times the density are sampled
      integer :: num_types = 0       !< number of species types
      integer :: num_bins = 0        !< number of bins
      real(wp) :: bin_width = 0.0_wp !< width of a bin
      real(wp) :: zmin = 0.0_wp !< minimum of density profile
      real(wp) :: zmax = 0.0_wp !< maximum of density profile
      real(wp), allocatable :: number_density(:,:) !< number density histogram
      real(wp), allocatable :: charge_density(:,:) !< charge density histogram

   end type MW_diagnostic_t

contains

   subroutine set_num_bins(this, num_bins, limits, zmin, zmax)
      implicit none
      ! Parameters
      ! ----------
      type(MW_diagnostic_t), intent(inout) :: this
      integer, intent(in) :: num_bins              !< number of bins
      real(wp), intent(in) :: zmin, zmax
      logical, intent(in) :: limits
      this%num_bins = num_bins
      if (limits) then
         this%zmin = zmin
         this%zmax = zmax
      end if
   end subroutine set_num_bins

   ! --------------------------------------------------------------------------------
   ! Define data type
   subroutine define_type(this, num_types, z_width)
      implicit none
      ! Parameters
      ! ----------
      type(MW_diagnostic_t), intent(inout) :: this
      integer, intent(in) :: num_types             !< number of species types
      real(wp), intent(in) :: z_width            !< width of a bin

      ! Local
      ! -----
      integer :: ierr
      integer :: num_bins

      num_bins = this%num_bins

      ! 1st dimension starts at 0 to store total sum
      allocate(this%number_density(0:num_types, num_bins), &
            this%charge_density(0:num_types, num_bins), &
            stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("allocate_arrays", "diagnostic.f90", ierr)
      end if

      if (this%zmin == 0.0_wp .and. this%zmax == 0.0_wp) then
         this%zmin = 0.0_wp
         this%zmax = z_width
      end if

      this%number_density(:,:) = 0.0_wp
      this%charge_density(:,:) = 0.0_wp
      this%num_bins = num_bins
      this%num_types = num_types
      this%bin_width = (this%zmax - this%zmin) / num_bins
      this%num_density_samples = 0

   end subroutine define_type

   ! --------------------------------------------------------------------------------
   ! Void data type
   subroutine void_type(this)
      implicit none
      ! Parameters
      ! ----------
      type(MW_diagnostic_t), intent(inout) :: this

      ! Local
      ! -----
      integer :: ierr

      if (allocated(this%number_density)) then
         deallocate(this%number_density, stat=ierr)
         if (ierr /= 0) call MW_errors_deallocate_error("void_type", "diagnostics.f90", ierr)
      end if

      if (allocated(this%charge_density)) then
         deallocate(this%charge_density, stat=ierr)
         if (ierr /= 0) call MW_errors_deallocate_error("void_type", "diagnostics.f90", ierr)
      end if
      this%num_bins = 0
      this%num_types = 0
      this%bin_width = 0.0_wp
      this%zmin = 0.0_wp
      this%zmax = 0.0_wp
      this%num_density_samples = 0

   end subroutine void_type

   ! --------------------------------------------------------------------------------
   ! Compute number and charge density
   subroutine compute_density(this, ions, electrodes, z_ions, z_elec, q_elec)
      implicit none
      ! Parameters
      ! ----------
      type(MW_diagnostic_t), intent(inout) :: this
      type(MW_ion_t), intent(in) :: ions(:)
      type(MW_electrode_t), intent(in) :: electrodes(:)
      real(wp), intent(in) :: z_ions(:)
      real(wp), intent(in) :: z_elec(:)
      real(wp), intent(in) :: q_elec(:)

      ! Local
      ! -----
      integer :: itype, iion, i, offset, ibin
      integer :: num_ion_types, num_elec_types

      num_ion_types = size(ions,1)
      num_elec_types = size(electrodes,1)
      do itype = 1, num_ion_types
         offset = ions(itype)%offset
         do iion = 1, ions(itype)%count
            i = iion + offset
            ibin = floor((z_ions(i) - this%zmin) / this%bin_width) + 1
            ibin = min(ibin, this%num_bins)
            ibin = max(ibin, 1)
            this%number_density(0, ibin) = this%number_density(0, ibin) + 1.0_wp
            this%number_density(itype, ibin) = this%number_density(itype, ibin) + 1.0_wp
            this%charge_density(0, ibin) = this%charge_density(0, ibin) + ions(itype)%charge
            this%charge_density(itype, ibin) = this%charge_density(itype, ibin) + ions(itype)%charge
         end do
      end do

      do itype = 1, num_elec_types
         offset = electrodes(itype)%offset
         do iion = 1, electrodes(itype)%count
            i = iion + offset
            ibin = floor((z_elec(i) - this%zmin) / this%bin_width) + 1
            ibin = min(ibin, this%num_bins)
            ibin = max(ibin, 1)
            this%number_density(0, ibin) = this%number_density(0, ibin) + 1.0_wp
            this%number_density(num_ion_types+itype, ibin) = this%number_density(num_ion_types+itype, ibin) + 1.0_wp
            this%charge_density(0, ibin) = this%charge_density(0, ibin) + q_elec(i)
            this%charge_density(num_ion_types+itype, ibin) = this%charge_density(num_ion_types+itype, ibin) + q_elec(i)
         end do
      end do

      this%num_density_samples = this%num_density_samples + 1
   end subroutine compute_density

end module MW_diagnostic
