text: from .database import kirkland_params
  from .optics import *
  from .utils import *
  from .cuda_kernels import ProbeKernels, PotentialKernels
  import numpy as np
  from scipy.special import k0
  import multiprocessing as mp
  import ctypes
  import sys
  from warnings import warn, catch_warnings, simplefilter
  import os
  
  try:
      import pyfftw
  except:
      pyfftw=None
  import pycuda
  import pycuda.driver as cuda
  import skcuda.fft as skfft
  import skcuda.cufft as cufft
  from time import time
  import h5py
  from mpi4py import MPI
  
  
  XYZ_dtype = [('atomic_number', 'i'), ('x', 'f'), ('y', 'f'), ('z', 'f'), ('occ', 'f'), ('DW', 'f')]
  
  def catch_warn():
      warn('complex_casting', RuntimeWarning)
  
  def unwrap(args):
      (msa, params), kwargs = args
      method_name = kwargs.pop('method')
      if method_name == 'build_slices':
          return msa.make_slice(params)
      elif method_name == 'propagate_beams':
          return msa.propagate_beam(params)
      elif method_name == 'build_probes':
          return msa.build_probe(probe_position=params[0], probe_dict=params[1])
  
  
  class MSA(object):
      def __init__(self, energy, semi_angle, supercell, sampling=np.array([512, 512]), max_angle=None, verbose=False,
                   debug=False, output_dir='', debye_waller=True):
          self.E = energy
          self.Lambda = voltage2Lambda(self.E*1e3)
          self.semi_ang = semi_angle
          self.verbose = verbose
          self.debug = debug
          self.output_dir = output_dir
          self.debye_waller = debye_waller
          # Load and set supercell properties
          if isinstance(supercell, np.ndarray):
              try:
                  check = [supercell[0][field] for field in ['x', 'y', 'z', 'occ', 'DW', 'atomic_number']]
                  supercell_arr = np.copy(supercell)
              except ValueError as err:
                  warn('Supercell numpy array does not have required fields.')
                  raise err
          elif isinstance(supercell, str):
              if os.path.exists(supercell):
                  try:
                      supercell_arr = np.genfromtxt(supercell, skip_footer=1, skip_header=2, dtype=XYZ_dtype)
                  except IOError as err:
                      warn('Unable to load supercell file into numpy array.')
                      raise err
              else:
                  warn('Path %s was not found.' % supercell)
                  sys.exit(0)
          else:
              warn('supercell format was not recognized.')
              sys.exit(0)
          self.supercell_xyz = np.column_stack([supercell_arr['x'], supercell_arr['y'], supercell_arr['z']]).astype(np.float32)
          self.supercell_Z = supercell_arr['atomic_number']
          self.supercell_dw = supercell_arr['DW']
          self.supercell_occ = supercell_arr['occ']
  
          # Adding random displacements due to thermal effects
          unique_Z = np.unique(self.supercell_Z)
          for Z in unique_Z:
              indx = np.where(self.supercell_Z == Z)[0]
              dw = self.supercell_dw[indx][0]
              displacements = float(self.debye_waller) * dw * np.random.standard_normal(size=self.supercell_xyz[indx].shape)
              self.supercell_xyz[indx] += displacements
  
          # Set simulation parameters
          self.dims = self.supercell_xyz.max(0) - self.supercell_xyz.min(0)
          if max_angle is None:
              self.sampling = sampling
              self.kmax = np.min(self.sampling/self.dims[:2])
              self.max_ang = self.kmax * self.Lambda
          else:
              self.max_ang = max_angle
              self.kmax = self.max_ang / self.Lambda
              self.sampling = np.floor(self.kmax * self.dims[:2]).astype(np.int)
              self.sampling = self.sampling[::-1]
          self.pix_size = self.dims[:2][::-1] / self.sampling
          self.kpix_size = self.kmax/self.sampling
          self.sigma = sigma_int(self.E*1e3)
          self.print_verbose('Simulation Parameters:\nSupercell dimensions xyz:%s (Å)\nReal, Reciprocal space pixel sizes:%s Å, %s 1/Å'
                '\nMax angle: %2.2f (rad)\nSampling in real and reciprocal space: %s pixels,\nThermal Effects: %s' %
                (format(np.round(self.dims, 2)), format(np.round(self.pix_size, 2)), format(np.round(self.kpix_size, 2)),
                 self.max_ang, format(self.sampling), format(self.debye_waller)))
  
      def print_debug(self, *args, **kwargs):
          if self.debug:
              print(*args, **kwargs)
  
      def print_verbose(self, *args, **kwargs):
          if self.verbose:
              print(*args, **kwargs)
  
      def calc_atomic_potentials(self, potential_range=8, oversample=2, kirkland=True):
          if kirkland:
              self.scattering_params = kirkland_params
              #TODO: Figure out how to calculate scattering potential using different methods
          start = potential_range / 2
          step_x, step_y = 1.j * np.floor(potential_range/self.pix_size).astype(np.int) * oversample
          grid_x, grid_y = np.mgrid[-start:start:step_x, -start:start:step_y]
          self.cached_pots = dict()
          for Z in np.unique(self.supercell_Z):
              self.cached_pots[Z] = self._get_potential(grid_x, grid_y, self.scattering_params[Z - 1], oversample)
  
      def build_potential_slices(self, slice_thickness):
          self.slice_t = slice_thickness
          num_slices = np.int(np.floor(self.dims[-1] / slice_thickness))
          tasks = [((self, slice_num), {'method': 'build_slices'}) for slice_num in range(num_slices)]
          processes = min(mp.cpu_count(), num_slices)
          chunk = np.int(np.floor(num_slices / processes))
          pool = mp.Pool(processes=processes, maxtasksperchild=1)
          jobs = pool.imap(unwrap, tasks, chunksize=chunk)
          potential_slices = np.array([j for j in jobs])
          pool.close()
          self.potential_slices = potential_slices
          self.print_verbose('Built %d potential slices with shape:%s pixels' % (self.potential_slices.shape[0],
                                                                    format(self.potential_slices.shape[1:])))
  
      def make_slice(self, args):
          slice_num = args
          mask = np.logical_and(self.supercell_xyz[:, -1] >= slice_num * self.slice_t,
                                self.supercell_xyz[:, -1] < (slice_num + 1) * self.slice_t)
          arr_slice = np.zeros(self.sampling, dtype=np.float32)
          for Z, xyz in zip(self.supercell_Z[mask], self.supercell_xyz[mask]):
              pot = self.cached_pots[Z]
              y_pix, x_pix = xyz[:2][::-1] / self.pix_size
              y_start, y_end = int(y_pix - pot.shape[0] / 2), int(y_pix + pot.shape[0] / 2)
              x_start, x_end = int(x_pix - pot.shape[1] / 2), int(x_pix + pot.shape[1] / 2)
              repl_y = slice(max(y_start, 0), y_end)
              repl_x = slice(max(x_start, 0), x_end)
              offset_y = abs(min(y_start, 0))
              offset_x = abs(min(x_start, 0))
              repl_shape = arr_slice[repl_y, repl_x].shape
              arr_slice[repl_y, repl_x] += pot[offset_y:repl_shape[0] + offset_y, offset_x:repl_shape[1] + offset_x]
          return arr_slice
  
      @staticmethod
      def _get_potential(grid_x, grid_y, params, oversample):
          eps = 0.1  # Å cutoff for Bessel function
          a_0 = 0.529  # Å Bohr radius
          e = 14.4  # V.Å electron charge
          r = np.sqrt(grid_x ** 2 + grid_y ** 2)
          prefac_1 = 4 * np.pi ** 2 * a_0 * e
          coeff = params
          a_coeff, b_coeff, c_coeff, d_coeff = coeff
          bessel_arg = 2 * np.pi * r[:, :, np.newaxis] * np.sqrt(b_coeff)
          bessel_eval = k0(bessel_arg + eps)
          sum_v = prefac_1 * a_coeff * bessel_eval + prefac_1 / 2. * c_coeff / d_coeff * np.exp(-np.pi ** 2 *
                                                                                       r[:, :, np.newaxis] ** 2 / d_coeff)
          v = sum_v.sum(-1)
          potential = bin_2d_array(v, dwnspl=oversample, mode='mean')
          potential -= potential.min()
          return potential.astype(np.float32)
  
      def build_probe(self, probe_position=np.array([0., 0.]), probe_dict={'smooth_apert': True, 'apert_smooth': 50, 'spherical_phase': True, 'aberration_dict':
                          {'C1': 0., 'C3': 0., 'C5': 0.},'scherzer': True}
                      ):
          self.probe_dict = probe_dict
          k_y, k_x = np.mgrid[-self.kmax/2: self.kmax/2: 1j*self.sampling[0],
                              -self.kmax/2: self.kmax/2: 1j * self.sampling[1]]
          k_rad = np.sqrt(k_x ** 2 + k_y ** 2)
          k_semi = self.semi_ang/self.Lambda
  
          # aperture function
          if probe_dict['smooth_apert']:
              aperture = 1 / (1 + np.exp(-2 * probe_dict['apert_smooth'] * (k_semi - k_rad)))
          else:
              aperture = np.heaviside(k_semi - k_rad, 0.5)
  
          # aberration
          if probe_dict['spherical_phase']:
              phase_error = spherical_phase_error(k_rad, self.Lambda, probe_dict['scherzer'],
                                                  **probe_dict['aberration_dict'])
  
          else:
              pass
              #TODO: implement non-rotationally invariant phase error
  
          # probe wavefunction
          psi_k = aperture * phase_error
          psi_k = psi_k.astype(np.complex64)
          y, x = probe_position
          kr = k_x * x + k_y * y
          phase_shift = np.exp(2 * np.pi * 1.j * kr).astype(np.complex64)
          psi_x = pyfftw.interfaces.numpy_fft.ifft2(psi_k * phase_shift)
          psi_x = pyfftw.interfaces.numpy_fft.fftshift(psi_x)
          # TODO: need to make fft library choice optional
          # psi_x = np.fft.ifft2(psi_k * phase_shift, norm='ortho')
          # psi_x = np.fft.fftshift(psi_x)
          psi_x /= np.sqrt(np.sum(np.abs(psi_x) ** 2))
          return psi_x.astype(np.complex64), psi_k.astype(np.complex64), aperture.astype(np.float32)
  
      def build_propagator(self):
          k_y, k_x = np.mgrid[-self.kmax / 2: self.kmax / 2: 1j * self.sampling[0],
                     -self.kmax / 2: self.kmax / 2: 1j * self.sampling[1]]
          k_rad_sq = k_x ** 2 + k_y ** 2
          propag = np.exp(-np.pi * 1.j * self.Lambda * self.slice_t * k_rad_sq)
          return propag.astype(np.complex64)
  
      @staticmethod
      def bandwidth_limit_mask(arr_shape, radius=0.5):
          # assumes square image
          grid_x, grid_y = np.mgrid[-arr_shape[0] // 2:arr_shape[0] // 2, -arr_shape[1] // 2:arr_shape[1] // 2]
          r_grid = np.sqrt(grid_x ** 2 + grid_y ** 2)
          bl_mask = np.heaviside(max(arr_shape[0], arr_shape[1]) * radius - r_grid, 0)
          return bl_mask.astype(np.float32)
  
      def generate_probe_positions(self, probe_step=np.array([0.1, 0.1]), probe_range=np.array([[0., 1.0], [0., 1.0]]), grid_steps=None):
          if grid_steps is not None:
              grid_range_start = (0.5 - self.dims[:2]/2)/2
              grid_range_stop = (0.5 + self.dims[:2]/2)/2
              x_pos, y_pos = np.mgrid[grid_range_start[0]:grid_range_stop[0]:-1j*grid_steps[0], 
                                      grid_range_start[1]:grid_range_stop[1]:-1j*grid_steps[1]]
              grid_steps_x, grid_steps_y = grid_steps.astype(np.int) 
              grid_range_x = np.array([grid_range_start[0], grid_range_stop[0]])
              grid_range_y = np.array([grid_range_start[1], grid_range_stop[1]])
              probe_pos = np.array([[-x, y] for y, x in zip(y_pos.flatten()[::-1], x_pos.flatten())])
  
          else:                            
              grid_steps_x, grid_steps_y = np.floor(np.diff(probe_range).flatten() * self.dims[:2] / probe_step).astype(np.int)
              grid_range_x, grid_range_y = [(probe_range[i] - np.ones((2,)) * 0.5) * self.dims[i]
                                          for i in range(2)]
              y_pos, x_pos = np.mgrid[grid_range_y[0]: grid_range_y[1]: -1j * grid_steps_y,
                          grid_range_x[0]: grid_range_x[1]: -1j * grid_steps_x]
              probe_pos = np.array([[y, -x] for y, x in zip(y_pos.flatten()[::-1], x_pos.flatten())])
          self.grid_steps = np.array([grid_steps_x, grid_steps_y])
          self.grid_range = np.array([grid_range_x, grid_range_y]).flatten()
          self.probe_positions = probe_pos
          self.num_probes = np.int32(probe_pos.shape[0])
         
      def multislice(self, probe_pos=np.array([0., 0.]), probe_grid=True, save_probes=True, bandwidth=1 / 3):
          # check for slices
          if isinstance(self.potential_slices, np.ndarray) is False:
              warn('Potential slices must be calculated first before calling multi_slice')
              return
          # check for probe
          if isinstance(self.probe_dict, dict) is False:
              warn('Probe wave function must be initialized first before calling multi_slice')
              return
          if probe_grid:
              # Define heuristics for python multiprocessing and FFTW multi-threading
              self.fftw_threads = int(self.sampling.max() // 512)
              processes = min(mp.cpu_count(), self.probe_positions.shape[0]) // self.fftw_threads
              chunk = int(np.floor(self.probe_positions.shape[0] / processes))
  
  
              # Put the potential slices in shared memory so all workers access it (asynchronously)
              global shared_slices
              shared_slices = mp.Array(ctypes.c_float, self.potential_slices.size, lock=False)
              temp = np.frombuffer(shared_slices, dtype=np.float32)
              for (i, pot) in enumerate(self.potential_slices):
                  temp[i * pot.size:(i + 1) * pot.size] = pot.flatten().astype(np.float32)
  
              tasks = (((self, (probe_num, pos, save_probes, probe_grid, bandwidth)), {'method': 'propagate_beams'})
                       for (probe_num, pos) in enumerate(self.probe_positions))
              pool = mp.Pool(processes=processes, initargs=(shared_slices,))
              jobs = pool.map(unwrap, tasks, chunksize=chunk)
              trans_probes = np.array([j for j in jobs])
              pool.close()
              pool.join()
          else:
              trans_probes = self.propagate_beam([None, probe_pos, save_probes, probe_grid, bandwidth])
  
          self.print_verbose('Propagated %d probe wavefunctions' % trans_probes.shape[0])
          self.trans_probes = trans_probes
          return self.trans_probes
  
      def propagate_beam(self, args):
          probe_num, probe_pos, save_probes, probe_grid, bandwidth = args
          self.print_debug('received params.')
          propag = self.build_propagator()
          blim_mask = self.bandwidth_limit_mask(propag.shape, radius=bandwidth)
          probe, _, _ = self.build_probe(probe_pos, self.probe_dict)
          probes = []
          probe_last = probe
          if probe_grid:
              slices = np.frombuffer(shared_slices, dtype=np.float32)
              slices = slices.reshape(self.potential_slices.shape)
              self.print_debug('fetched data from shared memory.')
          else:
              slices = self.potential_slices
          slices = np.exp(1.j * self.sigma * slices).astype(np.complex64)
  
          for (i, trans) in enumerate(slices[::-1]):
              t_psi = pyfftw.byte_align(trans * probe_last)
              fft_fwd = pyfftw.builders.fft2(t_psi, threads=self.fftw_threads, avoid_copy=True)
              temp = pyfftw.byte_align(fft_fwd() * blim_mask * propag)
              fft_bwd = pyfftw.builders.ifft2(temp, threads=self.fftw_threads, avoid_copy=True)
              probe_last = fft_bwd()
              if save_probes:
                  probes.append(probe_last)
          self.print_debug('finished beam propagation.')
  
          if save_probes:
              np.save(os.path.join(self.output_dir, 'probes_%d.npy') % probe_num, np.array(probes), allow_pickle=False)
              del probes
              self.print_verbose('finished with probe position: %s' % format(probe_pos))
  
          return probe_last
  
      def check_simulation(self):
          prob = np.sum([np.abs(probe)**2 for probe in self.trans_probes], axis=(1, 2))
          max_val = prob.max()
          min_val = prob.min()
          print('Max (Min) Integrated Intensity: %2.2f (%2.2f)' % (max_val, min_val))
          if max_val > 1.0 or min_val < 0.95:
              print('Significant deviation of the probability from unity is found.\n'
                    'Change the sampling and/or slice thickness!')
          return prob
  
  
  class MSAHybrid(MSA):
      def setup_device(self, gpu_rank=0):
          global ctx
          cuda.init()
          dev = cuda.Device(gpu_rank)
          ctx = dev.make_context()
          # ctx.attach()
          self.gpu_rank = gpu_rank
  
          import atexit
          def _clean_up():
              global ctx
              if ctx is not None:
                  try:#global ctx
                      #ctx.push()
                      ctx.pop()
                      ctx.detach()
                      #ctx = None
                  except:
                      pass
              from pycuda.tools import clear_context_caches
              clear_context_caches()
  
  
          atexit.register(_clean_up)
          return ctx # return context in case of manual clean-up
  
      def plan_simulation(self, num_probes=None):
          if num_probes is None:
              num_probes = self.num_probes
          self.print_verbose('Simulation requested %d probes simultaneously.' % self.num_probes)
          free_mem, tot_mem = pycuda.driver.mem_get_info()
          free_mem = free_mem/1024e6  # in GB
          mem_alloc = num_probes * np.prod(self.sampling) * 8 / 1024e6 + self.potential_slices.nbytes / 1024e6
          self.print_verbose('mem_alloc: %2.3f' % mem_alloc)
          self.print_verbose('free_mem: %2.3f' % free_mem)
          self.max_probes = num_probes
          mod = False
          while mem_alloc > 0.5 * free_mem:
              mod = True
              num_probes = int(num_probes * 0.75)
              mem_alloc = num_probes * np.prod(self.sampling) * 8 / 1024e6 + self.potential_slices.nbytes / 1024e6
              self.max_probes = num_probes
          if mod: self.print_verbose('Device can hold at most %d probes simultaneously.' % self.max_probes)
  
      def multislice(self, save_probes=True, bandwidth=1/3):
          # not supporting a single probe!!
          # check for potential slices
          if isinstance(self.potential_slices, np.ndarray) is False:
              warn('Potential slices must be calculated first before calling multi_slice')
              return
          # check for probe
          if isinstance(self.probe_dict, dict) is False:
              warn('Probe wave function must be initialized first before calling multi_slice')
              return
          # check for probe positions
          if isinstance(self.probe_positions, np.ndarray) is False:
              warn('probe positions must be initialized first before calling multi_slice')
              return
  
          # Initialize needed data
          slices = np.exp(1.j * self.sigma * self.potential_slices).astype(np.complex64)
          propag = self.build_propagator()
          mask = self.bandwidth_limit_mask(propag.shape, radius=bandwidth)
          t = time()
          self.probes = self.build_probes_cpu()
          sim_t = time() - t
          self.print_verbose('Spent %2.4f s building %d probes on cpu' % (sim_t, self.max_probes))
  
          # Copy over to device
          trans_gpu = pycuda.gpuarray.to_gpu_async(slices)
          mask_propag_gpu = pycuda.gpuarray.to_gpu_async(mask * propag)
          probes_gpu = pycuda.gpuarray.to_gpu_async(self.probes)
  
          # Setup fft plans
          # TODO: tile multiple fft plans
          t = time()
          plan = skfft.Plan(propag.shape, np.complex64, np.complex64, batch=self.max_probes)
          for slice_num in range(trans_gpu.shape[0]):
              for z_slice in range(probes_gpu.shape[0]):
                  probes_gpu[z_slice] *= trans_gpu[slice_num]
              cufft.fft(probes_gpu, probes_gpu, plan, True)
              for z_slice in range(probes_gpu.shape[0]):
                  probes_gpu[z_slice] *= mask_propag_gpu
              cufft.ifft(probes_gpu, probes_gpu, plan, False)
          cufft.fft(probes_gpu, probes_gpu, plan, True)  #return probe wavefunctions in reciprocal space
          sim_t = time() - t
          self.print_verbose('Propagated %d probes in %2.4f s' % (self.max_probes, sim_t))
          self.probes = probes_gpu.get()
  
          # free up device memory
          trans_gpu.gpudata.free()
          mask_propag_gpu.gpudata.free()
          probes_gpu.gpudata.free()
  
          # destroy cufft object
          cufft.cufftDestroy(plan.handle)
          return self.probes
  
      def build_probes_cpu(self):
          processes = min(mp.cpu_count(), self.probe_positions.shape[0]) // 4
          chunk = int(np.floor(self.probe_positions.shape[0] / processes))
          tasks = (((self, (pos, self.probe_dict)), {'method': 'build_probes'})
                   for pos in self.probe_positions)
          pool = mp.Pool(processes=processes)
          jobs = pool.map(unwrap, tasks, chunksize=chunk)
          probes = np.array([j[0] for j in jobs])
          pool.close()
          pool.join()
          return probes
  
  
  class MSAGPU(MSAHybrid):
      
      @staticmethod
      def clean_up(ctx=None, vars=None):
          if vars is not None:
              for var in vars:
                  if var is not None: 
                      var.free()
          if ctx is not None:
              ctx.pop()
              ctx.detach()
              ctx = None
              from pycuda.tools import clear_context_caches
              clear_context_caches()
      
  #     def build_potential_slices(self, slice_thickness):
  # #         # find number of slices and atomic sites per slice
  #         self.slice_t = slice_thickness
  #         self.num_slices = np.int32(np.floor(self.dims[-1] / slice_thickness))
  # #         masks = [np.logical_and(self.supercell_xyz[:, -1] >= slice_num * self.slice_t,
  # #                         self.supercell_xyz[:, -1] < (slice_num + 1) * self.slice_t)
  # #                         for slice_num in range(self.num_slices)]
  # #         # remap Z numbers to indices in cached atomic potential dictionary
  # #         unique_Z = np.unique(self.supercell_Z)
  # #         supercell_Z_idx = np.array([np.argmax(np.equal(unique_Z, Z_val)) for Z_val in self.supercell_Z])
  # #         Z_arr = np.array([supercell_Z_idx[mask] for mask in masks])
  # #         supercell_pix = self.supercell_xyz[:,:2]/self.pix_size[::-1]
  # #         yx_arr = np.array([supercell_pix[mask] for mask in masks])
  
  # #         # stack the input for all slices into a single 1-d array
  # #         Zxy_input = np.array([np.column_stack([Z, yx[:,1], yx[:,0]]).astype(np.int32).flatten()
  # #                                             for Z, yx in zip(Z_arr, yx_arr)])
  
  # #         # pad sites array to have equal num of element per slice
  # #         shapes = np.array([itm.shape for itm in Zxy_input])
  # #         pads = [(0, shapes.max() - itm.shape) for itm in Zxy_input]
  # #         padded = [np.pad(itm,pad,'constant',constant_values=-1) for pad, itm in zip(pads, Zxy_input)]
  # #         Zxy_input = np.vstack([np.pad(itm,pad,'constant',constant_values=-1) for pad, itm in zip(pads, Zxy_input)])
  
  # #         # stack atomic potentials of unique elements
  # #         atom_pot_stack = np.array([self.cached_pots[uq_Z] for uq_Z in unique_Z]).astype(np.float32)
  
  # #         # compile/load cuda kernels
  # #         self._load_potential_kernels([atom_pot_stack.shape[1], atom_pot_stack.shape[2]],
  # #                                                 Zxy_input.shape[-1])
  
  # #         # allocate memory
  # #         atom_pot_stack_d = cuda.to_device(atom_pot_stack)
  # #         sites_d = cuda.to_device(Zxy_input)
  # #         self.potential_slices = cuda.aligned_empty((int(self.num_slices),
  # #                      int(self.sampling[0]), int(self.sampling[1])), np.complex64)
  #         self.vars = []
  # #         self.vars.append(self.potential_slices.base)
  # #         self.potential_slices = cuda.register_host_memory(self.potential_slices)
  #         potential_slices_d = cuda.to_device(self.potential_slices)
  
  #         # store device allocation ref for later use
  #         self.pot_dev_ptr = potential_slices_d
  
  #         # get block, grid dimensions
  # #         block, grid = self._get_blockgrid([self.sampling[1], self.sampling[0], self.num_slices],
  # #                     mode='3D')
  # #         build_potential = self.pot_kernels['build_potential']
  
  # #         # build potential
  # #         build_potential(potential_slices_d, atom_pot_stack_d, sites_d,
  # #                         np.float32(self.sigma), block=block, grid=grid)
  #         ctx.synchronize()
  # #         cuda.memcpy_dtoh_async(self.potential_slices, potential_slices_d, cuda.Stream())
  
  # #         # free gpu memory
  # #         sites_d.free()
  # #         atom_pot_stack_d.free()
  #         self.print_verbose('Built %d potential slices with shape:%s pixels' % (self.potential_slices.shape[0],
  #                                                                   format(self.potential_slices.shape[1:])))
  #         # unregister host memory
  #         #self.potential_slices.base.unregister()
  #         #self.potential_slices.base.free()
  
      def _load_kernels(self):
          try:
              probe_kernels = ProbeKernels(sampling=self.sampling)
              self.kernels = probe_kernels.kernels
              self.print_verbose('CUDA C/C++ Probe Kernels compiled successfully')
          except cuda.CompileError:
              warn('CUDA C/C++ Probe Kernels did not compile successfully')
              raise cuda.CompileError
  
      def _load_potential_kernels(self, potential_shape, sites_size ):
          try:
              potential_kernels = PotentialKernels(sampling=self.sampling,
              num_slices=self.num_slices, potential_shape=potential_shape,
              sites_size=sites_size)
              self.pot_kernels = potential_kernels.kernels
              self.print_verbose('CUDA C/C++ Potential Kernels compiled successfully')
          except cuda.CompileError:
              warn('CUDA C/C++ Potential Kernels did not compile successfully')
              raise cuda.CompileError
  
      def build_probe(self, probe_dict={'smooth_apert': True, 'apert_smooth': 50, 'spherical_phase': True,
                                        'aberration_dict': {'C1': 0., 'C3': 0., 'C5': 0.},'scherzer': True}):
          # load kernels and prepare func args
          self._load_kernels()
          aber_dict = probe_dict['aberration_dict']
          k_semi = np.float32(self.semi_ang / self.Lambda)
          k_max, Lambda, C1, C3, C5 = [np.float32(itm) for itm in [self.kmax, self.Lambda, aber_dict['C1'],
                                                                   aber_dict['C3'], aber_dict['C5']]]
          if probe_dict['scherzer']:
              C1, _ = scherzer_params(self.Lambda, aber_dict['C3'])
              C1 = np.float32(C1)
  
          # grab appropriate kernels
          self.multwise_2d_func = self.kernels['mult_wise_c2d_re2d']
          fftshift_func = self.kernels['fftshift_2d']
          if probe_dict['smooth_apert']:
              apert_func = self.kernels['soft_aperture']
          else:
              apert_func = self.kernels['hard_aperture']
          if probe_dict['spherical_phase']:
              phase_func = self.kernels['spherical_phase_error']
          else:
              pass
              # TODO: implement non-rotationally invariant phase error kernel
  
          # define block/grid threads
          shape_x = np.int32(self.sampling[1])
          shape_y = np.int32(self.sampling[0])
          block, grid = self._get_blockgrid([shape_x, shape_y], mode='2D')
  
          # allocate memory
          self.apert = np.empty(self.sampling, dtype=np.float32)
          #self.apert = cuda.register_host_memory(self.apert)
          apert_d = cuda.mem_alloc(self.apert.nbytes)
          self.psi_k = np.empty(self.sampling, dtype=np.complex64)
          #self.psi_k = cuda.register_host_memory(self.psi_k)
          self.psi = np.empty_like(self.psi_k)
          #self.psi = cuda.register_host_memory(self.psi)
          psi_k_d = cuda.mem_alloc(self.psi_k.nbytes)
          psi_x_d = cuda.mem_alloc(self.psi_k.nbytes)
  
          # build a probe in k-space
          apert_func(apert_d, k_max, k_semi, shape_x, shape_y, block=block, grid=grid, shared=0, stream=cuda.Stream())
          #ctx.synchronize()
          phase_func(psi_k_d, k_max, Lambda, C1, C3, C5, shape_x, shape_y, block=block, grid=grid, shared=0, stream=cuda.Stream())
          #ctx.synchronize()
          self.multwise_2d_func(psi_k_d, apert_d, shape_x, shape_y, block=block, grid=grid)
          #ctx.synchronize()
          cuda.memcpy_dtoh_async(self.psi_k, psi_k_d, cuda.Stream())
          cuda.memcpy_dtoh_async(self.apert, apert_d, cuda.Stream())
  
          # build probe in x-space
          fft_plan = skfft.Plan(self.psi_k.shape, np.complex64, np.complex64, batch=1)
          cufft.cufftExecC2C(fft_plan.handle, int(psi_k_d), int(psi_x_d), cufft.CUFFT_INVERSE)
          ctx.synchronize()
          fftshift_func(psi_x_d, shape_y, block=block, grid=grid)
          #ctx.synchronize()
          cuda.memcpy_dtoh_async(self.psi, psi_x_d, cuda.Stream())
          self.normalization = np.sum(np.abs(self.psi, dtype=np.float64)**2)
          self.psi /= np.sqrt(self.normalization)
          cuda.memcpy_dtoh_async(self.psi_k, psi_k_d, cuda.Stream())
  
          # free gpu mem
          psi_x_d.free()
          psi_k_d.free()
          apert_d.free()
          cufft.cufftDestroy(fft_plan.handle)
          
      @staticmethod
      def _get_blockgrid(shapes, mode='2D'):
          # define block/grid threads
          if mode == '3D':
              shape_x = shapes[0]
              shape_y = shapes[1]
              shape_z = shapes[2]
              blk_zsize = 1
              blk_xsize = 32
              blk_ysize = 32
              grid_xsize = int(round((shape_x + blk_xsize - 1) / blk_xsize))
              grid_ysize = int(round((shape_y + blk_ysize - 1) / blk_ysize))
              grid_zsize = int(round((shape_z + blk_zsize - 1) / blk_zsize))
              block_3d = (blk_xsize, blk_ysize, blk_zsize)
              grid_3d = (grid_xsize, grid_ysize, grid_zsize)
              return block_3d, grid_3d
          if mode == '2D':
              shape_x = shapes[0]
              shape_y = shapes[1]
              block_2d = (32, 32, 1)
              grid_2d = (int(round((shape_x + block_2d[0] - 1) / 32)), int(round((shape_y + block_2d[1] - 1) / 32)), 1)
              return block_2d, grid_2d
          if mode == '1D':
              shape_x = shapes[0]
              shape_y = shapes[1]
              shape_z = shapes[2]
              block_1d = (int(min(shape_x * shape_y, 1024)), 1, 1)
              grid_1d = (int((shape_x * shape_y) / block_1d[0]), int(shape_z), 1)
              return block_1d, grid_1d
  
      def multislice(self, bandwidth=1/3, unified_mem=False, batch_size=None, transmit=True):
          """
  
          :param bandwidth:
          :return:
          """
          ## TODO: when batch_size != self.num_probes, probe positions get scrambled!!!
          if batch_size == None:
              batch_size = self.num_probes 
          # checks
          if isinstance(self.potential_slices, np.ndarray) is False:
              warn('Potential slices must be calculated first before calling multi_slice\n. '
                   'Call method build_potential_slices().')
              return
          if isinstance(self.psi_k, np.ndarray) is False:
              warn('Probe in k-space must be initiated first before calling multi_slice\n. '
                   'Call method build_probe().')
              return
          if isinstance(self.probe_positions, np.ndarray) is False:
              warn('Probe positions must be calculated first before calling multislice\n. '
                   'Call method generate_probe_positions().')
              return
  
          # load cuda kernels and prepare arguments
          shape_x = np.int32(self.sampling[1])
          shape_y = np.int32(self.sampling[0])
          num_probes = self.max_probes
          k_semi = np.float32(self.semi_ang / self.Lambda)
  
          # define block/grid threads
          block_3d, grid_3d = self._get_blockgrid([self.sampling[1], self.sampling[0], num_probes], mode='3D')
          block_2d, grid_2d = self._get_blockgrid([self.sampling[1], self.sampling[0], num_probes], mode='2D')
          self.print_debug('block, grid:', block_3d, grid_3d)
  
          # allocate memory
          # self.probes = np.empty((self.num_probes, shape_y, shape_x), dtype=np.complex64)
          self.propag = cuda.aligned_empty((int(self.sampling[0]), int(self.sampling[1])), np.complex64)
          self.vars.append(self.propag.base)
          self.propag = cuda.register_host_memory(self.propag)
          propag_d = cuda.to_device(self.propag)
          self.mask = cuda.aligned_empty((int(self.sampling[0]), int(self.sampling[1])), np.float32)
          self.vars.append(self.mask.base)
          self.mask = cuda.register_host_memory(self.mask)
          mask_d = cuda.to_device(self.mask)
          grid_steps_d = cuda.to_device(self.grid_steps.astype(np.int32))
          grid_range_d = cuda.to_device(self.grid_range.astype(np.float32))
          psi_k_d = cuda.to_device(self.psi_k)
          if unified_mem:
              self.probes = cuda.managed_empty(shape=(int(self.num_probes), int(self.sampling[0]), int(self.sampling[1])),
                                          dtype=np.complex64, mem_flags=cuda.mem_attach_flags.GLOBAL)
          else:
              # pinned memory is default
              self.probes = cuda.aligned_empty((int(self.num_probes), int(self.sampling[0]), int(self.sampling[1])), np.complex64)
              self.vars.append(self.probes.base)
              self.probes = cuda.register_host_memory(self.probes)
          ones = cuda.aligned_zeros((int(self.sampling[0]), int(self.sampling[1])), np.complex64) + 1
          self.vars.append(ones.base)
          ones = cuda.register_host_memory(ones)
          ones_d = cuda.mem_alloc(ones.nbytes)
          cuda.memcpy_htod_async(ones_d, ones, cuda.Stream())
  
  
          # grab needed kernels
          propag_func = self.kernels['propagator']
          mask_func = self.kernels['hard_aperture']
          multwise_func = self.kernels['mult_wise_c2d_re2d']
  
          # 1. Build propagator and bandwidth limiting mask
          mask_func(mask_d, np.float32(self.kmax), np.float32(bandwidth * self.kmax), shape_x, shape_y, block=block_2d,
                    grid=grid_2d, shared=0, stream=cuda.Stream())
          #ctx.synchronize()
          propag_func(propag_d, np.float32(self.kmax), np.float32(self.slice_t), np.float32(self.Lambda), shape_x, shape_y,
                      block=block_2d, grid=grid_2d, shared=0, stream=cuda.Stream())
          #ctx.synchronize()
          multwise_func(propag_d, mask_d, shape_x, shape_y, block=block_2d, grid=grid_2d, shared=0)
          cuda.memcpy_dtoh_async(self.propag, propag_d, cuda.Stream())
          cuda.memcpy_dtoh_async(self.mask, mask_d, cuda.Stream())
  
          # 1. Generate batch fft plans, create host/device pointers, and launch cuda streams
          num_phases = self.num_probes//self.max_probes
          phases = []
          self.results = []
          t = time()
          num_stream=0
          for i in range(num_phases + 1):
              phase = slice(i * self.max_probes, (i+1) * self.max_probes)
              if i == num_phases :
                  phase = slice(i * self.max_probes, self.num_probes)
              phases.append(phase)
          self.print_debug('Simulation split into %d serial phases.' % len(phases))
          self.print_debug('Phases: %s ' %format(phases))
          for (i, phase) in enumerate(phases):
              t = time()
              if self.probes[phase].shape[0] == 0: break
              self.print_verbose('Starting simulation phase #%d' % i)
              num_batches = (phase.stop - phase.start + 1)//batch_size
              batches = []
              streams = []
              probes_d = []
              norm_consts = []
              plans = []
          # 2. Generate batch fft plans and create pinned memory pointers and cuda streams
              for batch_num in range(num_batches+1):
                  if batch_num == num_batches:
                      slice_obj = slice(batch_num * batch_size, None)
                  else:
                      slice_obj = slice(batch_num * batch_size, (batch_num + 1) * batch_size)
                  batches.append(slice_obj)
                  stream = cuda.Stream()
                  num_stream += 1
                  self.print_debug(' Launch Stream # %d' % num_stream)
                  streams.append(stream)
                  num_probes = np.int32(self.probe_positions[phase][slice_obj].shape[0])
                  if num_probes == 0: break
                  probes_d.append(cuda.mem_alloc(int(num_probes*np.prod(self.sampling)*8)))
                  plans.append(skfft.Plan(self.sampling, np.complex64, np.complex64, batch=num_probes, stream=stream))
                  norm_consts.append(cuda.mem_alloc(np.empty(num_probes,dtype=np.float32).nbytes))
          # 3. Propagate Beams
              for batch, stream, probe_d, plan,  norm_const in zip(batches, streams, probes_d, plans, norm_consts):
                  num_probes = np.int32(self.probe_positions[phase][batch].shape[0])
                  self.print_debug('batch: %s' % format(batch))
                  self.__propagate_beams(num_probes, batch, probe_d, propag_d, psi_k_d, norm_const, grid_steps_d, grid_range_d,
                                       self.probes[phase][batch], plan, ones_d, stream, transmit=transmit)
              ctx.synchronize()
          # 4. clean-up
              for plan, probe_d, norm_const in zip(plans, probes_d, norm_consts):
                 cufft.cufftDestroy(plan.handle)
                 probe_d.free()
                 norm_const.free()
                 # ctx.synchronize()
                 del probe_d, norm_const, plan
              ctx.synchronize()
              self.print_verbose('finished simulation phase #%d' % i)
              self.probes /= self.normalization
              # self.probes[phase][batch] = self.probes[phase][batch]/self.normalization
              sim_t = time()-t
              self.print_verbose('Propagated %d probes in %2.4f s' % (self.probe_positions[phase].shape[0], sim_t))
          
          if transmit:
              with catch_warnings():
                  simplefilter('ignore')
                  catch_warn()
                  self.probes = self.probes.astype(np.float32) # discard imaginary
          # clean up device variables
  #         self.vars.append(self.pot_dev_ptr) # don't de-allocate here to allow for simulations with different probe/exp. params
  #         self.pot_dev_ptr.free()
          mask_d.free()
          propag_d.free()
          psi_k_d.free()
          grid_range_d.free()
          grid_steps_d.free()
          ones_d.free()
          return None
      
      def build_potential_slices(self, slice_thickness):
          self.slice_t = slice_thickness
          self.num_slices = np.int32(np.floor(self.dims[-1] / slice_thickness))
          tasks = [((self, slice_num), {'method': 'build_slices'}) for slice_num in range(self.num_slices)]
          processes = min(mp.cpu_count(), self.num_slices)
          chunk = np.int(np.floor(self.num_slices / processes))
          pool = mp.Pool(processes=processes)
          jobs = pool.imap(unwrap, tasks, chunksize=chunk)
          self.potential_slices = np.array([j for j in jobs])
          pool.close()
          self.print_verbose('Built %d potential slices with shape:%s pixels' % (self.potential_slices.shape[0],
                                                                    format(self.potential_slices.shape[1:])))
  #         self.pot_dev_ptr = [cuda.to_device(pot_slice) for pot_slice in self.potential_slices]
          self.pot_slices_d = [cuda.to_device(pot_slice) for pot_slice in self.potential_slices]
          self.vars = []
  #         print(self.potential_slices)
  
  
      def __propagate_beams(self, num_probes, batch, psi_pos_d, propag_d, psi_k_d,
                          norm_const_d, grid_steps_d, grid_range_d,
                          psi_x_pos_pin, fft_plan_probe, ones_d, stream, transmit=True):
          """
          :param batch:
          :param psi_pos_d:
          :param propag_d:
          :param psi_k_d:
          :param norm_const_d:
          :param grid_steps_d:
          :param grid_range_d:
          :param psi_x_pos_pin:
          :param fft_plan_probe:
          :param stream:
          :return:
          """
          # define block/grid threads
          block_3d, grid_3d = self._get_blockgrid([self.sampling[1], self.sampling[0], num_probes], mode='3D')
          block_1d, grid_1d = self._get_blockgrid([self.sampling[1], self.sampling[0], num_probes], mode='1D')
          self.print_debug('block, grid:', block_3d, grid_3d)
  
          # grab needed kernels
          probe_stack_func = self.kernels['probes_stack']
          multwise2d_stack_func = self.kernels['mult_wise_c3d_c2d']
          multwise_stack_func = self.kernels['mult_wise_c3d_c2d']
          multwise_stack_func.prepare(("P", "P", np.int32, np.float32))
  #         multwise_stack_func = self.kernels['mult_wise_c3d_c3d_ind']
  #         multwise_stack_func.prepare(("P", "P", np.int32, np.int32, np.float32))
          multwise2d_stack_func.prepare(("P", "P", np.int32, np.float32))
          fftshift_func = self.kernels['fftshift_2d_stack']
          # norm_const_func = self.kernels['norm_const']
          normalize_func = self.kernels['normalize']
          mod_square_func = self.kernels['mod_square_stack']
  
          #1. build probes
          probe_stack_func(psi_pos_d, psi_k_d, num_probes, np.float32(self.kmax), grid_steps_d, grid_range_d,
                           block=block_3d, grid=grid_3d, shared=0, stream=stream)
          #ctx.synchronize()
          cufft.cufftExecC2C(fft_plan_probe.handle, int(psi_pos_d), int(psi_pos_d), cufft.CUFFT_INVERSE)
          #ctx.synchronize()
          fftshift_func(psi_pos_d, num_probes, block=block_3d, grid=grid_3d, shared=0, stream=stream)
          #ctx.synchronize()
          ### TODO: Check why cuda normalization is off
  
          # 2. Propagate probes through atomic potential
          # ctx.synchronize()
          if transmit:
              for i in range(self.num_slices):
                  # self.print_debug('Atomic potential slice #%d' % i)
                  multwise_stack_func.prepared_async_call(grid_3d, block_3d, stream, psi_pos_d, self.pot_slices_d[i], num_probes,
                                                      np.float32(1/np.prod(self.sampling)))
  
  #                 multwise_stack_func.prepared_async_call(grid_3d, block_3d, stream, psi_pos_d, self.pot_dev_ptr[i], num_probes, np.int32(i),
  #                                                         np.float32(1/np.prod(self.sampling)))
  #                 ctx.synchronize()
                  cufft.cufftExecC2C(fft_plan_probe.handle, int(psi_pos_d), int(psi_pos_d), cufft.CUFFT_FORWARD)
  #                 ctx.synchronize()
                  multwise2d_stack_func.prepared_async_call(grid_3d, block_3d, stream, psi_pos_d, propag_d, num_probes,
                                                          np.float32(1.0))
  #                 ctx.synchronize()
                  cufft.cufftExecC2C(fft_plan_probe.handle, int(psi_pos_d), int(psi_pos_d), cufft.CUFFT_INVERSE)
  #                 ctx.synchronize()
  
  
              multwise2d_stack_func.prepared_async_call(grid_3d, block_3d, stream, psi_pos_d, ones_d, num_probes,
                                                      np.float32(1 / (np.prod(self.sampling))))
              cufft.cufftExecC2C(fft_plan_probe.handle, int(psi_pos_d), int(psi_pos_d), cufft.CUFFT_FORWARD)
              multwise2d_stack_func.prepared_async_call(grid_3d, block_3d, stream, psi_pos_d, ones_d, num_probes,
                                                      np.float32(np.sqrt(np.prod(self.sampling))))
              mod_square_func(psi_pos_d, num_probes, block=block_3d, grid=grid_3d, stream=stream)
          # ctx.synchronize()
          cuda.memcpy_dtoh_async(psi_x_pos_pin, psi_pos_d, stream=stream)
  
  
  class MSAMPI(MSAGPU):
      def __init__(self, *args, **kwargs):
          global comm
          comm = MPI.COMM_WORLD
          self.size = comm.Get_size()
          self.rank = comm.Get_rank()
          super(MSAMPI, self).__init__(*args, **kwargs)
  
      def print_rank(self, *args, **kwargs):
          if self.rank == 0:
              print(*args, **kwargs)
  
      def print_verbose(self, *args, **kwargs):
          if self.verbose:
              self.print_rank(*args, **kwargs)
  
      def generate_probe_positions(self, *args, **kwargs):
          super(MSAMPI, self).generate_probe_positions(*args, **kwargs)
          data_parts = []
          part = self.probe_positions.shape[0] // self.size
          self.total_num_probes = self.probe_positions.shape[0]
          for rank in range(self.size):
              slice_obj = slice(rank * part, (rank + 1)* part)
              if rank == self.size - 1:
                  slice_obj = slice( rank * part, None)
              data_parts.append(slice_obj)
          self.data_parts = data_parts
          if self.rank == 0:
              data = [self.probe_positions[part] for part in data_parts]
              self.data_size = np.sum([dat.shape[0] for dat in data])
          else:
              data = None
          data = comm.scatter(data, root=0)
          self.probe_positions = data
          print("rank %d: my # of probe_positions %d" % (self.rank,
          self.probe_positions.shape[0]))
          self.num_probes = np.int32(self.probe_positions.shape[0])
  
      def multislice(self, *args, **kwargs):
          try:
              # h5_write= kwargs.pop('h5_write')
              h5_file = kwargs.pop('h5_write')
              h5_write = isinstance(h5_file, h5py.File)
          except KeyError:
              h5_write = False
          super(MSAMPI, self).multislice(*args, **kwargs)
          if h5_write:
              self.print_rank('writing to h5 file...')
              dset = h5_file.create_dataset('4D_CBED', (self.total_num_probes,
                                  self.sampling[0], self.sampling[1]), dtype=self.probes.dtype)
              # Workaround the large parallel I/O limitation
              ### TODO: (low priority) test if doing below in chunks over indices < 2GB speeds up the writing.
              for i, p in enumerate(range(*self.data_parts[self.rank].indices(self.total_num_probes))):
                  dset[p, :, :] = self.probes[i, :, :]
              self.print_rank('finished writing h5 file.')
          else:
              receive_buff = None
              self.print_verbose('rank %d: shape of calculated probes = %s' %(self.rank, format(self.probes.shape)))
              if self.rank == 0:
                  buff_shape = (self.total_num_probes, self.sampling[0], self.sampling[1])
                  self.print_rank('receive buffer shape: %s' %format(buff_shape))
                  receive_buff = np.empty(buff_shape, dtype=self.probes.dtype)
              # Calculate buffer size for a vector gather op
              count = []
              for data_part in self.data_parts:
                  start = data_part.start
                  if start is None:
                      start = 0
                  stop = data_part.stop
                  if stop is None:
                      stop = self.total_num_probes
                  count.append(stop - start)
              count = np.array(count)
              displ = np.cumsum(count)
              displ = np.append(displ[::-1],0)[::-1]
              xy_offset = np.prod(self.sampling)
              count *= xy_offset
              count = count.astype(np.int64)
              displ *= xy_offset
              displ = displ.astype(np.int64)
              self.print_verbose("rank %d: my # probes: %s %s, count: %d, displ: %d" %(self.rank, format(self.probes.shape), self.probes.dtype, count[self.rank], displ[self.rank]))
              # gather
              # comm.Gatherv(self.probes, [receive_buff, count, displ[:-1], MPI.C_FLOAT_COMPLEX], root=0)
              comm.Gatherv(self.probes, [receive_buff, count, displ[:-1], MPI.FLOAT], root=0)
  
  
              if self.rank == 0: self.print_rank('Gathered results of shape: %s' %format(receive_buff.shape))
  