# Author: Olivier
# run "%qui qt" magic command from Ipython prompt for interactive mode
from pathlib import Path
import pickle

import numpy as np
from mayavi import mlab
from oneibl.one import ONE

import ibllib.plots
from iblutil.util import Bunch
import ibllib.atlas as atlas
from atlaselectrophysiology import rendering
import brainbox.io.one as bbone

cache_dir = '/datadisk/FlatIron/queries'
eids = ['aad23144-0e52-4eac-80c5-c4ee2decb198',
        '572a95d1-39ca-42e1-8424-5c9ffcb2df87',
        '57fd2325-67f4-4d45-9907-29e77d3043d7',
        '38d95489-2e82-412a-8c1a-c5377b5f1555',
        '4153bd83-2168-4bd4-a15c-f7e82f3f73fb',
        '614e1937-4b24-4ad3-9055-c8253d089919',
        'ee5c418c-e7fa-431d-8796-b2033e000b75']
probes = ['probe01', 'probe01', 'probe00', 'probe01', 'probe00', 'probe01', 'probe00']

# fetching data part
brain_atlas = atlas.AllenAtlas(25)
file_pickle = Path(cache_dir).joinpath('repeated_sites_channels.pkl')
if file_pickle.exists() or False:
    ins = pickle.load(open(file_pickle, 'rb'))
else:
    one = ONE()
    ins = Bunch({'eid': eids, 'probe_label': probes,
                 'insertion': [], 'channels': [], 'session': []})
    for eid, probe_label in zip(ins.eid, ins.probe_label):
        traj = one.alyx.rest('trajectories', 'list', session=eid,
                             provenance='Histology track', probe=probe_label)[0]
        ses = one.alyx.rest('sessions', 'read', id=eid)
        channels = bbone.load_channel_locations(eid=ses, one=one, probe=probe_label)[probe_label]
        insertion = atlas.Insertion.from_dict(traj)
        ins.insertion.append(insertion)
        ins.channels.append(channels)
        ins.session.append(ses)
    pickle.dump(ins, open(file_pickle, 'wb'))

# Display part
fig = rendering.figure(grid=False)  # set grid=True for ugly axes
for m in np.arange(len(ins.eid)):
    print(ins.session[m]['subject'], ins.session[m]['start_time'][:10],
          ins.session[m]['number'], ins.probe_label[m])
    color = ibllib.plots.color_cycle(m)
    mlapdv = brain_atlas.xyz2ccf(ins.insertion[m].xyz)
    # display the trajectories
    mlab.plot3d(mlapdv[:, 1], mlapdv[:, 2], mlapdv[:, 0],
                line_width=1, tube_radius=10, color=color)
    xyz_channels = np.c_[ins.channels[m].x, ins.channels[m].y, ins.channels[m].z]
    mlapdv_channels = brain_atlas.xyz2ccf(xyz_channels)
    # display the channels locations
    mlab.points3d(mlapdv_channels[:, 1], mlapdv_channels[:, 2], mlapdv_channels[:, 0],
                  color=color, scale_factor=50)
    # setup the labels at the top of the trajectories
    mlab.text3d(mlapdv[0, 1], mlapdv[0, 2], mlapdv[0, 0] - 500, ins.session[m]['subject'],
                line_width=4, color=tuple(color), figure=fig, scale=150)
