"""
Find certification recording sessions
=====================================
Use ONE to get the training status of a chosen subject or all subjects within a lab.
Training status is computed based on performance over latest 3 sessions (default) or last 3
sessions before a specified date.
"""
# Author: Gaelle Chapuis

# import modules
import numpy as np
from one.api import ONE
one = ONE()

dataset_types = ['spikes.times',
                 'spikes.clusters']

# eid1, det1 = one.search(project='ibl_certif_neuropix_recording',
#                         dataset=dataset_types, details=True)

eids, det = one.search(task_protocol='ephys_certification',
                       dataset=dataset_types, details=True)

sub = [p['subject'] for p in det]
# sub_unique = list(set(sub))

lab = [p['lab'] for p in det]
# lab_unique = list(set(lab))

# task = [p['task_protocol'] for p in det]
# task_unique = list(set(task))

# -- How many animals were used per lab
su, ind_su = np.unique(sub, return_index=True)
lab_arr = np.array(lab)
lu = lab_arr[ind_su]

for i_su in range(0, len(su)):
    # Find how many recordings were made with this animals
    sess_id = np.where(np.array(sub) == su[i_su])
    # Display
    tr_pl = one.alyx.rest('trajectories', 'list', subject=su[i_su], provenance='Planned')
    tr_tr = one.alyx.rest('trajectories', 'list', subject=su[i_su], provenance='Histology track')
    print(f'Subject: {su[i_su]} - Lab: {lu[i_su]} - N session: {len(sess_id[0])}'
          f' - N planned traces: {len(tr_pl)} - N tracked traces: {len(tr_tr)}')

# TODO -- How many recording sessions were done per lab
# prob_des = one.load_dataset(eids[0], 'probes.description.json')
# n_probe = len(prob_des[0])
