"""
Finding brain regions of clusters
=================================

Finds the brain region that each cluster has been assigned to and creates a plot of the number of
clusters per unique brain region for the probe. If the ephys aligned trajectory is
available, the channels will be obtained from this trajectory, otherwise the non-aligned histology
track is used. If neither are available, channels have not been allocated and a warning is given in
the terminal. In this case no plot is generated.
"""


from one.api import ONE
import numpy as np
import matplotlib.pyplot as plt
from ibllib.atlas import BrainRegions

one = ONE(base_url='https://openalyx.internationalbrainlab.org', silent=True)

# Specify subject, date and probe we are interested in
subject = 'CSHL049'
date = '2020-01-08'
sess_no = 1
probe_label = 'probe00'
eid = one.search(subject=subject, date=date, number=sess_no)[0]

cluster_chans = one.load_dataset(eid, 'clusters.channels.npy', collection=f'alf/{probe_label}')


aligned_traj = one.alyx.rest('trajectories', 'list', subject=subject, session=eid,
                             probe=probe_label, provenance='Ephys aligned histology track')

if len(aligned_traj) > 0:
    print('Getting channels for provenance ' + aligned_traj[0]['provenance'])

    channels = one.alyx.rest('channels', 'list', trajectory_estimate=aligned_traj[0]['id'])

    chans = {'atlas_id': np.array([ch['brain_region'] for ch in channels]),
             'x': np.array([ch['x'] for ch in channels]) / 1e6,
             'y': np.array([ch['y'] for ch in channels]) / 1e6,
             'z': np.array([ch['z'] for ch in channels]) / 1e6,
             'axial_um': np.array([ch['axial'] for ch in channels]),
             'lateral_um': np.array([ch['lateral'] for ch in channels])}

else:
    histology_traj = one.alyx.rest('trajectories', 'list', subject=subject, session=eid,
                                   probe=probe_label, provenance='Histology track')
    if len(histology_traj) > 0:
        print('Getting channels for provenance ' + histology_traj[0]['provenance'])

        channels = one.alyx.rest('channels', 'list', trajectory_estimate=histology_traj[0]['id'])

        chans = {'atlas_id': np.array([ch['brain_region'] for ch in channels]),
                 'x': np.array([ch['x'] for ch in channels]) / 1e6,
                 'y': np.array([ch['y'] for ch in channels]) / 1e6,
                 'z': np.array([ch['z'] for ch in channels]) / 1e6,
                 'axial_um': np.array([ch['axial'] for ch in channels]),
                 'lateral_um': np.array([ch['lateral'] for ch in channels])}

    else:
        print(f'No histology or ephys aligned trajectory for session: {eid} and '
              f'probe: {probe_label}, no channels available')
        chans = None


if chans is not None:
    r = BrainRegions()
    chans['acronym'] = r.get(ids=chans['atlas_id']).acronym
    chans['rgb'] = r.get(ids=chans['atlas_id']).rgb
    cluster_brain_region = chans['acronym'][cluster_chans]
    cluster_colour = chans['rgb'][cluster_chans]
    cluster_xyz = np.c_[chans['x'], chans['y'], chans['z']][cluster_chans]
    regions, idx, n_clust = np.unique(cluster_brain_region, return_counts=True, return_index=True)

    region_cols = cluster_colour[idx, :]
    fig, ax = plt.subplots()
    ax.bar(x=np.arange(len(regions)), height=n_clust, tick_label=regions, color=region_cols / 255)
    ax.set_xlabel('Brain region acronym')
    ax.set_ylabel('No. of clusters')
    plt.show()
