"""
Scatter plot of spike depth, amplitude and firing rate
==========================
Example of how to plot scatter plot of spike depths vs spike times vs spike firing rate
"""

from one.api import ONE
from brainbox.io.one import SpikeSortingLoader
from ibllib.atlas import AllenAtlas

import matplotlib.pyplot as plt
import numpy as np
from brainbox.ephys_plots import scatter_amp_depth_fr_plot, plot_brain_regions

one = ONE()
ba = AllenAtlas()
pid = 'da8dfec1-d265-44e8-84ce-6ae9c109b8bd'

# Load in spikesorting data
sl = SpikeSortingLoader(pid=pid, one=one, atlas=ba)
spikes, clusters, channels = sl.load_spike_sorting()
clusters = sl.merge_clusters(spikes, clusters, channels)

# Find the index of good spikes
good_idx = np.isin(spikes['clusters'], clusters['cluster_id'][clusters['label'] == 1])


# Make plot
fig, axs = plt.subplots(1, 3, gridspec_kw={'width_ratios': [3, 3, 1], 'wspace': 1})
_ = scatter_amp_depth_fr_plot(spikes['amps'], spikes['clusters'], spikes['depths'], spikes['times'], title='all units',
                              display=True, ax=axs[0])
_ = scatter_amp_depth_fr_plot(spikes['amps'][good_idx], spikes['clusters'][good_idx], spikes['depths'][good_idx],
                              spikes['times'][good_idx], title='good_units', display=True, ax=axs[1])
_ = plot_brain_regions(channels['atlas_id'], channel_depths=channels['axial_um'], ax=axs[2], label='right', display=True)

# Clean up plot
axs[1].get_yaxis().set_visible(False)

plt.show()
