code: import numpy as np
  import datajoint as dj
  from PIL import ImageColor
  from collections import Counter
  
  import matplotlib as mpl
  import matplotlib.pyplot as plt
  import seaborn as sns
  import itertools
  import pandas as pd
  
  from pipeline import experiment, ephys, psth, lab, histology, ccf, psth_foraging
  
  from pipeline.plot.util import (_plot_with_sem, _extract_one_stim_dur,
                                  _plot_stacked_psth_diff, _plot_avg_psth, _jointplot_w_hue)
  from pipeline.plot import unit_psth
  from pipeline.util import (_get_units_hemisphere, _get_trial_event_times,
                             _get_stim_onset_time, _get_clustering_method)
  
  from . import PhotostimError
  
  _plt_xmin = -3
  _plt_xmax = 2
  
  
  def plot_clustering_quality(probe_insertion, clustering_method=None, axs=None):
      probe_insertion = probe_insertion.proj()
  
      if clustering_method is None:
          try:
              clustering_method = _get_clustering_method(probe_insertion)
          except ValueError as e:
              raise ValueError(str(e) + '\nPlease specify one with the kwarg "clustering_method"')
  
      amp, snr, spk_rate, isi_violation = (ephys.Unit * ephys.UnitStat * ephys.ProbeInsertion.InsertionLocation
                                           & probe_insertion & {'clustering_method': clustering_method}).fetch(
          'unit_amp', 'unit_snr', 'avg_firing_rate', 'isi_violation')
  
      metrics = {'amp': amp,
                 'snr': snr,
                 'isi': np.array(isi_violation) * 100,  # to percentage
                 'rate': np.array(spk_rate)}
      label_mapper = {'amp': 'Amplitude',
                      'snr': 'Signal to noise ratio (SNR)',
                      'isi': 'ISI violation (%)',
                      'rate': 'Firing rate (spike/s)'}
  
      fig = None
      if axs is None:
          fig, axs = plt.subplots(2, 3, figsize = (12, 8))
          fig.subplots_adjust(wspace=0.4)
  
      assert axs.size == 6
  
      for (m1, m2), ax in zip(itertools.combinations(list(metrics.keys()), 2), axs.flatten()):
          ax.plot(metrics[m1], metrics[m2], '.k')
          ax.set_xlabel(label_mapper[m1])
          ax.set_ylabel(label_mapper[m2])
  
          # cosmetic
          ax.spines['right'].set_visible(False)
          ax.spines['top'].set_visible(False)
  
      return fig
  
  
  def plot_unit_characteristic(probe_insertion, clustering_method=None, axs=None):
      probe_insertion = probe_insertion.proj()
  
      if clustering_method is None:
          try:
              clustering_method = _get_clustering_method(probe_insertion)
          except ValueError as e:
              raise ValueError(str(e) + '\nPlease specify one with the kwarg "clustering_method"')
  
      if clustering_method in ('kilosort2'):
          q_unit = (ephys.Unit * ephys.ProbeInsertion.InsertionLocation.proj('depth') * ephys.UnitStat
                    * lab.ElectrodeConfig.Electrode.proj() * lab.ProbeType.Electrode.proj('x_coord', 'y_coord')
                    & probe_insertion & {'clustering_method': clustering_method} & 'unit_quality != "all"').proj(
              ..., x='x_coord', y='y_coord')
      else:
          q_unit = (ephys.Unit * ephys.ProbeInsertion.InsertionLocation.proj('depth') * ephys.UnitStat
                    & probe_insertion & {'clustering_method': clustering_method} & 'unit_quality != "all"').proj(
              ..., x='unit_posx', y='unit_posy')
  
      amp, snr, spk_rate, x, y, insertion_depth = q_unit.fetch(
          'unit_amp', 'unit_snr', 'avg_firing_rate', 'x', 'y', 'depth')
  
      metrics = pd.DataFrame(list(zip(*(amp/amp.max(), snr/snr.max(), spk_rate/spk_rate.max(),
                                        x, insertion_depth.astype(float) + y))))
      metrics.columns = ['amp', 'snr', 'rate', 'x', 'y']
  
      # --- prepare for plotting
      shank_count = (ephys.ProbeInsertion & probe_insertion).aggr(lab.ElectrodeConfig.Electrode * lab.ProbeType.Electrode,
                                                                  shank_count='count(distinct shank)').fetch1('shank_count')
      m_scale = get_m_scale(shank_count)
  
      ymin = metrics.y.min() - 100
      ymax = metrics.y.max() + 200
      xmax = 1.3 * metrics.x.max()
      xmin = -1/6*xmax
      cosmetic = {'legend': None,
                  'linewidth': 1.75,
                  'alpha': 0.9,
                  'facecolor': 'none', 'edgecolor': 'k'}
  
      # --- plot
      fig = None
      if axs is None:
          fig, axs = plt.subplots(1, 3, figsize=(10, 8))
          fig.subplots_adjust(wspace=0.6)
  
      assert axs.size == 3
  
      sns.scatterplot(data=metrics, x='x', y='y', s=metrics.amp*m_scale, ax=axs[0], **cosmetic)
      sns.scatterplot(data=metrics, x='x', y='y', s=metrics.snr*m_scale, ax=axs[1], **cosmetic)
      sns.scatterplot(data=metrics, x='x', y='y', s=metrics.rate*m_scale, ax=axs[2], **cosmetic)
  
      # manually draw the legend
      lg_ypos = ymax
      data = pd.DataFrame({'x': [0.1*xmax, 0.4*xmax, 0.75*xmax], 'y': [lg_ypos, lg_ypos, lg_ypos],
                           'size_ratio': np.array([0.2, 0.5, 0.8])})
      for ax, ax_maxval in zip(axs.flatten(), (amp.max(), snr.max(), spk_rate.max())):
          sns.scatterplot(data=data, x='x', y='y', s=data.size_ratio*m_scale, ax=ax, **dict(cosmetic, facecolor='k'))
          for _, r in data.iterrows():
              ax.text(r['x']-4, r['y']+70, (r['size_ratio']*ax_maxval).astype(int))
  
      # cosmetic
      for title, ax in zip(('Amplitude', 'SNR', 'Firing rate'), axs.flatten()):
          ax.spines['right'].set_visible(False)
          ax.spines['top'].set_visible(False)
          ax.set_title(title)
          ax.set_xlim((xmin, xmax))
          ax.plot([0.5*xmin, xmax], [lg_ypos-80, lg_ypos-80], '-k')
          ax.set_ylim((ymin, ymax + 150))
  
      return fig
  
  
  def plot_unit_selectivity(probe_insertion, clustering_method=None, axs=None):
      probe_insertion = probe_insertion.proj()
  
      if clustering_method is None:
          try:
              clustering_method = _get_clustering_method(probe_insertion)
          except ValueError as e:
              raise ValueError(str(e) + '\nPlease specify one with the kwarg "clustering_method"')
  
      if clustering_method in ('kilosort2'):
          q_unit = (psth.PeriodSelectivity * ephys.Unit * ephys.ProbeInsertion.InsertionLocation
                    * lab.ElectrodeConfig.Electrode.proj() * lab.ProbeType.Electrode.proj('x_coord', 'y_coord')
                    * experiment.Period & probe_insertion & {'clustering_method': clustering_method}
                    & 'period_selectivity != "non-selective"').proj(..., x='unit_posx', y='unit_posy').proj(
              ..., x='x_coord', y='y_coord')
      else:
          q_unit = (psth.PeriodSelectivity * ephys.Unit * ephys.ProbeInsertion.InsertionLocation
                    * experiment.Period & probe_insertion & {'clustering_method': clustering_method}
                    & 'period_selectivity != "non-selective"').proj(..., x='unit_posx', y='unit_posy')
  
      attr_names = ['unit', 'period', 'period_selectivity', 'contra_firing_rate',
                    'ipsi_firing_rate', 'x', 'y', 'depth']
      selective_units = q_unit.fetch(*attr_names)
      selective_units = pd.DataFrame(selective_units).T
      selective_units.columns = attr_names
      selective_units.period_selectivity.astype('category')
  
      # --- account for insertion depth (manipulator depth)
      selective_units.y = selective_units.depth.values.astype(float) + selective_units.y
  
      # --- get ipsi vs. contra firing rate difference
      f_rate_diff = np.abs(selective_units.ipsi_firing_rate - selective_units.contra_firing_rate)
      selective_units['f_rate_diff'] = f_rate_diff / f_rate_diff.max()
  
      # --- prepare for plotting
      shank_count = (ephys.ProbeInsertion & probe_insertion).aggr(lab.ElectrodeConfig.Electrode * lab.ProbeType.Electrode,
                                                                  shank_count='count(distinct shank)').fetch1('shank_count')
      m_scale = get_m_scale(shank_count)
  
      cosmetic = {'legend': None,
                  'linewidth': 0.0001}
      ymin = selective_units.y.min() - 100
      ymax = selective_units.y.max() + 100
      xmax = 1.3 * selective_units.x.max()
      xmin = -1/6*xmax
  
      # a bit of hack to get the 'open circle'
      pts = np.linspace(0, np.pi * 2, 24)
      circ = np.c_[np.sin(pts) / 2, -np.cos(pts) / 2]
      vert = np.r_[circ, circ[::-1] * .7]
  
      open_circle = mpl.path.Path(vert)
  
      # --- plot
      fig = None
      if axs is None:
          fig, axs = plt.subplots(1, 3, figsize=(10, 8))
          fig.subplots_adjust(wspace=0.6)
  
      assert axs.size == 3
  
      for (title, df), ax in zip(((p, selective_units[selective_units.period == p])
                                  for p in ('sample', 'delay', 'response')), axs):
          sns.scatterplot(data=df, x='x', y='y',
                          s=df.f_rate_diff.values.astype(float)*m_scale,
                          hue='period_selectivity', marker=open_circle,
                          palette={'contra-selective': 'b', 'ipsi-selective': 'r'},
                          ax=ax, **cosmetic)
          contra_p = (df.period_selectivity == 'contra-selective').sum() / len(df) * 100
          # cosmetic
          ax.spines['right'].set_visible(False)
          ax.spines['top'].set_visible(False)
          ax.set_title(f'{title}\n% contra: {contra_p:.2f}\n% ipsi: {100-contra_p:.2f}')
          ax.set_xlim((xmin, xmax))
          ax.set_xlabel('x')
          ax.set_ylabel('y')
          ax.set_ylim((ymin, ymax))
  
      return fig
  
  
  def plot_unit_bilateral_photostim_effect(probe_insertion, clustering_method=None, axs=None):
      probe_insertion = probe_insertion.proj()
  
      if not (psth.TrialCondition().get_trials('all_noearlylick_both_alm_stim') & probe_insertion):
          raise PhotostimError('No Bilateral ALM Photo-stimulation present')
  
      if clustering_method is None:
          try:
              clustering_method = _get_clustering_method(probe_insertion)
          except ValueError as e:
              raise ValueError(str(e) + '\nPlease specify one with the kwarg "clustering_method"')
  
      dv_loc = (ephys.ProbeInsertion.InsertionLocation & probe_insertion).fetch1('depth')
  
      no_stim_cond = (psth.TrialCondition
                      & {'trial_condition_name':
                         'all_noearlylick_nostim'}).fetch1('KEY')
  
      bi_stim_cond = (psth.TrialCondition
                      & {'trial_condition_name':
                         'all_noearlylick_both_alm_stim'}).fetch1('KEY')
  
      units = ephys.Unit & probe_insertion & {'clustering_method': clustering_method} & 'unit_quality != "all"'
  
      metrics = pd.DataFrame(columns=['unit', 'x', 'y', 'frate_change'])
  
      # get photostim onset and duration
      stim_durs = np.unique((experiment.Photostim & experiment.PhotostimEvent
                             * psth.TrialCondition().get_trials('all_noearlylick_both_alm_stim')
                             & probe_insertion).fetch('duration'))
      stim_dur = _extract_one_stim_dur(stim_durs)
      stim_time = _get_stim_onset_time(units, 'all_noearlylick_both_alm_stim')
  
      # XXX: could be done with 1x fetch+join
      for u_idx, unit in enumerate(units.fetch('KEY', order_by='unit')):
          if clustering_method in ('kilosort2'):
              x, y = (ephys.Unit * lab.ElectrodeConfig.Electrode.proj()
                      * lab.ProbeType.Electrode.proj('x_coord', 'y_coord') & unit).fetch1('x_coord', 'y_coord')
          else:
              x, y = (ephys.Unit & unit).fetch1('unit_posx', 'unit_posy')
  
          # obtain unit psth per trial, for all nostim and bistim trials
          nostim_trials = ephys.Unit.TrialSpikes & unit & psth.TrialCondition.get_trials(no_stim_cond['trial_condition_name'])
          bistim_trials = ephys.Unit.TrialSpikes & unit & psth.TrialCondition.get_trials(bi_stim_cond['trial_condition_name'])
  
          nostim_psths, nostim_edge = psth.compute_unit_psth(unit, nostim_trials.fetch('KEY'), per_trial=True)
          bistim_psths, bistim_edge = psth.compute_unit_psth(unit, bistim_trials.fetch('KEY'), per_trial=True)
  
          # compute the firing rate difference between contra vs. ipsi within the stimulation time window
          ctrl_frate = np.array([nostim_psth[np.logical_and(nostim_edge >= stim_time,
                                                            nostim_edge <= stim_time + stim_dur)].mean()
                                 for nostim_psth in nostim_psths])
          stim_frate = np.array([bistim_psth[np.logical_and(bistim_edge >= stim_time,
                                                            bistim_edge <= stim_time + stim_dur)].mean()
                                 for bistim_psth in bistim_psths])
  
          frate_change = (stim_frate.mean() - ctrl_frate.mean()) / ctrl_frate.mean()
          frate_change = abs(frate_change) if frate_change < 0 else 0.0001
  
          metrics.loc[u_idx] = (int(unit['unit']), x, float(dv_loc) + y, frate_change)
  
      metrics.frate_change = metrics.frate_change / metrics.frate_change.max()
  
      # --- prepare for plotting
      shank_count = (ephys.ProbeInsertion & probe_insertion).aggr(lab.ElectrodeConfig.Electrode * lab.ProbeType.Electrode,
                                                                  shank_count='count(distinct shank)').fetch1('shank_count')
      m_scale = get_m_scale(shank_count)
  
      fig = None
      if axs is None:
          fig, axs = plt.subplots(1, 1, figsize=(4, 8))
  
      xmax = 1.3 * metrics.x.max()
      xmin = -1/6*xmax
  
      cosmetic = {'legend': None,
                  'linewidth': 1.75,
                  'alpha': 0.9,
                  'facecolor': 'none', 'edgecolor': 'k'}
  
      sns.scatterplot(data=metrics, x='x', y='y', s=metrics.frate_change*m_scale,
                      ax=axs, **cosmetic)
  
      axs.spines['right'].set_visible(False)
      axs.spines['top'].set_visible(False)
      axs.set_title('% change')
      axs.set_xlim((xmin, xmax))
  
      return fig
  
  
  def plot_pseudocoronal_slice(probe_insertion, shank_no=1):
      # ---- Electrode sites ----
      annotated_electrodes = (lab.ElectrodeConfig.Electrode * lab.ProbeType.Electrode
                              * ephys.ProbeInsertion
                              * histology.ElectrodeCCFPosition.ElectrodePosition
                              & probe_insertion & {'shank': shank_no})
  
      electrode_coords = np.array(list(zip(*annotated_electrodes.fetch(
          'ccf_z', 'ccf_y', 'ccf_x', order_by='ccf_y'))))  # (AP, DV, ML)
      probe_track_coords = np.array(list(zip(*(histology.LabeledProbeTrack.Point
                                               & probe_insertion & {'shank': shank_no}).fetch(
          'ccf_z', 'ccf_y', 'ccf_x', order_by='ccf_y'))))
  
      voxel_res = ccf.CCFLabel.CCF_R3_20UM_RESOLUTION
      lr_max, dv_max, _ = ccf.get_ccf_xyz_max()
  
      pseudocoronal_points, shank_ccfs = histology.retrieve_pseudocoronal_slice(probe_insertion, shank_no)
  
      dv_pts, lr_pts, ap_pts, color_codes = pseudocoronal_points.T
      dv_pts = dv_pts.astype(int)
      lr_pts = lr_pts.astype(int)
      color_codes = color_codes.astype(str)
  
      # ---- paint annotation color code ----
      coronal_slice = np.full((dv_max + 1, lr_max + 1, 3), np.nan)
      for color in set(color_codes):
          matched_ind = np.where(color_codes == color)[0]
          dv_ind = dv_pts[matched_ind]  # rows
          lr_ind = lr_pts[matched_ind]  # cols
          try:
              c_rgb = ImageColor.getcolor("#" + color, "RGB")
          except ValueError as e:
              print(str(e))
              continue
          coronal_slice[dv_ind, lr_ind, :] = np.full((len(matched_ind), 3), c_rgb)
  
      # ---- paint the interpolated track of this probe/shank in gray ----
      in_probe_range = np.logical_and(shank_ccfs[:, 1] >= probe_track_coords[:, 1].min(),
                                      shank_ccfs[:, 1] <= probe_track_coords[:, 1].max())
      in_electrode_range = np.logical_and(shank_ccfs[:, 1] >= electrode_coords[:, 1].min(),
                                          shank_ccfs[:, 1] <= electrode_coords[:, 1].max())
  
      tracks_coords = shank_ccfs[np.logical_and(in_probe_range, ~in_electrode_range), :]
      coronal_slice[tracks_coords[:, 1], tracks_coords[:, 0], :] = np.full(
          (tracks_coords.shape[0], 3), ImageColor.getcolor("#FFFFFF", "RGB"))
  
      # ---- paint electrode sites on this probe/shank in black ----
      coronal_slice[electrode_coords[:, 1], electrode_coords[:, 2], :] = np.full(
          (electrode_coords.shape[0], 3), ImageColor.getcolor("#080808", "RGB"))
  
      # ---- downsample the 2D slice to the voxel resolution ----
      coronal_slice = coronal_slice[::voxel_res, ::voxel_res, :]
  
      # paint outside region white
      nan_r, nan_c = np.where(np.nansum(coronal_slice, axis=2) == 0)
      coronal_slice[nan_r, nan_c, :] = np.full((len(nan_r), 3), ImageColor.getcolor("#FFFFFF", "RGB"))
  
      # ---- plot ----
      fig, ax = plt.subplots(1, 1)
      ax.imshow(coronal_slice.astype(np.uint8), extent=[0, lr_max, dv_max, 0])
  
      ax.invert_xaxis()
      ax.set_xticks([])
      ax.set_yticks([])
      ax.spines['right'].set_visible(False)
      ax.spines['top'].set_visible(False)
      ax.spines['bottom'].set_visible(False)
      ax.spines['left'].set_visible(False)
  
      return fig
  
  
  def plot_driftmap(probe_insertion, clustering_method=None, shank_no=1):
      probe_insertion = probe_insertion.proj()
  
      assert histology.InterpolatedShankTrack & probe_insertion
  
      if clustering_method is None:
          try:
              clustering_method = _get_clustering_method(probe_insertion)
          except ValueError as e:
              raise ValueError(str(e) + '\nPlease specify one with the kwarg "clustering_method"')
  
      units = (ephys.Unit * lab.ElectrodeConfig.Electrode
               & probe_insertion & {'clustering_method': clustering_method}
               & 'unit_quality != "all"')
      units = (units.proj('spike_times', 'spike_depths', 'unit_posy')
               * ephys.ProbeInsertion.proj()
               * lab.ProbeType.Electrode.proj('shank') & {'shank': shank_no})
  
      # ---- ccf region ----
      annotated_electrodes = (lab.ElectrodeConfig.Electrode * lab.ProbeType.Electrode
                              * ephys.ProbeInsertion
                              * histology.ElectrodeCCFPosition.ElectrodePosition
                              * ccf.CCFAnnotation * ccf.CCFBrainRegion.proj(..., annotation='region_name')
                              & probe_insertion & {'shank': shank_no})
      pos_y, ccf_y, color_code = annotated_electrodes.fetch(
          'y_coord', 'ccf_y', 'color_code', order_by='y_coord DESC')
  
      # CCF position of most ventral recording site
      last_electrode_site = np.array((histology.InterpolatedShankTrack.DeepestElectrodePoint
                                      & probe_insertion & {'shank': shank_no}).fetch1(
          'ccf_x', 'ccf_y', 'ccf_z'))
      # CCF position of the brain surface where this shank crosses
      brain_surface_site = np.array((histology.InterpolatedShankTrack.BrainSurfacePoint
                                     & probe_insertion & {'shank': shank_no}).fetch1(
          'ccf_x', 'ccf_y', 'ccf_z'))
  
      # CCF position of most ventral recording site, with respect to the brain surface
      y_ref = -np.linalg.norm(last_electrode_site - brain_surface_site)
  
      # ---- spikes ----brain_surface_site
      spike_times, spike_depths = units.fetch('spike_times', 'spike_depths', order_by='unit')
  
      spike_times = np.hstack(spike_times)
      spike_depths = np.hstack(spike_depths)
  
      # histogram
      # time_res = 10    # time resolution: 1sec
      # depth_res = 10  # depth resolution: 10um
      #
      # spike_bins = np.arange(0, spike_times.max() + time_res, time_res)
      # depth_bins = np.arange(spike_depths.min() - depth_res, spike_depths.max() + depth_res, depth_res)
  
      # time-depth 2D histogram
      time_bin_count = 1000
      depth_bin_count = 200
  
      spike_bins = np.linspace(0, spike_times.max(), time_bin_count)
      depth_bins = np.linspace(0, np.nanmax(spike_depths), depth_bin_count)
  
      spk_count, spk_edges, depth_edges = np.histogram2d(spike_times, spike_depths, bins=[spike_bins, depth_bins])
      spk_rates = spk_count / np.mean(np.diff(spike_bins))
      spk_edges = spk_edges[:-1]
      depth_edges = depth_edges[:-1]
  
      # region colorcode, by depths
      binned_hexcodes = []
  
      y_spacing = np.abs(np.nanmedian(np.where(np.diff(pos_y)==0, np.nan, np.diff(pos_y))))
      anno_depth_bins = np.arange(0, depth_bins[-1], y_spacing)
      for s, e in zip(anno_depth_bins[:-1], anno_depth_bins[1:]):
          hexcodes = color_code[np.logical_and(pos_y > s, pos_y <= e)]
          if len(hexcodes):
              binned_hexcodes.append(Counter(hexcodes).most_common()[0][0])
          else:
              binned_hexcodes.append('FFFFFF')
  
      region_rgba = np.array([list(ImageColor.getcolor("#" + chex, "RGBA")) for chex in binned_hexcodes])
      region_rgba = np.repeat(region_rgba[:, np.newaxis, :], 10, axis=1)
  
      # canvas setup
      fig = plt.figure(figsize=(16, 8))
      grid = plt.GridSpec(12, 12)
  
      ax_main = plt.subplot(grid[1:, 0:9])
      ax_cbar = plt.subplot(grid[0, 0:9])
      ax_spkcount = plt.subplot(grid[1:, 9:11])
      ax_anno = plt.subplot(grid[1:, 11:])
  
      # -- plot main --
      im = ax_main.imshow(spk_rates.T, aspect='auto', cmap='gray_r',
                          extent=[spike_bins[0], spike_bins[-1], depth_bins[-1], depth_bins[0]])
      # cosmetic
      ax_main.invert_yaxis()
      ax_main.set_xlabel('Time (sec)')
      ax_main.set_ylabel('Distance from tip sites (um)')
      ax_main.set_ylim(depth_edges[0], depth_edges[-1])
      ax_main.spines['right'].set_visible(False)
      ax_main.spines['top'].set_visible(False)
  
      cb = fig.colorbar(im, cax=ax_cbar, orientation='horizontal')
      cb.outline.set_visible(False)
      cb.ax.xaxis.tick_top()
      cb.set_label('Firing rate (Hz)')
      cb.ax.xaxis.set_label_position('top')
  
      # -- plot spikecount --
      ax_spkcount.plot(spk_count.sum(axis=0) / 10e3, depth_edges, 'k')
      ax_spkcount.set_xlabel('Spike count (x$10^3$)')
      ax_spkcount.set_yticks([])
      ax_spkcount.set_ylim(depth_edges[0], depth_edges[-1])
  
      ax_spkcount.spines['right'].set_visible(False)
      ax_spkcount.spines['top'].set_visible(False)
      ax_spkcount.spines['bottom'].set_visible(False)
      ax_spkcount.spines['left'].set_visible(False)
  
      # -- plot colored region annotation
      ax_anno.imshow(region_rgba, aspect='auto',
                     extent=[0, 10, (anno_depth_bins[-1] + y_ref) / 1000, (anno_depth_bins[0] + y_ref) / 1000])
  
      ax_anno.invert_yaxis()
  
      ax_anno.spines['right'].set_visible(False)
      ax_anno.spines['top'].set_visible(False)
      ax_anno.spines['bottom'].set_visible(False)
      ax_anno.spines['left'].set_visible(False)
  
      ax_anno.set_xticks([])
      ax_anno.yaxis.tick_right()
      ax_anno.set_ylabel('Depth in the brain (mm)')
      ax_anno.yaxis.set_label_position('right')
  
      return fig
  
  
  def plot_stacked_contra_ipsi_psth(units, axs=None):
      units = units.proj()
  
      # get event start times: sample, delay, response
      period_names, period_starts = _get_trial_event_times(['sample', 'delay', 'go'], units, 'good_noearlylick_hit')
  
      hemi = _get_units_hemisphere(units)
  
      conds_i = (psth.TrialCondition
                 & {'trial_condition_name':
                    'good_noearlylick_left_hit' if hemi == 'left' else 'good_noearlylick_right_hit'}).fetch1('KEY')
  
      conds_c = (psth.TrialCondition
                 & {'trial_condition_name':
                    'good_noearlylick_right_hit' if hemi == 'left' else 'good_noearlylick_left_hit'}).fetch1('KEY')
  
      sel_i = (ephys.Unit * psth.UnitSelectivity
               & 'unit_selectivity = "ipsi-selective"' & units)
  
      sel_c = (ephys.Unit * psth.UnitSelectivity
               & 'unit_selectivity = "contra-selective"' & units)
  
      # ipsi selective ipsi trials
      psth_is_it = (psth.UnitPsth * sel_i.proj('unit_posy') & conds_i).fetch(order_by='unit_posy desc')
  
      # ipsi selective contra trials
      psth_is_ct = (psth.UnitPsth * sel_i.proj('unit_posy') & conds_c).fetch(order_by='unit_posy desc')
  
      # contra selective contra trials
      psth_cs_ct = (psth.UnitPsth * sel_c.proj('unit_posy') & conds_c).fetch(order_by='unit_posy desc')
  
      # contra selective ipsi trials
      psth_cs_it = (psth.UnitPsth * sel_c.proj('unit_posy') & conds_i).fetch(order_by='unit_posy desc')
  
      fig = None
      if axs is None:
          fig, axs = plt.subplots(1, 2, figsize=(20, 20))
      assert axs.size == 2
  
      _plot_stacked_psth_diff(psth_cs_ct, psth_cs_it, ax=axs[0], vlines=period_starts, flip=True)
  
      axs[0].set_title('Contra-selective Units')
      axs[0].set_ylabel('Unit (by depth)')
      axs[0].set_xlabel('Time to go (s)')
  
      _plot_stacked_psth_diff(psth_is_it, psth_is_ct, ax=axs[1], vlines=period_starts)
  
      axs[1].set_title('Ipsi-selective Units')
      axs[1].set_ylabel('Unit (by depth)')
      axs[1].set_xlabel('Time to go (s)')
  
      return fig
  
  
  def plot_avg_contra_ipsi_psth(units, axs=None):
      units = units.proj()
  
      # get event start times: sample, delay, response
      period_names, period_starts = _get_trial_event_times(['sample', 'delay', 'go'], units, 'good_noearlylick_hit')
  
      hemi = _get_units_hemisphere(units)
  
      good_unit = ephys.Unit & 'unit_quality != "all"'
  
      conds_i = (psth.TrialCondition
                 & {'trial_condition_name':
                    'good_noearlylick_left_hit' if hemi == 'left' else 'good_noearlylick_right_hit'}).fetch('KEY')
  
      conds_c = (psth.TrialCondition
                 & {'trial_condition_name':
                    'good_noearlylick_right_hit' if hemi == 'left' else 'good_noearlylick_left_hit'}).fetch('KEY')
  
      sel_i = (ephys.Unit * psth.UnitSelectivity
               & 'unit_selectivity = "ipsi-selective"' & units)
  
      sel_c = (ephys.Unit * psth.UnitSelectivity
               & 'unit_selectivity = "contra-selective"' & units)
  
      psth_is_it = (((psth.UnitPsth & conds_i)
                     * ephys.Unit.proj('unit_posy'))
                    & good_unit.proj() & sel_i.proj()).fetch(
                        'unit_psth', order_by='unit_posy desc')
  
      psth_is_ct = (((psth.UnitPsth & conds_c)
                     * ephys.Unit.proj('unit_posy'))
                    & good_unit.proj() & sel_i.proj()).fetch(
                        'unit_psth', order_by='unit_posy desc')
  
      psth_cs_ct = (((psth.UnitPsth & conds_c)
                     * ephys.Unit.proj('unit_posy'))
                    & good_unit.proj() & sel_c.proj()).fetch(
                        'unit_psth', order_by='unit_posy desc')
  
      psth_cs_it = (((psth.UnitPsth & conds_i)
                     * ephys.Unit.proj('unit_posy'))
                    & good_unit.proj() & sel_c.proj()).fetch(
                        'unit_psth', order_by='unit_posy desc')
  
      fig = None
      if axs is None:
          fig, axs = plt.subplots(1, 2, figsize=(16, 6))
      assert axs.size == 2
  
      _plot_avg_psth(psth_cs_it, psth_cs_ct, period_starts, axs[0],
                     'Contra-selective')
      _plot_avg_psth(psth_is_it, psth_is_ct, period_starts, axs[1],
                     'Ipsi-selective')
  
      ymax = max([ax.get_ylim()[1] for ax in axs])
      for ax in axs:
          ax.set_ylim((0, ymax))
  
      return fig
  
  
  def plot_psth_photostim_effect(units, condition_name_kw=['both_alm'], axs=None):
      """
      For the specified `units`, plot PSTH comparison between stim vs. no-stim with left/right trial instruction
      The stim location (or other appropriate search keywords) can be specified in `condition_name_kw` (default: both ALM)
      """
      units = units.proj()
  
      fig = None
      if axs is None:
          fig, axs = plt.subplots(1, 2, figsize=(16, 6))
      assert axs.size == 2
  
      hemi = _get_units_hemisphere(units)
  
      # no photostim:
      psth_n_l = psth.TrialCondition.get_cond_name_from_keywords(['_nostim', '_left'])[0]
      psth_n_r = psth.TrialCondition.get_cond_name_from_keywords(['_nostim', '_right'])[0]
  
      psth_n_l = (psth.UnitPsth * psth.TrialCondition & units
                  & {'trial_condition_name': psth_n_l} & 'unit_psth is not NULL').fetch('unit_psth')
      psth_n_r = (psth.UnitPsth * psth.TrialCondition & units
                  & {'trial_condition_name': psth_n_r} & 'unit_psth is not NULL').fetch('unit_psth')
  
      # with photostim
      psth_s_l = psth.TrialCondition.get_cond_name_from_keywords(condition_name_kw + ['_stim_left'])[0]
      psth_s_r = psth.TrialCondition.get_cond_name_from_keywords(condition_name_kw + ['_stim_right'])[0]
  
      psth_s_l = (psth.UnitPsth * psth.TrialCondition & units
                  & {'trial_condition_name': psth_s_l} & 'unit_psth is not NULL').fetch('unit_psth')
      psth_s_r = (psth.UnitPsth * psth.TrialCondition & units
                  & {'trial_condition_name': psth_s_r} & 'unit_psth is not NULL').fetch('unit_psth')
  
      # get event start times: sample, delay, response
      period_names, period_starts = _get_trial_event_times(['sample', 'delay', 'go'], units, 'good_noearlylick_hit')
  
      # get photostim onset and duration
      stim_trial_cond_name = psth.TrialCondition.get_cond_name_from_keywords(condition_name_kw + ['_stim'])[0]
      stim_durs = np.unique((experiment.Photostim & experiment.PhotostimEvent
                             * psth.TrialCondition().get_trials(stim_trial_cond_name)
                             & units).fetch('duration'))
      stim_dur = _extract_one_stim_dur(stim_durs)
      stim_time = _get_stim_onset_time(units, stim_trial_cond_name)
  
      if hemi == 'left':
          psth_s_i = psth_s_l
          psth_n_i = psth_n_l
          psth_s_c = psth_s_r
          psth_n_c = psth_n_r
      else:
          psth_s_i = psth_s_r
          psth_n_i = psth_n_r
          psth_s_c = psth_s_l
          psth_n_c = psth_n_l
  
      _plot_avg_psth(psth_n_i, psth_n_c, period_starts, axs[0],
                     'Control')
      _plot_avg_psth(psth_s_i, psth_s_c, period_starts, axs[1],
                     'Photostim')
  
      # cosmetic
      ymax = max([ax.get_ylim()[1] for ax in axs])
      for ax in axs:
          ax.set_ylim((0, ymax))
          ax.set_xlim([_plt_xmin, _plt_xmax])
  
      # add shaded bar for photostim
      axs[1].axvspan(stim_time, stim_time + stim_dur, alpha=0.3, color='royalblue')
  
      return fig
  
  
  def plot_coding_direction(units, time_period=None, label=None, axs=None):
      # get event start times: sample, delay, response
      period_names, period_starts = _get_trial_event_times(['sample', 'delay', 'go'], units, 'good_noearlylick_hit')
  
      _, proj_contra_trial, proj_ipsi_trial, time_stamps, _ = psth.compute_CD_projected_psth(
          units.fetch('KEY'), time_period=time_period)
  
      fig = None
      if axs is None:
          fig, axs = plt.subplots(1, 1, figsize=(8, 6))
  
      # plot
      _plot_with_sem(proj_contra_trial, time_stamps, ax=axs, c='b')
      _plot_with_sem(proj_ipsi_trial, time_stamps, ax=axs, c='r')
  
      for x in period_starts:
          axs.axvline(x=x, linestyle = '--', color = 'k')
      # cosmetic
      axs.spines['right'].set_visible(False)
      axs.spines['top'].set_visible(False)
      axs.set_ylabel('CD projection (a.u.)')
      axs.set_xlabel('Time (s)')
      if label:
          axs.set_title(label)
  
      return fig
  
  
  def plot_paired_coding_direction(unit_g1, unit_g2, labels=None, time_period=None):
      """
      Plot trial-to-trial CD-endpoint correlation between CD-projected trial-psth from two unit-groups (e.g. two brain regions)
      Note: coding direction is calculated on selective units, contra vs. ipsi, within the specified time_period
      """
      _, proj_contra_trial_g1, proj_ipsi_trial_g1, time_stamps, unit_g1_hemi = psth.compute_CD_projected_psth(
          unit_g1.fetch('KEY'), time_period=time_period)
      _, proj_contra_trial_g2, proj_ipsi_trial_g2, time_stamps, unit_g2_hemi = psth.compute_CD_projected_psth(
          unit_g2.fetch('KEY'), time_period=time_period)
  
      # get event start times: sample, delay, response
      period_names, period_starts = _get_trial_event_times(['sample', 'delay', 'go'], unit_g1, 'good_noearlylick_hit')
  
      if labels:
          assert len(labels) == 2
      else:
          labels = ('unit group 1', 'unit group 2')
  
      # plot projected trial-psth
      fig, axs = plt.subplots(1, 2, figsize=(16, 6))
  
      _plot_with_sem(proj_contra_trial_g1, time_stamps, ax=axs[0], c='b')
      _plot_with_sem(proj_ipsi_trial_g1, time_stamps, ax=axs[0], c='r')
      _plot_with_sem(proj_contra_trial_g2, time_stamps, ax=axs[1], c='b')
      _plot_with_sem(proj_ipsi_trial_g2, time_stamps, ax=axs[1], c='r')
  
      # cosmetic
      for ax, label in zip(axs, labels):
          for x in period_starts:
              ax.axvline(x=x, linestyle = '--', color = 'k')
          ax.spines['right'].set_visible(False)
          ax.spines['top'].set_visible(False)
          ax.set_ylabel('CD projection (a.u.)')
          ax.set_xlabel('Time (s)')
          ax.set_title(label)
  
      # plot trial CD-endpoint correlation - if 2 unit-groups are from 2 hemispheres,
      #   then contra-ipsi definition is based on the first group
      p_start, p_end = time_period
      contra_cdend_1 = proj_contra_trial_g1[:, np.logical_and(time_stamps >= p_start, time_stamps < p_end)].mean(axis=1)
      ipsi_cdend_1 = proj_ipsi_trial_g1[:, np.logical_and(time_stamps >= p_start, time_stamps < p_end)].mean(axis=1)
      if unit_g1_hemi == unit_g1_hemi:
          contra_cdend_2 = proj_contra_trial_g2[:, np.logical_and(time_stamps >= p_start, time_stamps < p_end)].mean(axis=1)
          ipsi_cdend_2 = proj_ipsi_trial_g2[:, np.logical_and(time_stamps >= p_start, time_stamps < p_end)].mean(axis=1)
      else:
          contra_cdend_2 = proj_ipsi_trial_g2[:, np.logical_and(time_stamps >= p_start, time_stamps < p_end)].mean(axis=1)
          ipsi_cdend_2 = proj_contra_trial_g2[:, np.logical_and(time_stamps >= p_start, time_stamps < p_end)].mean(axis=1)
  
      c_df = pd.DataFrame([contra_cdend_1, contra_cdend_2]).T
      c_df.columns = labels
      c_df['trial-type'] = 'contra'
      i_df = pd.DataFrame([ipsi_cdend_1, ipsi_cdend_2]).T
      i_df.columns = labels
      i_df['trial-type'] = 'ipsi'
      df = c_df.append(i_df)
  
      jplot = _jointplot_w_hue(data=df, x=labels[0], y=labels[1], hue= 'trial-type', colormap=['b', 'r'],
                               figsize=(8, 6), fig=None, scatter_kws=None)
      jplot['fig'].show()
  
      return fig
  
  
  # ========== Foraging task ==========
  def plot_unit_period_fit(linear_model='Q_rel + Q_tot + rpe'):
  #%%
      # linear_model='Q_c + Q_i + rpe'
      q_unit = ((ephys.Unit * ephys.ClusterMetric * ephys.UnitStat * ephys.MAPClusterMetric.DriftMetric)
                & 'presence_ratio > 0.95'
                & 'amplitude_cutoff < 0.1'
                & 'isi_violation < 0.5' 
                & 'unit_amp > 70'
                # & 'drift_metric < 0.1'
                )
  
      q_hist = (q_unit * histology.ElectrodeCCFPosition.ElectrodePosition) * ccf.CCFAnnotation
      q_unit_n = dj.U('annotation').aggr(q_hist, area_num_units='count(*)')
      q_hist *= q_unit_n
  
      lvs = (psth_foraging.LinearModel.X & {'multi_linear_model': linear_model}).fetch('var_name')
      q_all = ((psth_foraging.UnitPeriodLinearFit
                * psth_foraging.UnitPeriodLinearFit.Param
                * q_hist)
                & {'multi_linear_model': linear_model})   
  
  #%%
      # -- Heatmap ---
      # for lv in zip(lvs):
      #     fig, ax = plt.subplots(1, 1, figsize=(13, 15))
  
      #     df = pd.DataFrame((q_all & {'var_name': lv}).proj('beta', 'p', 't', 'area_num_units').fetch())
      #     df = df.pivot(index=['subject_id', 'session', 'insertion_number',
      #                          'clustering_method', 'unit', 'annotation', 'area_num_units'], 
      #                   columns='period', values='t')
      #     df.sort_values(by=['area_num_units', 'iti_all'], ascending=False , inplace=True)
      #     df = df.reset_index().drop(['subject_id', 'session', 'insertion_number', 
      #                                'clustering_method', 'unit', 'area_num_units'], axis=1)
      #     df = df.set_index('annotation')
      #     sns.heatmap(df, ax=ax, cmap='coolwarm')
      #     ax.set_position([0.5, 0.1, 0.2, 0.8])
  
  #%%
      # -- t distribution --
      epochs = ['delay', 'go_to_end', 'iti_first_2', 'iti_last_2']
      fig, axs = plt.subplots(len(lvs), len(epochs), figsize=(4*len(epochs), 4*len(lvs)))
      areas = q_unit_n.fetch(order_by='area_num_units desc', format='frame')
  
      # Areas that have most number of neurons
      areas = list(areas.index[:10])
  
      for i, lv in enumerate(lvs):
          for j, ep in enumerate(epochs):
              ax = axs[i, j]
              ax.axhline(y=0.95, color='k', linestyle=':')
              ax.axvline(x=1.96, color='k', linestyle=':')
  
              for area in areas:
                  this_ts = (q_all & {'var_name': lv, 'period': ep, 'annotation': area}).fetch('t')
                  values, bin = np.histogram(np.abs(this_ts), 100)
                  ax.plot(bin[:-1], np.cumsum(values)/len(this_ts), label=f'{area}, n = {len(this_ts)}')
  
              ax.set(xlim=(0, 10))
              ax.label_outer()
  
              if i == 0:
                  ax.set_title(ep)
              if j == 0:
                  ax.set_ylabel(lv)
              if i == len(lvs) - 1 and j == 0:
                  ax.set_xlabel('|t value|')
      
      ax.legend(bbox_to_anchor=(-1,3), loc='upper left')
  
  #%%
      # -- ipsi and contra action value weights --
      fig, axs = plt.subplots(1,2, figsize=(8,4))
      if linear_model == 'Q_c + Q_i + rpe':
          lvs = ['ipsi_action_value', 'contra_action_value'] 
      elif linear_model == 'Q_l + Q_r + rpe':
          lvs = ['left_action_value', 'right_action_value']
      else:
          lvs = ['relative_action_value_ic', 'total_action_value']
  
      for j, ep in enumerate(['go_to_end', 'iti_all']):
          ax = axs[j]
  
          for area in areas:
              # if not 'thalamus' in area:
              #     continue
  
              ax.axhline(y=-2, color='k', ls='--', lw=.5)
              ax.axhline(y=2, color='k', ls='--', lw=.5)
              ax.axvline(x=-2, color='k', ls='--', lw=.5)
              ax.axvline(x=2, color='k', ls='--', lw=.5)
  
              df = pd.DataFrame((q_all
                                 & {'annotation': area}
                                 & {'period': ep}).proj('beta', 'p', 'area_num_units', 't').fetch())
  
              betas = df.pivot(index=['subject_id', 'session', 'insertion_number',
                                  'clustering_method', 'unit', 'annotation', 'area_num_units'], 
                                  columns='var_name', values='t')
              ps = df.pivot(index=['subject_id', 'session', 'insertion_number',
                                  'clustering_method', 'unit', 'annotation', 'area_num_units'], 
                                  columns='var_name', values='p')
              sizes = 2 + 2 * np.sum(ps.values < 0.05, axis=1)
              ax.scatter(x=betas[lvs[0]], y=betas[lvs[1]], s=sizes)
              ax.set_xlim([-20, 20])
              ax.set_ylim([-20, 20])
              ax.set_xlabel(lvs[0])
              ax.set_ylabel(lvs[1])
              ax.set_title(ep)
              ax.label_outer()
  
              # sns.scatterplot(data=betas, x='ipsi_action_value', y='contra_action_value',
              #                 hue='annotation', sizes=sizes, legend=False)
  
  
  #%%
  
  def plot_example_cells(sort_lv = 'relative_action_value_ic', 
                         sort_ep = 'iti_all',
                         best_n = 10, linear_model='Q_rel + Q_tot + rpe'):
      
  #%%
      q_unit = ((ephys.Unit * ephys.ClusterMetric * ephys.UnitStat * ephys.MAPClusterMetric.DriftMetric)
                & 'presence_ratio > 0.95'
                & 'amplitude_cutoff < 0.1'
                & 'isi_violation < 0.5' 
                & 'unit_amp > 100'
                # & 'drift_metric < 0.1'
                )
  
      q_hist = (q_unit * histology.ElectrodeCCFPosition.ElectrodePosition) * ccf.CCFAnnotation
      q_unit_n = dj.U('annotation').aggr(q_hist, area_num_units='count(*)')
      q_hist *= q_unit_n
  
      lvs = (psth_foraging.LinearModel.X & {'multi_linear_model': linear_model}).fetch('var_name')
      q_all = ((psth_foraging.UnitPeriodLinearFit
                * psth_foraging.UnitPeriodLinearFit.Param
                * q_hist)
                & {'multi_linear_model': linear_model})   
  
      # Best n (absolute value)
      best_models = (q_all & f'var_name = "{sort_lv}"' & f'period = "{sort_ep}"').proj(
          'actual_behavior_model', abs_t='abs(t)').fetch(order_by='abs_t desc', limit=best_n, format='frame')
  
      for unit_key in best_models.reset_index().to_dict('records'):
          unit_psth.plot_unit_psth_choice_outcome(unit_key)
          unit_psth.plot_unit_psth_latent_variable_quantile(unit_key, 
              model_id=unit_key['actual_behavior_model'])
          unit_psth.plot_unit_period_tuning(unit_key)
  
  
  #%%
  # =========== HELPER ==============
  
  def get_m_scale(shank_count):
      return 1350 - 150*shank_count
  
apis: numpy.sum numpy.diff matplotlib.pyplot.GridSpec numpy.nansum matplotlib.path.Path matplotlib.pyplot.figure numpy.logical_and numpy.abs numpy.cos numpy.histogram2d numpy.where numpy.linspace numpy.repeat matplotlib.pyplot.subplots numpy.arange numpy.hstack numpy.linalg.norm numpy.cumsum numpy.nanmax pandas.DataFrame matplotlib.pyplot.subplot numpy.array numpy.sin numpy.full