"""
Get all trajectories from given provenance,
that have better QC than given status,
plot as scatter
"""

from ibl_pipeline import acquisition
from ibl_pipeline.analyses import behavior as behavior_analysis
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from one.api import ONE


traj = ONE().alyx.rest('trajectories', 'list', provenance='Planned',
                       django='probe_insertion__session__project__name__'
                              'icontains,ibl_neuropixel_brainwide_01,'
                              'probe_insertion__session__qc__lt,50')

# Ephys aligned histology track, Histology track, Micro-manipulator, Planned

# QC_CHOICES = [
#     (50, 'CRITICAL',),
#     (40, 'FAIL',),
#     (30, 'WARNING',),
#     (0, 'NOT_SET',),
#     (10, 'PASS',),
# ]

eids_traj = [p['session']['id'] for p in traj]
print(f'N traj with QC : {len(eids_traj)}')

# DATAJOINT query to combine with behavioral criterion
data_all = (acquisition.Session & [{'session_uuid': i_e} for i_e in eids_traj] &
            (behavior_analysis.SessionTrainingStatus & 'good_enough_for_brainwide_map=1'))
data_eids = data_all.proj('session_uuid')
df = data_eids.fetch(format='frame').reset_index()

eids_good = df['session_uuid'].values.tolist()
eids_good = [str(_) for _ in eids_good]
print(f'N traj good : {len(eids_good)}')

# Get ml / ap of only those that are good
traj_dict = {str(p['session']['id']): (p['x'], p['y']) for p in traj}
ml = np.array([traj_dict[eid][0] for eid in eids_good])
ap = np.array([traj_dict[eid][1] for eid in eids_good])

# Read CSV containing all x / y positions to be done
data = pd.read_csv(
    "/Users/gaelle/Documents/Git/Scrapbook/Needles/Firstpassmap_x_y.csv")
ap_fm = data['ap_um']
ml_fm = data['ml_um']

# Plot
fig, ax = plt.subplots()
ax.scatter(ap_fm, ml_fm, color='black', alpha=0.1)
ax.scatter(ap, ml, color='green', alpha=0.4)

ax.set_xlim(4000, -8000)  # decreasing x
