import matplotlib.pyplot as plt
import numpy as np
from one.api import ONE
from neurodsp.utils import fcn_cosine

import ibllib.atlas as atlas
from ibllib.pipes.histology import coverage

ba = atlas.AllenAtlas()
# trajs = ONE().alyx.rest('trajectories', 'list', provenance='Ephys aligned histology track')

trajs = ONE().alyx.rest('trajectories', 'list', provenance='Micro-manipulator',
                        django='probe_insertion__session__project__name__icontains,'
                               'ibl_neuropixel_brainwide_01,'
                               'probe_insertion__session__qc__lt,40')

full_coverage = coverage(trajs)

fig, axs = plt.subplots(2, 2)
ax = ba.plot_hslice(-4000 * 1e-6, volume=full_coverage, ax=axs[0, 0])
ax.set_title("horizontal slice at dv=-4mm")
ax = ba.plot_sslice(ml_coordinate=-0.002, volume=full_coverage, ax=axs[0, 1])
ax.set_title("sagittal slice at ml=-2mm")
ax = ba.plot_cslice(ap_coordinate=-0.003, volume=full_coverage, ax=axs[1, 1])
ax.set_title("coronal slice at ap=-3mm")

axs[1, 0].plot(np.linspace(0, 200), 1 - fcn_cosine([100, 150])(np.linspace(0, 200)))
axs[1, 0].set_xlabel('distance from nearest active site (um)')
axs[1, 0].set_ylabel('weight')
