text: import os
  import shutil
  import numpy as np
  import pandas as pd
  import itertools as it
  import umap
  import pywt
  from statsmodels.stats.diagnostic import lilliefors
  from copy import deepcopy
  from scipy.spatial.distance import mahalanobis
  from scipy import linalg
  from scipy.signal import find_peaks
  from scipy.stats import sem
  from sklearn.mixture import GaussianMixture
  from sklearn.decomposition import PCA
  from blechpy.utils import write_tools as wt, print_tools as pt, math_tools as mt, userIO
  from blechpy.dio import h5io
  from blechpy.analysis import clustering, spike_analysis as sas
  from blechpy.plotting import data_plot as dplt
  import datetime as dt
  
  
  def get_detection_threshold(filt_el):
      m = np.mean(filt_el)
      th = 5.0*np.median(np.abs(filt_el)/0.6745)
      return m-th
  
  def detect_spikes(filt_el, spike_snapshot = [0.5, 1.0], fs = 30000.0, thresh = None):
      '''Detects spikes in the filtered electrode trace and return the waveforms
      and spike_times
  
      Parameters
      ----------
      filt_el : np.array, 1-D
          filtered electrode trace
      spike_snapshot : list
          2-elements, [ms before spike minimum, ms after spike minimum] 
          time around spike to snap as waveform
      fs : float, sampling rate in Hz
  
      Returns
      -------
      waves : np.array
          matrix of de-jittered, spike waveforms, upsampled by 10x, row for each spike
      times : np.array
          array of spike times in samples
      threshold: float
          spike detection threshold
      '''
      # get indices of spike snapshot, expand by .1 ms in each direction
      snapshot = np.arange(-(spike_snapshot[0]+0.1)*fs/1000,
                           1+(spike_snapshot[1]+0.1)*fs/1000).astype('int64')
      if thresh is None:
          thresh = get_detection_threshold(filt_el)
  
      pos = np.where(filt_el <= thresh)[0]
      consecutive = mt.group_consecutives(pos)
  
      waves = []
      times = []
      for idx in consecutive:
          minimum = idx[np.argmin(filt_el[idx])]
          spike_idx = minimum + snapshot
          if spike_idx[0] >= 0 and spike_idx[-1] < len(filt_el):
              waves.append(filt_el[spike_idx])
              times.append(minimum)
  
      if len(waves) == 0:
          return None, None
  
      waves_dj, times_dj = clustering.dejitter(np.array(waves), np.array(times), spike_snapshot, fs)
      return waves_dj, times_dj, thresh
  
  
  def implement_pca(scaled_slices):
      pca = PCA()
      pca_slices = pca.fit_transform(scaled_slices)
      return pca_slices, pca.explained_variance_ratio_
  
  
  def implement_umap(waves, n_pc=3, n_neighbors=30, min_dist=0.0):
      reducer = umap.UMAP(n_components=n_pc,
                          n_neighbors=n_neighbors,
                          min_dist=min_dist)
      return reducer.fit_transform(waves)
  
  
  def implement_wavelet_transform(waves, n_pc=10):
      coeffs = pywt.wavedec(waves, 'haar', axis=1)
      all_coeffs = np.column_stack(coeffs)
      k_stats = np.zeros((all_coeffs.shape[1],))
      p_vals = np.ones((all_coeffs.shape[1],))
      for i, c in enumerate(all_coeffs.T):
          k_stats[i], p_vals[i] = lilliefors(c, dist='norm')
  
      idx = np.argsort(p_vals)
      return all_coeffs[:, idx[:n_pc]]
  
  
  def compute_waveform_metrics(waves, n_pc=3, umap=False):
      '''Make clustering data array with columns:
           - amplitudes, energy, slope, pc1, pc2, pc3, etc
      Parameters
      ----------
      waves : np.array
          waveforms with a row for each spike waveform
      n_pc : int (optional)
          number of principal components to include in data array
  
      Returns
      -------
      np.array
      '''
      data = np.zeros((waves.shape[0], 3))
      for i, wave in enumerate(waves):
          data[i,0] = np.min(wave)
          data[i,1] = np.sqrt(np.sum(wave**2))/len(wave)
          peaks = find_peaks(wave)[0]
          minima = np.argmin(wave)
          if not any(peaks < minima):
              maxima = np.argmax(wave[:minima])
          else:
              maxima = max(peaks[np.where(peaks < minima)[0]])
  
          data[i,2] = (wave[minima]-wave[maxima])/(minima-maxima)
  
      # Scale waveforms to energy before running PCA
      if umap:
          pc_waves = implement_umap(waves, n_pc=n_pc)
      else:
          scaled_waves = scale_waveforms(waves, energy=data[:,1])
          pc_waves, _ = implement_pca(scaled_waves)
  
      data = np.hstack((data, pc_waves[:,:n_pc]))
      data_columns = ['amplitude', 'energy', 'spike_slope']
      data_columns.extend(['PC%i' % i for i in range(n_pc)])
      return data, data_columns
  
  
  def get_waveform_amplitudes(waves):
      '''Returns array of waveform amplitudes
  
      Parameters
      ----------
      waves : np.array, matrix of waveforms, with row for each spike
  
      Returns
      -------
      np.array
      '''
      return np.min(waves,axis = 1)
  
  
  def get_waveform_energy(waves):
      '''Returns array of waveform energies
  
      Parameters
      ----------
      waves : np.array, matrix of waveforms, with row for each spike
  
      Returns
      -------
      np.array
      '''
      energy = np.sqrt(np.sum(waves**2, axis=1))/waves.shape[1]
      return energy
  
  
  def get_spike_slopes(waves):
      '''Returns array of spike slopes (initial downward slope of spike)
  
      Parameters
      ----------
      waves : np.array, matrix of waveforms, with row for each spike
  
      Returns
      -------
      np.array
      '''
      slopes = np.zeros((waves.shape[0],))
      for i, wave in enumerate(waves):
          peaks = find_peaks(wave)[0]
          minima = np.argmin(wave)
          if not any(peaks < minima):
              maxima = np.argmax(wave[:minima])
          else:
              maxima = max(peaks[np.where(peaks < minima)[0]])
  
          slopes[i] = (wave[minima]-wave[maxima])/(minima-maxima)
  
      return slopes
  
  
  def get_ISI_and_violations(spike_times, fs, rec_map=None):
      '''returns array of ISIs in ms and # of 1ms and 2ms violations
  
      Parameters
      ----------
      spike_time  numpy.array
      fs : float, sampling rate in Hz
      rec_map : np.array (optional)
          if not passed, it is assumed all spike times are from same recording
          if passed, spike times are split into recordings and ISIs are computed
          per recording.
          If fs is different for each recording, fs should be a dict with keys as
          rec ids in rec_map
  
      Returns
      -------
      np.array : ISIs
      int : 1ms violations
      int : 2ms violations
      '''
      if rec_map is not None:
          if not isinstance(fs, dict):
              fs = dict.fromkeys(np.unique(rec_map), fs)
  
          ISIs = np.array([])
          violations1 = 0
          violations2 = 0
          for i in np.unique(rec_map):
              idx = np.where(rec_map == i)[0]
              tmp_isi, v1, v2 = get_ISI_and_violations(spike_times[idx], fs[i])
              violations1 += v1
              violations2 += v2
              ISIs = np.concatenate((ISIs, tmp_isi))
  
      else:
          fs = float(fs/1000.0)
          ISIs = np.ediff1d(np.sort(spike_times))/fs
          violations1 = np.sum(ISIs < 1.0)
          violations2 = np.sum(ISIs < 2.0)
  
      return ISIs, violations1, violations2
  
  
  def scale_waveforms(waves, energy=None):
      '''Scales each waveform to its own energy
  
      Parameters
      ----------
      waves : np.array, matrix of waveforms, with row for each spike
      energy : np.array (optional)
          array of waveform energies, saves computation time
  
      Returns
      -------
      np.array
      '''
      if energy is None:
          energy = get_waveform_energy(waves)
      elif len(energy) != waves.shape[0]:
          raise ValueError(('Energies must correspond to each waveforms.'
                            'Different lengths are not allowed'))
  
      scaled_slices = np.zeros(waves.shape)
      for i, w in enumerate(zip(waves, energy)):
          scaled_slices[i] = w[0]/w[1]
  
      return scaled_slices
  
  
  def get_mahalanobis_distances_to_cluster(data, model, clusters, target_cluster):
      '''computes mahalanobis distance from spikes in target_cluster to all clusters
      in GMM model
  
      Parameters
      ----------
      data : np.array, data used to train GMM
      model : fitted GMM model
      clusters : np.array, maps data points to clusters
      target_cluster : int, cluster for which to compute distances
  
      Returns
      -------
      np.array
      '''
      unique_clusters = np.unique(abs(clusters))
      out_distances = dict.fromkeys(unique_clusters)
      cluster_idx = np.where(clusters == target_cluster)[0]
      for other_cluster in unique_clusters:
          mahalanobis_dist = np.zeros((len(cluster_idx),))
          other_cluster_mean = model.means_[other_cluster, :]
          other_cluster_covar_I = linalg.inv(model.covariances_[other_cluster, :, :])
          for i, idx in enumerate(cluster_idx):
              mahalanobis_dist[i] = mahalanobis(data[idx, :],
                                                other_cluster_mean,
                                                other_cluster_covar_I)
  
          out_distances[other_cluster] = mahalanobis_dist
  
      return out_distances
  
  
  def get_recording_cutoff(filt_el, sampling_rate, voltage_cutoff,
                           max_breach_rate, max_secs_above_cutoff,
                           max_mean_breach_rate_persec, **kwargs):
      breach_idx = np.where(filt_el > voltage_cutoff)[0]
      breach_rate = float(len(breach_idx)*int(sampling_rate))/len(filt_el)
      # truncate to nearest second and make 1 sec bins
      filt_el = filt_el[:int(sampling_rate)*int(len(filt_el)/sampling_rate)]
      test_el = np.reshape(filt_el, (-1, int(sampling_rate)))
      breaches_per_sec = [len(np.where(test_el[i] > voltage_cutoff)[0])
                          for i in range(len(test_el))]
      breaches_per_sec = np.array(breaches_per_sec)
      secs_above_cutoff = len(np.where(breaches_per_sec > 0)[0])
      if secs_above_cutoff == 0:
          mean_breach_rate_persec = 0
      else:
          mean_breach_rate_persec = np.mean(breaches_per_sec[np.where(breaches_per_sec > 0)[0]])
  
      # And if they all exceed the cutoffs, assume that the headstage fell off mid-experiment
      recording_cutoff = int(len(filt_el)/sampling_rate) # cutoff in seconds
      if (breach_rate >= max_breach_rate and
          secs_above_cutoff >= max_secs_above_cutoff and
          mean_breach_rate_persec >= max_mean_breach_rate_persec):
          # Find the first 1 second epoch where the number of cutoff breaches is
          # higher than the maximum allowed mean breach rate 
          recording_cutoff = np.where(breaches_per_sec > max_mean_breach_rate_persec)[0][0]
          # cutoff is still in seconds since 1 sec bins
  
      return recording_cutoff
  
  
  def UMAP_METRICS(waves, n_pc):
      return compute_waveform_metrics(waves, n_pc, umap=True)
  
  
  class SpikeDetection(object):
      '''Interface to manage spike detection and data extraction in preparation
      for GMM clustering. Intended to help create and access the neccessary
      files. If object will detect is file already exist to avoid re-creation
      unless overwrite is specified as True.
      '''
  
      def __init__(self, file_dir, electrode, params=None, overwrite=False):
          # Setup paths to files and directories needed
          self._file_dir = file_dir
          self._electrode = electrode
          self._out_dir = os.path.join(file_dir, 'spike_detection',
                                       'electrode_%i' % electrode)
          self._data_dir = os.path.join(self._out_dir, 'data')
          self._plot_dir = os.path.join(self._out_dir, 'plots')
          self._files = {'params': os.path.join(file_dir,'analysis_params', 'spike_detection_params.json'),
                         'spike_waveforms': os.path.join(self._data_dir, 'spike_waveforms.npy'),
                         'spike_times' : os.path.join(self._data_dir, 'spike_times.npy'),
                         'energy' : os.path.join(self._data_dir, 'energy.npy'),
                         'spike_amplitudes' : os.path.join(self._data_dir, 'spike_amplitudes.npy'),
                         'pca_waveforms' : os.path.join(self._data_dir, 'pca_waveforms.npy'),
                         'slopes' : os.path.join(self._data_dir, 'spike_slopes.npy'),
                         'recording_cutoff' : os.path.join(self._data_dir, 'cutoff_time.txt'),
                         'detection_threshold' : os.path.join(self._data_dir, 'detection_threshold.txt')}
  
          self._status = dict.fromkeys(self._files.keys(), False)
          self._referenced = True
  
          # Delete existing data if overwrite is True
          if overwrite and os.path.isdir(self._out_dir):
              shutil.rmtree(self._out_dir)
  
          # See what data already exists
          self._check_existing_files()
  
          # Make directories if needed
          if not os.path.isdir(self._out_dir):
              os.makedirs(self._out_dir)
  
          if not os.path.isdir(self._data_dir):
              os.makedirs(self._data_dir)
  
          if not os.path.isdir(self._plot_dir):
              os.makedirs(self._plot_dir)
  
          if not os.path.isdir(os.path.join(file_dir, 'analysis_params')):
              os.makedirs(os.path.join(file_dir, 'analysis_params'))
  
          # grab recording cutoff time if it already exists
          # cutoff should be in seconds
          self.recording_cutoff = None
          if os.path.isfile(self._files['recording_cutoff']):
              self._status['recording_cutoff'] = True
              with open(self._files['recording_cutoff'], 'r') as f:
                  self.recording_cutoff = float(f.read())
  
          self.detection_threshold = None
          if os.path.isfile(self._files['detection_threshold']):
              self._status['detection_threshold'] = True
              with open(self._files['detection_threshold'], 'r') as f:
                  self.detection_threshold = float(f.read())
  
          # Read in parameters
          # Parameters passed as an argument will overshadow parameters saved in file
          # Input parameters should be formatted as dataset.clustering_parameters
          if params is None and os.path.isfile(self._files['params']):
              self.params = wt.read_dict_from_json(self._files['params'])
          elif params is None:
              raise FileNotFoundError('params must be provided if spike_detection_params.json does not exist.')
          else:
              self.params = {}
              self.params['voltage_cutoff'] = params['data_params']['V_cutoff for disconnected headstage']
              self.params['max_breach_rate'] = params['data_params']['Max rate of cutoff breach per second']
              self.params['max_secs_above_cutoff'] = params['data_params']['Max allowed seconds with a breach']
              self.params['max_mean_breach_rate_persec'] = params['data_params']['Max allowed breaches per second']
              band_lower = params['bandpass_params']['Lower freq cutoff']
              band_upper = params['bandpass_params']['Upper freq cutoff']
              self.params['bandpass'] = [band_lower, band_upper]
              snapshot_pre = params['spike_snapshot']['Time before spike (ms)']
              snapshot_post = params['spike_snapshot']['Time after spike (ms)']
              self.params['spike_snapshot'] = [snapshot_pre, snapshot_post]
              self.params['sampling_rate'] = params['sampling_rate']
              # Write params to json file
              wt.write_dict_to_json(self.params, self._files['params'])
              self._status['params'] = True
  
      def _check_existing_files(self):
          '''Checks which files already exist and updates _status so as to avoid
          re-creation later
          '''
          for k, v in self._files.items():
              if os.path.isfile(v):
                  self._status[k] = True
              else:
                  self._status[k] = False
  
      def run(self):
          status = self._status
          file_dir = self._file_dir
          electrode = self._electrode
          params = self.params
          fs = params['sampling_rate']
  
          # Check if this even needs to be run
          if all(status.values()):
              return electrode, 1, self.recording_cutoff
  
          # Grab referenced electrode or raw if ref is not available
          ref_el = h5io.get_referenced_trace(file_dir, electrode)
          if ref_el is None:
              print('Could not find referenced data for electrode %i. Using raw.' % electrode)
              self._referenced = False
              ref_el = h5io.get_raw_trace(file_dir, electrode)
              if ref_el is None:
                  raise KeyError('Neither referenced nor raw data found for electrode %i in %s' % (electrode, file_dir))
  
          # Filter electrode trace
          filt_el = clustering.get_filtered_electrode(ref_el, freq=params['bandpass'],
                                                 sampling_rate = fs)
          del ref_el
          # Get recording cutoff
          if not status['recording_cutoff']:
              self.recording_cutoff = get_recording_cutoff(filt_el, **params)
              with open(self._files['recording_cutoff'], 'w') as f:
                  f.write(str(self.recording_cutoff))
  
              status['recording_cutoff'] = True
              fn = os.path.join(self._plot_dir, 'cutoff_time.png')
              dplt.plot_recording_cutoff(filt_el, fs, self.recording_cutoff,
                                         out_file=fn)
  
          # Truncate electrode trace, deal with early cutoff (<60s)
          if self.recording_cutoff < 60:
              print('Immediate Cutoff for electrode %i...exiting' % electrode)
              return electrode, 0, self.recording_cutoff
  
          filt_el = filt_el[:int(self.recording_cutoff*fs)]
  
          if not status['detection_threshold']:
              threshold = get_detection_threshold(filt_el)
              self.detection_threshold = threshold
              with open(self._files['detection_threshold'], 'w') as f:
                  f.write(str(threshold))
  
              status['detection_threshold'] = True
  
          if status['spike_waveforms'] and status['spike_times']:
              waves = np.load(self._files['spike_waveforms'])
              times = np.load(self._files['spike_times'])
          else:
              # Detect spikes and get dejittered times and waveforms
              # detect_spikes returns waveforms upsampled by 10x and times in units
              # of samples
              waves, times, threshold = detect_spikes(filt_el, params['spike_snapshot'], fs, thresh=self.detection_threshold)
              if waves is None:
                  print('No waveforms detected on electrode %i' % electrode)
                  return electrode, 0, self.recording_cutoff
  
              # Save waveforms and times
              np.save(self._files['spike_waveforms'], waves)
              np.save(self._files['spike_times'], times)
              status['spike_waveforms'] = True
              status['spike_times'] = True
  
          # Get various metrics and scale waveforms
          if not status['spike_amplitudes']:
              amplitudes = get_waveform_amplitudes(waves)
              np.save(self._files['spike_amplitudes'], amplitudes)
              status['spike_amplitudes'] = True
  
          if not status['slopes']:
              slopes = get_spike_slopes(waves)
              np.save(self._files['slopes'], slopes)
              status['slopes'] = True
  
          if not status['energy']:
              energy = get_waveform_energy(waves)
              np.save(self._files['energy'], energy)
              status['energy'] = True
          else:
              energy=None
  
          # get pca of scaled waveforms
          if not status['pca_waveforms']:
              scaled_waves = scale_waveforms(waves, energy=energy)
              pca_waves, explained_variance_ratio = implement_pca(scaled_waves)
  
              # Plot explained variance
              fn = os.path.join(self._plot_dir, 'pca_variance.png')
              dplt.plot_explained_pca_variance(explained_variance_ratio,
                                               out_file = fn)
  
          return electrode, 1, self.recording_cutoff
  
      def get_spike_waveforms(self):
          '''Returns spike waveforms if they have been extracted, None otherwise
          Dejittered waveforms upsampled to 10 x sampling_rate
  
          Returns
          -------
          numpy.array
          '''
          if os.path.isfile(self._files['spike_waveforms']):
              return np.load(self._files['spike_waveforms'])
          else:
              return None
  
      def get_spike_times(self):
          '''Returns spike times if they have been extracted, None otherwise
          In units of samples.
  
          Returns
          -------
          numpy.array
          '''
          if os.path.isfile(self._files['spike_times']):
              return np.load(self._files['spike_times'])
          else:
              return None
  
      def get_energy(self):
          '''Returns spike energies if they have been extracted, None otherwise
  
          Returns
          -------
          numpy.array
          '''
          if os.path.isfile(self._files['energy']):
              return np.load(self._files['energy'])
          else:
              return None
  
      def get_spike_amplitudes(self):
          '''Returns spike amplitudes if they have been extracted, None otherwise
  
          Returns
          -------
          numpy.array
          '''
          if os.path.isfile(self._files['spike_amplitudes']):
              return np.load(self._files['spike_amplitudes'])
          else:
              return None
  
      def get_spike_slopes(self):
          '''Returns spike slopes if they have been extracted, None otherwise
  
          Returns
          -------
          numpy.array
          '''
          if os.path.isfile(self._files['slopes']):
              return np.load(self._files['slopes'])
          else:
              return None
  
      def get_pca_waveforms(self):
          '''Returns pca of sclaed spike waveforms if they have been extracted,
          None otherwise
          Dejittered waveforms upsampled to 10 x sampling_rate, scaled to energy
          and transformed via PCA
  
          Returns
          -------
          numpy.array
          '''
          if os.path.isfile(self._files['spike_waveforms']):
              return np.load(self._files['spike_waveforms'])
          else:
              return None
  
      def get_clustering_metrics(self, n_pc=3):
          '''Returns array of metrics to use for feature based clustering
          Row for each waveform with columns:
              - amplitude, energy, spike slope, PC1, PC2, etc
          '''
          amplitude = self.get_spike_amplitudes()
          energy = self.get_energy()
          slopes = self.get_spike_slopes()
          pca_waves = self.get_pca_waveforms()
          out = np.vstack((amplitude, energy, slopes)).T
          out = np.hstack((out, pca_waves[:,:n_pc]))
          return out
  
      def __str__(self):
          out = []
          out.append('SpikeDetection\n--------------')
          out.append('Recording Directory: %s' % self._file_dir)
          out.append('Electrode: %i' % self._electrode)
          out.append('Output Directory: %s' % self._out_dir)
          out.append('###################################\n')
          out.append('Status:')
          out.append(pt.print_dict(self._status))
          out.append('-------------------\n')
          out.append('Parameters:')
          out.append(pt.print_dict(self.params))
          out.append('-------------------\n')
          out.append('Data files:')
          out.append(pt.print_dict(self._files))
          return '\n'.join(out)
  
  
  class BlechClust(object):
      def __init__(self, rec_dirs, electrode, out_dir=None, params=None,
                   overwrite=False, no_write=False, n_pc=3,
                   data_transform=compute_waveform_metrics):
          '''Recording directories should be ordered to make spike sorting easier later on
          '''
          if isinstance(rec_dirs, str):
              rec_dirs = [rec_dirs]
  
          rec_dirs = [x[:-1] if x.endswith(os.sep) else x for x in rec_dirs]
          self.rec_dirs = rec_dirs
          self.electrode = electrode
          self._data_transform = data_transform
          self._n_pc = n_pc
          if out_dir is None:
              if len(rec_dirs) > 1:
                  top = os.path.dirname(rec_dirs[0])
                  out_dir = os.path.join(top, 'BlechClust', 'electrode_%i' % electrode)
              else:
                  out_dir = os.path.join(rec_dirs[0], 'BlechClust', 'electrode_%i' % electrode)
  
          if overwrite:
              shutil.rmtree(out_dir)
  
          # Make directories
          self.out_dir = out_dir
          self._plot_dir = os.path.join(out_dir, 'plots')
          self._data_dir = os.path.join(out_dir, 'clustering_results')
          if not os.path.isdir(out_dir):
              os.makedirs(out_dir)
  
          if not os.path.isdir(self._data_dir):
              os.mkdir(self._data_dir)
  
          if not os.path.isdir(self._plot_dir):
              os.mkdir(self._plot_dir)
  
          # Check files
          params_file = os.path.join(out_dir, 'BlechClust_params.json')
          map_file = os.path.join(self._data_dir, 'spike_id.npy')
          key_file = os.path.join(self._data_dir, 'rec_key.json')
          results_file = os.path.join(self._data_dir, 'clustering_results.json')
          self._files = {'params': params_file, 'spike_map': map_file,
                         'rec_key': key_file, 'clustering_results': results_file}
          self.params = params
          self._load_existsing_data()
  
          if self._rec_key is None and not no_write:
              # Create new rec key
              rec_key = {x:y for x,y in enumerate(self.rec_dirs)}
              self._rec_key = rec_key
              wt.write_dict_to_json(rec_key, self._files['rec_key'])
          elif self._rec_key is None:
              ValueError('Existing rec_key not found and no_write is enabled')
  
          # Check to see if spike detection is already completed on all recording directories
          spike_check = self._check_spike_detection()
          if not all(spike_check):
              invalid = [rec_dirs[i] for i, x in enumerate(spike_check) if x==False]
              error_str = '\n\t'.join(invalid)
              raise ValueError('Spike detection has not been run on:\n\t%s' % error_str)
  
      def _load_existsing_data(self):
          params = self.params
          file_check = self._check_existing_files()
  
          # Check params files and create if new params are passed
          if file_check['params']:
              self.params = wt.read_dict_from_json(self._files['params'])
  
          # Make new params or overwrite existing with passed params
          if params is None and not file_check['params']:
              raise ValueError(('Params file does not exists at %s. Must provide'
                                ' clustering parameters.') % self._files['params'])
          elif params is not None:
              self.params['max_clusters'] = params['clustering_params']['Max Number of Clusters']
              self.params['max_iterations'] = params['clustering_params']['Max Number of Iterations']
              self.params['threshold'] = params['clustering_params']['Convergence Criterion']
              self.params['num_restarts'] = params['clustering_params']['GMM random restarts']
              self.params['wf_amplitude_sd_cutoff'] = params['data_params']['Intra-cluster waveform amp SD cutoff']
              wt.write_dict_to_json(self.params, self._files['params'])
  
          # Deal with existing rec key
          if file_check['rec_key']:
              rec_dirs = self.rec_dirs
              rec_key = wt.read_dict_from_json(self._files['rec_key'])
              rec_key = {int(x): y for x,y in rec_key.items()}
              if len(rec_key) != len(rec_dirs):
                  raise ValueError('Rec key does not match rec dirs')
  
              # Correct rec key in case rec_dir roots have changed
              for rd in rec_dirs:
                  rn = os.path.basename(rd)
                  dn = os.path.dirname(rd)
                  kd = [(x, y) for x,y in rec_key.items() if rn in y]
                  if len(kd) == 0:
                      raise ValueError('%s not found in rec_key' % rn)
  
                  kd = kd[0]
                  if kd[1] != rd:
                      rec_key[kd[0]] = rd
  
              inverted = {v:k for k,v in rec_key.items()}
              self.rec_dirs = sorted(self.rec_dirs, key=lambda i: inverted[i])
              self._rec_key = rec_key
          else:
              self._rec_key = None
  
          # Check is clustering has already been done, load results
          if file_check['clustering_results']:
              self.results = wt.read_pandas_from_table(self._files['clustering_results'])
              expected_results = np.arange(2, self.params['max_clusters'] + 1)
              if not all([x in self.results['clusters'] for x in expected_results]):
                  self.clustered = False
              else:
                  self.clustered = True
  
          else:
              self.results = None
              self.clustered = False
  
      def _check_existing_files(self):
          out = dict.fromkeys(self._files.keys(), False)
          for k,v in self._files.items():
              if os.path.isfile(v):
                  out[k] = True
  
          return out
  
      def _check_spike_detection(self):
          '''Check to see if spike detection has been run on all recording directories
          '''
          out = []
          for rec in  self.rec_dirs:
              try:
                  spike_detect = SpikeDetection(rec, self.electrode)
                  if all(spike_detect._status):
                      out.append(True)
                  else:
                      out.append(False)
  
              except FileNotFoundError:
                  out.append(False)
  
          return out
  
      def run(self, n_pc=None, overwrite=False):
          if self.clustered and not overwrite:
              return True
  
          if n_pc is None:
              n_pc = self._n_pc
  
          GMM = ClusterGMM(self.params['max_iterations'],
                           self.params['num_restarts'], self.params['threshold'])
  
          # Collect data from all recordings
          waveforms, spike_times, spike_map, fs, offsets = self.get_spike_data()
  
          # Save array to map spikes and predictions back to original recordings
          np.save(self._files['spike_map'], spike_map)
  
          data, data_columns = self._data_transform(waveforms, n_pc)
          amplitudes = get_waveform_amplitudes(waveforms)
  
          # Run GMM for each number of clusters from 2 to max_clusters
          tested_clusters = np.arange(2, self.params['max_clusters']+1)
          clust_results = pd.DataFrame(columns=['clusters','converged',
                                                'BIC','spikes_per_cluster'],
                                       index=tested_clusters)
          for n_clust in tested_clusters:
              data_dir = os.path.join(self._data_dir, '%i_clusters' % n_clust)
              plot_dir = os.path.join(self._plot_dir, '%i_clusters' % n_clust)
              wave_plot_dir = os.path.join(self._plot_dir, '%i_clusters_waveforms_ISIs' % n_clust)
              bic_file = os.path.join(data_dir, 'bic.npy')
              pred_file = os.path.join(data_dir, 'predictions.npy')
  
              if os.path.isfile(bic_file) and os.path.isfile(pred_file) and not overwrite:
                  bic = np.load(bic_file)
                  predictions = np.load(pred_file)
                  spikes_per_clust = [len(np.where(predictions == c)[0])
                                      for c in np.unique(predictions)]
                  clust_results.loc[n_clust] = [n_clust, True, bic, spikes_per_clust]
                  continue
  
              if not os.path.isdir(wave_plot_dir):
                  os.makedirs(wave_plot_dir)
  
              if not os.path.isdir(data_dir):
                  os.makedirs(data_dir)
  
              if not os.path.isdir(plot_dir):
                  os.makedirs(plot_dir)
  
              model, predictions, bic = GMM.fit(data, n_clust)
              if model is None:
                  clust_results.loc[n_clust] = [n_clust, bic, False, [0]]
                  # Nothing converged
                  continue
  
              # Go through each cluster and throw out any spikes too far from the
              # mean
              spikes_per_clust = []
              for c in range(n_clust):
                  idx = np.where(predictions == c)[0]
                  mean_amp = np.mean(amplitudes[idx])
                  sd_amp = np.std(amplitudes[idx])
                  cutoff_amp = mean_amp - (sd_amp * self.params['wf_amplitude_sd_cutoff'])
                  rejected_idx = np.array([i for i in idx if amplitudes[i] <= cutoff_amp])
                  if len(rejected_idx) > 0:
                      predictions[rejected_idx] = -1
  
                  idx = np.where(predictions == c)[0]
                  spikes_per_clust.append(len(idx))
  
                  if len(idx) == 0:
                      continue
  
                  # Plot waveforms and ISIs of cluster
                  ISIs, violations_1ms, violations_2ms = get_ISI_and_violations(spike_times[idx], fs, spike_map[idx])
                  cluster_waves = waveforms[idx]
                  cluster_times = spike_times[idx]
                  isi_fn = os.path.join(wave_plot_dir, 'Cluster%i_ISI.png' % c)
                  wave_fn = os.path.join(wave_plot_dir, 'Cluster%i_waveforms.png' % c)
                  title_str = ('Cluster%i\nviolations_1ms = %i, '
                               'violations_2ms = %i\n'
                               'Number of waveforms = %i' %
                               (c, violations_1ms, violations_2ms, len(idx)))
                  dplt.plot_waveforms(cluster_waves, title=title_str, save_file=wave_fn)
                  if len(ISIs) > 0:
                      dplt.plot_ISIs(ISIs, total_spikes=len(idx), save_file=isi_fn)
  
  
              clust_results.loc[n_clust] = [n_clust, True, bic, spikes_per_clust]
  
              # Plot feature pairs
              feature_pairs = it.combinations(list(range(data.shape[1])), 2)
              for f1, f2 in feature_pairs:
                  fn = '%sVS%s.png' % (data_columns[f1], data_columns[f2])
                  fn = os.path.join(plot_dir, fn)
                  dplt.plot_cluster_features(data[:, [f1,f2]], predictions,
                                             x_label = data_columns[f1],
                                             y_label = data_columns[f2],
                                             save_file = fn)
  
              # For each cluster plot mahanalobis distances to all other clusters
              for c in range(n_clust):
                  distances = get_mahalanobis_distances_to_cluster(data,  model,
                                                                   predictions, c)
                  fn = os.path.join(plot_dir, 'Mahalanobis_cluster%i.png' % c)
                  title = ('Mahalanobis distance of Cluster %i from all other clusters' % c)
                  dplt.plot_mahalanobis_to_cluster(distances, title=title, save_file=fn)
  
              # Save data
              np.save(bic_file, bic)
              np.save(pred_file, predictions)
  
          # Save results table
          self.results = clust_results
          wt.write_pandas_to_table(clust_results,
                                   self._files['clustering_results'],
                                   overwrite=True)
          self.clustered = True
          return True
  
      def get_spike_data(self):
          # Collect data from all recordings
          tmp_waves = []
          tmp_times = []
          tmp_id = []
          fs = dict.fromkeys(self._rec_key.keys())
          offsets = dict.fromkeys(self._rec_key.keys())
          offset = 0
          for i in sorted(self._rec_key.keys()):
              rec = self._rec_key[i]
              spike_detect = SpikeDetection(rec, self.electrode)
              t = spike_detect.get_spike_times()
              fs[i] = spike_detect.params['sampling_rate']
              if t is None:
                  offsets[i] = int(offset)
                  offset = offset + 3*fs[i]
                  continue
  
              tmp_waves.append(spike_detect.get_spike_waveforms())
              tmp_times.append(t)
              tmp_id.append(np.ones((t.shape[0],))*i)
              offsets[i] = int(offset)
              offset = offset + max(t) + 3*fs[i]
  
          waveforms = np.vstack(tmp_waves)
          spike_times = np.hstack(tmp_times)
          spike_map = np.hstack(tmp_id)
  
          # Double check that spike_map matches up with existing spike_map
          if os.path.isfile(self._files['spike_map']):
              orig_map = np.load(self._files['spike_map'])
              if len(orig_map) != len(spike_map):
                  raise ValueError('Spike detection has changed, please re-cluster with overwrite=True')
  
          return waveforms, spike_times, spike_map, fs, offsets
  
      def get_clusters(self, solution_num, cluster_nums):
          if not isinstance(cluster_nums, list):
              cluster_nums = [cluster_nums]
  
          waveforms, times, spike_map, fs, offsets = self.get_spike_data()
          predictions = self.get_predictions(solution_num)
          out = []
          for c in cluster_nums:
              idx = np.where(predictions == c)[0]
              if len(idx)==0:
                  continue
  
              tmp_clust = SpikeCluster('Cluster_%i' % c,
                                       self.electrode,
                                       solution_num,
                                       c,
                                       1,
                                       waveforms[idx],
                                       times[idx],
                                       spike_map[idx],
                                       self._rec_key.copy(),
                                       fs.copy(),
                                       offsets.copy(),
                                       manipulations='')
              out.append(tmp_clust)
  
          return out
  
      def get_predictions(self, n_clusters):
          fn = os.path.join(self._data_dir, '%i_clusters' % n_clusters,
                            'predictions.npy')
          if os.path.isfile(fn):
              return np.load(fn)
          else:
              return None
  
  
  class ClusterGMM(object):
      def __init__(self, n_iters, n_restarts, thresh):
          self.params = {'iterations': n_iters,
                         'restarts': n_restarts,
                         'thresh': thresh}
  
      def fit(self, data, n_clusters):
          min_bic = None
          best_model = None
          if n_clusters is not None:
              self.params['clusters'] = n_clusters
  
          for i in range(self.params['restarts']):
              model = GaussianMixture(n_components = self.params['clusters'],
                                      covariance_type = 'full',
                                      tol = self.params['thresh'],
                                      random_state = i,
                                      max_iter = self.params['iterations'])
              model.fit(data)
              if model.converged_:
                  new_bic = model.bic(data)
                  if min_bic is None:
                      min_bic = model.bic(data)
                      best_model = model
                  elif new_bic < min_bic:
                      best_model = model
                      min_bic = new_bic
  
          predictions = best_model.predict(data)
          self._model = best_model
          self._predictions = predictions
          self._bic = min_bic
          return best_model, predictions, min_bic
  
  
  class SpikeSorter(object):
      def __init__(self, rec_dirs, electrode, clustering_dir=None, shell=False):
          if isinstance(rec_dirs, str):
              rec_dirs = [rec_dirs]
  
          rec_dirs = [x[:-1] if x.endswith(os.sep) else x for x in rec_dirs]
          self.rec_dirs = rec_dirs
          self.electrode = electrode
          if clustering_dir is None:
              if len(rec_dirs) > 1:
                  top = os.path.dirname(rec_dirs[0])
                  clustering_dir = os.path.join(top, 'BlechClust', 'electrode_%i' % electrode)
              else:
                  clustering_dir = os.path.join(rec_dirs[0], 'BlechClust', 'electrode_%i' % electrode)
  
          self.clustering_dir = clustering_dir
          try:
              clust = BlechClust(rec_dirs, electrode, out_dir = clustering_dir, no_write=True)
          except FileNotFoundError:
              clust = None
  
          if clust is None or not clust.clustered:
              raise ValueError('Recordings have not been clustered yet.')
  
          # Match recording directory ordering to clustering
          self.rec_dirs = clust.rec_dirs
          self.clustering = clust
          self._current_solution = None
          self._active = None
          self._last_saved = None
          self._previous = None
          self._shell = shell
          self._split_results = None
          self._split_starter = None
          self._split_index = None
          self._last_umap_embedding = None
          self._last_action = None
          self._last_popped = None  # Dict of indices to clusters
          self._last_added = None  # List of indices
  
          thresh = []
          for rd in rec_dirs:
              sd = SpikeDetection(rd, electrode)
              thresh.append(sd.detection_threshold)
  
          self._detection_thresholds = thresh
  
      def undo(self):
          if self._last_action is None:
              return
  
          if self._last_action == 'save':
              self.undo_last_save()
              return
  
          # Remove last added
          for k in reversed(sorted(self._last_added)):
              self._active.pop(k)
  
          # Insert previous clusters
          for k in sorted(self._last_popped.keys()):
              self._active.insert(k, self._last_popped[k])
  
          # reset
          self._last_action = None
          self._last_popped = None
          self._last_added = None
  
      def set_active_clusters(self, solution_num):
          self._current_solution = solution_num
          cluster_nums = list(range(solution_num))
          clusters = self.clustering.get_clusters(solution_num, cluster_nums)
          if len(clusters) == 0:
              raise ValueError('Solution or clusters not found')
  
          self._active = clusters
          self._last_action = None
          self._last_popped = None
          self._last_added = None
  
      def save_clusters(self, target_clusters, single_unit, pyramidal, interneuron):
          '''Saves active clusters as cells, write them to the h5_files in the
          appropriate recording directories
  
          Parameters
          ----------
          target_clusters: list of int
              indicies of active clusters to save
          single_unit : list of bool
              elements in list must correspond to elements in active clusters
          pyramidal : list of bool
          interneuron : list of bool
          '''
          if self._active is None:
              return
  
          if any([i >= len(self._active) for i in target_clusters]):
              raise ValueError('Target cluster is out of range.')
  
          n_clusters = len(target_clusters)
          if (len(single_unit) != n_clusters or len(pyramidal) != n_clusters or
              len(interneuron) != n_clusters):
              raise ValueError('Length of input lists must match number of '
                               'active clusters. Expected %i' % n_clusters)
  
          self._last_action = 'save'
          self._last_popped = {i: self._active[i] for i in target_clusters}
          self._last_added = []
          clusters = [self._active[i] for i in target_clusters]
          rec_key = self.clustering._rec_key
          self._last_saved = dict.fromkeys(rec_key.keys(), None)
  
          for clust, single, pyr, intr in zip(clusters, single_unit,
                                              pyramidal, interneuron):
              for i, rec in rec_key.items():
                  idx = np.where(clust['spike_map'] == i)[0]
                  if len(idx) == 0:
                      continue
  
                  waves = clust['spike_waveforms'][idx]
                  times = clust['spike_times'][idx]
                  unit_name = h5io.add_new_unit(rec, self.electrode, waves,
                                                times, single, pyr, intr)
                  if self._last_saved[i] is None:
                      self._last_saved[i] = [unit_name]
                  else:
                      self._last_saved[i].append(unit_name)
  
                  metrics_dir = os.path.join(rec,'sorted_unit_metrics', unit_name)
                  if not os.path.isdir(metrics_dir):
                      os.makedirs(metrics_dir)
  
                  # Write cluster info to file
                  print_clust = clust.copy()
                  for k,v in clust.items():
                      if isinstance(v, np.ndarray):
                          print_clust.pop(k)
  
                  print_clust.pop('rec_key')
                  print_clust.pop('fs')
                  clust_info_file = os.path.join(metrics_dir, 'cluster.info')
                  with open(clust_info_file, 'a+') as log:
                      print('%s sorted on %s'
                            % (unit_name,
                               dt.datetime.today().strftime('%m/%d/%y %H:%M')),
                            file=log)
                      print('Cluster info: \n----------', file=log)
                      print(pt.print_dict(print_clust), file=log)
                      print('Saved metrics to %s' % metrics_dir, file=log)
                      print('--------------\n', file=log)
  
          userIO.tell_user('Target clusters successfully saved to recording '
                           'directories.', shell=True)
          self._active = [self._active[i] for i in range(len(self._active))
                          if i not in target_clusters]
  
      def undo_last_save(self):
          if self._last_saved is None:
              return
  
          rec_key = self.clustering._rec_key
          last_saved = self._last_saved
          for i, rec in rec_key.items():
              for unit in reversed(np.sort(last_saved[i])):
                  h5io.delete_unit(rec, unit)
  
          for k in sorted(self._last_popped.keys()):
              self._active.insert(k, self._last_popped[k])
  
          self._last_saved = None
          self._last_popped = None
          self._last_added = None
          self._last_action = None
  
      def split_cluster(self, target_clust, n_iter, n_restart, thresh, n_clust,
                        store_split=False, umap=False):
          '''splits the target active cluster using a GMM
          '''
          if target_clust >= len(self._active):
              raise ValueError('Invalid target. Only %i active clusters' % len(self._active))
  
          cluster = self._active.pop(target_clust)
          self._split_starter = cluster
          self._split_index = target_clust
  
          try:
              GMM = ClusterGMM(n_iter, n_restart, thresh)
              waves = cluster['spike_waveforms']
              data, data_columns = compute_waveform_metrics(waves, umap=umap)
              model, predictions, bic = GMM.fit(data, n_clust)
              new_clusts = []
              for i in np.unique(predictions):
                  idx = np.where(predictions == i)[0]
                  edit_str = (cluster['manipulations'] + '\nSplit %s into %i '
                              'clusters. This is sub-cluster %i'
                              % (cluster['Cluster_Name'], n_clust, i))
                  tmp_clust = SpikeCluster(cluster['Cluster_Name'] + '-%i' % i,
                                           cluster['electrode_num'],
                                           cluster['solution_num'],
                                           cluster['cluster_num'],
                                           cluster['cluster_id']*10+i,
                                           waves[idx],
                                           cluster['spike_times'][idx],
                                           cluster['spike_map'][idx],
                                           cluster['rec_key'].copy(),
                                           cluster['fs'].copy(),
                                           cluster['offsets'].copy(),
                                           manipulations=edit_str)
                  new_clusts.append(tmp_clust)
  
              # Plot cluster and ask to choose which to keep
              figs = []
              for i, c in enumerate(new_clusts):
                  _, viol_1ms, viol_2ms = get_ISI_and_violations(c['spike_times'], c['fs'], c['spike_map'])
                  plot_title = ('Index: %i\n1ms violations: %i, 2ms violations: %i\n'
                                'Total Waveforms: %i'
                                % (i, viol_1ms, viol_2ms, len(c['spike_times'])))
                  tmp_fig, _ = dplt.plot_waveforms(c['spike_waveforms'], title=plot_title)
                  figs.append(tmp_fig)
                  tmp_fig.show()
  
              f2 = dplt.plot_waveforms_pca([c['spike_waveforms'] for c in new_clusts])
              figs.append(f2)
              f2.show()
          except:
              # So cluster isn't lost with error
              self._active.insert(target_clust, cluster)
              self._split_starter = None
              self._split_index = None
              raise
  
          if store_split:
              self._split_results = new_clusts
              return new_clusts
          else:
              self._split_starter = None
              self._split_index = None
              selection_list = ['all'] + ['%i' % i for i in range(len(new_clusts))]
              prompt = 'Select split clusters to keep\nCancel to reset.'
              ans = userIO.select_from_list(prompt, selection_list,
                                            multi_select=True, shell=self._shell)
              if ans is None or 'all' in ans:
                  print('Reset to before split')
                  self._active.insert(target_clust, cluster)
              else:
                  keepers = [new_clusts[int(i)] for i in ans]
                  start_idx = len(self._active)
                  self._last_added = list(range(start_idx, start_idx+len(keepers)))
                  self._last_popped = {target_clust: cluster}
                  self._last_action = 'split'
                  self._active.extend(keepers)
  
              return True
  
      def set_split(self, choices):
          if self._split_starter is None:
              raise ValueError('Not split stored.')
  
          if len(choices) == 0:
              self._active.insert(self._split_index, self._split_starter)
          else:
              keepers = [self._split_results[i] for i in choices]
              start_idx = len(self._active)
              self._last_added = list(range(start_idx, start_idx+len(keepers)))
              self._last_popped = {self._split_index: self._split_starter}
              self._last_action = 'split'
              self._active.extend(keepers)
  
          self._split_index = None
          self._split_results = None
          self._split_starter = None
  
      def merge_clusters(self, target_clusters):
          if any([i >= len(self._active) for i in target_clusters]):
              raise ValueError('Target cluster is out of range.')
  
          new_clust = []
          self._last_popped = {}
          self._last_action = 'merge'
          self._last_added = []
          for c in target_clusters:
              self._last_popped[c] = self._active[c]
  
              if len(new_clust) == 0:
                  new_clust = deepcopy(self._active[c])
                  continue
  
              clust = self._active[c]
              sm1 = new_clust['spike_map']
              sm2 = clust['spike_map']
              st1 = new_clust['spike_times']
              st2 = clust['spike_times']
              sw1 = new_clust['spike_waveforms']
              sw2 = clust['spike_waveforms']
  
              spike_map = np.hstack((sm1, sm2))
              spike_times = np.hstack((st1, st2))
              spike_waveforms = np.vstack((sw1, sw2))
  
              # Re-order to spike_map
              idx = np.argsort(spike_map)
              spike_map = spike_map[idx]
              spike_times = spike_times[idx]
              spike_waveforms = spike_waveforms[idx]
  
              # Re-order so spike_times within a reocrding are in order
              times = []
              waves = []
              new_map = []
              for i in np.unique(spike_map):
                  idx = np.where(spike_map == i)[0]
                  st = spike_times[idx]
                  sw = spike_waveforms[idx]
                  sm = spike_map[idx]
                  idx2 = np.argsort(st)
                  st = st[idx2]
                  sw = sw[idx2]
                  sm = sm[idx2]
                  times.append(st)
                  waves.append(sw)
                  new_map.append(sm)
  
              times = np.hstack(times)
              waves = np.vstack(waves)
              spike_map = np.hstack(new_map)
              del new_map, spike_times, spike_waveforms
  
              new_clust['spike_map'] = spike_map
              new_clust['spike_times'] = times
              new_clust['spike_waveforms'] = waves
              new_clust['manipulations'] += '\nMerged with %s.' % clust['Cluster_Name']
              new_clust['Cluster_Name'] += '+' + clust['Cluster_Name'].replace('Cluster_','')
  
          self._active = [self._active[i] for i in range(len(self._active))
                          if i not in target_clusters]
  
          self._last_added = [len(self._active)]
          self._active.append(new_clust)
  
      def discard_clusters(self, target_clusters):
          if isinstance(target_clusters, int):
              target_clusters = [target_clusters]
  
          if len(target_clusters) == 0:
              return
  
          self._last_action = 'discard'
          self._last_popped = {i: self._active[i] for i in target_clusters}
          self._last_added = []
          self._active = [self._active[i] for i in range(len(self._active))
                          if i not in target_clusters]
  
      def plot_clusters_waveforms(self, target_clusters):
          if len(target_clusters) == 0:
              return
  
          for i in target_clusters:
              c = self._active[i]
              isi, v1, v2 = get_ISI_and_violations(c['spike_times'], c['fs'], c['spike_map'])
              title = ('Index : %i\n1ms violations: %0.1f, 2ms violations: %0.1f'
                       '\ntotal waveforms: %i'
                       % (i, v1, v2, len(c['spike_waveforms'])))
              fig, ax = dplt.plot_waveforms(c['spike_waveforms'], title=title,
                                            threshold=self._detection_thresholds[0])
              fig.show()
  
      def split_by_rec(self, target_cluster):
          if isinstance(target_cluster, list) and len(target_cluster) != 1:
              return
          elif isinstance(target_cluster, list):
              target_cluster = target_clsuter[0]
  
          clust = self._active[target_cluster]
          sm = clust['spike_map']
          recs = np.unique(sm)
          if len(sm) == 1:
              return
          else:
              clust = self._active.pop(target_cluster)
              st = clust['spike_times']
              sw = clust['spike_waveforms']
              keepers = []
              for i in recs:
                  idx = np.where(sm == i)[0]
                  new_clust = deepcopy(clust)
                  new_clust['spike_times'] = st[idx]
                  new_clust['spike_waveforms'] = sw[idx, :]
                  new_clust['cluster_id'] = clust['cluster_id']*10 + i
                  new_clust['spike_map'] = sm[idx]
                  new_clust['manipulations'] = '\nSplit by recording'
                  keepers.append(new_clust)
  
          start_idx = len(self._active)
          self._last_added = list(range(start_idx, start_idx+len(keepers)))
          self._last_popped = {target_cluster: clust}
          self._last_action = 'split'
          self._active.extend(keepers)
  
      def plot_cluster_waveforms_by_rec(self, target_cluster):
          if isinstance(target_cluster, list) and len(target_cluster) != 1:
              return
          elif isinstance(target_cluster, list):
              target_cluster = target_cluster[0]
  
          c = self._active[target_cluster]
          sm = c['spike_map']
          for i in np.unique(sm):
              idx = np.where(sm==i)[0]
              waves = c['spike_waveforms'][idx, :]
              isi, v1, v2 = get_ISI_and_violations(c['spike_times'][idx], c['fs'][i])
              title = ('Index : %i, Rec: %i\n1ms violations: %0.1f, 2ms violations: %0.1f'
                       '\ntotal waveforms: %i'
                       % (target_cluster, i, v1, v2, len(waves)))
              fig, ax = dplt.plot_waveforms(waves, title=title)
              fig.show()
  
      def plot_cluster_waveforms_over_time(self, target_cluster, interval):
          if isinstance(target_cluster, list) and len(target_cluster) != 1:
              return
          elif isinstance(target_cluster, list):
              target_cluster = target_cluster[0]
  
          c = self._active[target_cluster]
          spike_times = c.get_spike_time_vector('s')
          start_times = np.arange(spike_times[0], spike_times[-1]+1, interval)
          if len(start_times) > 10:
              userIO.tell_user('This would open more than 10 figures, choose a larger interval')
              return
  
          if len(start_times) == 0:
              return
  
          for i, start_time in enumerate(start_times):
              idx = np.where((spike_times >= start_time) & (spike_times < start_time+interval))[0]
              waves = c['spike_waveforms'][idx,:]
              isi, v1, v2 = get_ISI_and_violations(c['spike_times'][idx], c['fs'][0])
              title = ('Index : %i, Rec: %i\n1ms violations: %0.1f, 2ms violations: %0.1f'
                       '\ntotal waveforms: %i'
                       % (target_cluster, i, v1, v2, len(waves)))
              fig, ax = dplt.plot_waveforms(waves, title=title,
                                            threshold=self._detection_thresholds[0])
              fig.show()
  
      def plot_clusters_pca(self, target_clusters):
          if len(target_clusters) == 0:
              return
  
          waves = [self._active[i]['spike_waveforms'] for i in target_clusters]
          fig = dplt.plot_waveforms_pca(waves, cluster_ids=target_clusters)
          fig.show()
  
      def plot_clusters_umap(self, target_clusters):
          if len(target_clusters) == 0:
              return
  
          waves = [self._active[i]['spike_waveforms'] for i in target_clusters]
          fig = dplt.plot_waveforms_umap(waves, cluster_ids=target_clusters)
          fig.show()
  
      def plot_clusters_wavelets(self, target_clusters):
          if len(target_clusters) == 0:
              return
  
          waves = [self._active[i]['spike_waveforms'] for i in target_clusters]
          fig, ax = dplt.plot_waveforms_wavelet_tranform(waves,
                                                         cluster_ids=target_clusters,
                                                         n_pc=4)
          fig.show()
  
      def plot_clusters_raster(self, target_clusters):
          if len(target_clusters) == 0:
              return
  
          clusters = [self._active[i] for i in target_clusters]
          spike_times = []
          spike_waves = []
          vlines = {}
          for c in clusters:
              # Adjust spike times by offset so recordings are not overlapping
              st = c.get_spike_time_vector(units='s')
              vlines = {i: c['offsets'][i] / c['fs'][i] for i in c['fs'].keys()}
              spike_times.append(st)
              spike_waves.append(c['spike_waveforms'])
  
          fig, ax = dplt.plot_spike_raster(spike_times, spike_waves, target_clusters)
          ax.set_xlabel('Time (s)')
          for x in vlines.values():
              ax.axvline(x, color='black', linewidth=2)
  
          fig.show()
  
      def plot_clusters_ISI(self, target_clusters):
          if len(target_clusters) == 0:
              return
  
          for i in target_clusters:
              cluster = self._active[i]
              isi, v1, v2 = get_ISI_and_violations(cluster['spike_times'],
                                                   cluster['fs'],
                                                   cluster['spike_map'])
              fig, ax = dplt.plot_ISIs(isi, total_spikes=len(cluster['spike_times']))
              title= ax.get_title()
              title = 'Index: %i\n%s' % (i, title)
              ax.set_title(title)
              fig.show()
  
      def plot_clusters_acorr(self, target_clusters):
          if len(target_clusters) == 0 or not all([x < len(self._active) for x in target_clusters]):
              return
  
          for i in target_clusters:
              cluster = self._active[i]
              acf, bin_centers, edges = sas.spike_time_acorr(cluster.get_spike_time_vector(units='ms'))
              fig, ax = dplt.plot_correlogram(acf, bin_centers, edges)
              title = 'Index: %i\nAutocorrelogram' % (i)
              ax.set_title(title)
              fig.show()
  
      def plot_clusters_xcorr(self, target_clusters):
          if len(target_clusters) == 0 or not all([x < len(self._active) for x in target_clusters]):
              return
  
          pairs = it.combinations(target_clusters, 2)
          for x, y in pairs:
              clust1 = self._active[x]
              clust2 = self._active[y]
              xcf, bin_centers, edges = sas.spike_time_xcorr(clust1.get_spike_time_vector(units='ms'),
                                                      clust2.get_spike_time_vector(units='ms'))
              fig, ax = dplt.plot_correlogram(xcf, bin_centers, edges)
              title = 'Cross-correlogram\n%i vs %i' % (x, y)
              ax.set_title(title)
              fig.show()
  
      def get_mean_waveform(self, target_cluster):
          '''Returns mean waveform of target_cluster in active clusters. Also
          returns St. Dev. of waveforms
          '''
          cluster = self._active[target_cluster]
          return cluster.get_mean_waveform()
  
      def get_possible_solutions(self):
          results = self.clustering.results.dropna()
          converged = list(results[results['converged']].index)
          return converged
  
  
  class SpikeCluster(dict):
      def __init__(self, name, electrode, solution, cluster, cluster_id, waves, times,
                   spike_map, rec_key, fs={0: 30000}, offsets={0:0}, manipulations=''):
          # Confirm spike_map, rec_key, fs and offsets are all in sync
          rec_nums = np.unique(spike_map)
          if (not all([x in rec_key.keys() for x in rec_nums]) or
              not all([x in fs.keys() for x in rec_nums]) or
              not all([x in offsets.keys() for x in rec_nums])):
              raise ValueError('rec_key, fs and offsets must have entries for '
                               'each unique element of spike_map')
  
          # Confirm same number of waves, times and map entries
          if waves.shape[0] != len(times) or len(times) != len(spike_map):
              raise ValueError('Must have same number of waves, times and map entries')
  
          super(SpikeCluster, self).__init__(Cluster_Name=name,
                                             electrode_num=electrode,
                                             solution_num=solution,
                                             cluster_num = cluster,
                                             cluster_id=cluster_id,
                                             spike_waveforms=waves,
                                             spike_times=times,
                                             spike_map=spike_map,
                                             rec_key=rec_key,
                                             fs=fs,
                                             offsets=offsets,
                                             manipulations=manipulations)
  
      def delete_spikes(self, idx, msg=None):
          self['spike_waveforms'] = np.delete(self['spike_waveforms'],
                                                 idx, axis=0)
          self['spike_times'] = np.delete(self['spike_times'], idx)
          self['spike_map'] = np.delete(self['spike_map'], idx)
          print('deleted %i spikes.' % len(idx))
          if msg is not None:
              self['manipulations'] += '/n' + msg + '\n-Removed %i spikes' % len(idx)
  
      def get_spike_time_vector(self, units='samples'):
          '''Return vector of all spike times with offsets added if multiple
          recordings are present
  
          Parameters
          ----------
          units : {'samples' (default), 'ms', 's'}, units for spike times returned
  
          Returns
          -------
          numpy.ndarray
          '''
          if units.lower() == 'ms':
              times = np.array([(a + self['offsets'][b]) / (self['fs'][b] / 1000)
                                for a, b in
                                zip(self['spike_times'].astype('float64'), self['spike_map'])])
          elif units.lower() == 's':
              times = np.array([(a + self['offsets'][b]) / self['fs'][b]
                                for a, b in
                                zip(self['spike_times'].astype('float64'), self['spike_map'])])
          elif units.lower() == 'samples':
              times = np.array([(a + self['offsets'][b]) for a, b in
                                zip(self['spike_times'], self['spike_map'])])
          else:
              raise ValueError('units must be either samples or ms')
  
          return times
  
      def __eq__(self, other):
          times1 = self.get_spike_time_vector(units='samples')
          times2 = other.get_spike_time_vector(units='samples')
          times1 = np.sort(times1)
          times2= np.sort(times2)
          return np.array_equal(times1, times2)
  
      def get_mean_waveform(self):
          '''Returns mean waveform of cluster. Also
          returns St. Dev. of waveforms and number of waveforms
  
          Returns
          -------
          np.ndarray, np.ndarray, int
          mean waveform, st. dev of waveforms, number of waveforms 
          '''
          mean_wave = np.mean(self['spike_waveforms'], axis=0)
          std_wave = np.std(self['spike_waveforms'], axis=0)
          n_waves = self['spike_waveforms'].shape[0]
          return mean_wave, std_wave, n_waves
  
      # TODO: Finish this section
      def _dist(self, other=None):
          if other is None:
              other = self
  
          pass
  
      def _add(self, other):
          pass
  
      def _subtract(self, other):
          pass
  
      def _divide(self, N, method='pca'):
          pass
  
      def _is_subcluster(self, other):
          '''less than?'''
          pass
  
      def _is_supercluster(self, other):
          ''' greater than?'''
          pass
  
  