!pip install vfb-connect --upgrade

# Import libs and initialise API objects
from vfb_connect.cross_server_tools import VfbConnect
import pandas as pd

vc = VfbConnect()

import pymaid
import navis

navis.set_pbars(jupyter=False)
pymaid.set_pbars(jupyter=False)

# Connect to the VFB CATMAID server hosting the FAFB data
rm = pymaid.connect_catmaid(server="https://fafb.catmaid.virtualflybrain.org/", api_token=None, max_threads=10)

# Test call to see if connection works 
print(f'Server is running CATMAID version {rm.catmaid_version}')

# Many functions return JSON-compatible nested data structures. This function coverts them to DataFrame.
def summary_2_df(summary, sort=None):
    """Convert summary to DataFrame.  Optionally specify a set of columns to sort as a list of strings"""
    if sort:
        return pd.DataFrame.from_records(summary).sort_values(sort)
    else:
        return pd.DataFrame.from_records(summary)

# lets take some examples from a discovery query on the previous spreadhseet
sayin_tab = pd.DataFrame.from_records(vc.get_instances_by_dataset('Sayin2019', summary=True))
sayin_tab

oct_VPM3 = summary_2_df(vc.get_instances('octopaminergic VPM3 neuron', summary=True))
oct_VPM3

oct_VPM3_images = vc.neo_query_wrapper.get_images(oct_VPM3['id'], stomp=True, template='JRC2018Unisex', image_folder = 'oct_VPM3b')
oct_VPM3_images

oct_VPM3_images = vc.get_images_by_type('octopaminergic VPM3 neuron', stomp=True, template='JRC2018Unisex', image_folder = 'oct_VPM3')
nl = navis.read_swc('./oct_VPM3')
navis.plot3d(nl)

summary_2_df(vc.neo_query_wrapper.get_anatomical_individual_TermInfo(['VFB_001012ay'], summary=True))

similar_to_EPG6L1 = vc.get_similar_neurons('VFB_001012ay')
similar_to_EPG6L1 

EPG6L1_similar_images = vc.neo_query_wrapper.get_images(list(similar_to_EPG6L1['id']) + ['VFB_001012ay'], 
                                template='JRC2018Unisex',
                                image_folder='EPG6L1',
                                stomp=True) # Note - stomp allows an existing folder to be over-written


EPG6L1_similar_images

list(EPG6L1_similar_images['parents_label'])

import trimesh

jrc2018u = trimesh.load_remote('https://github.com/schlegelp/navis-flybrains/raw/main/flybrains/meshes/JRC2018U.ply')
jrc2018u = navis.Volume(jrc2018u, color=(.9, .9, .9, .1), name='JRC2018U')
jrc2018u

import matplotlib.pyplot as plt

# nl = navis.read_swc('./EPG6L1')
fig, ax = navis.plot2d([nl, jrc2018u], figsize=(12, 12), method='3d_complex')

ax.elev = -90
ax.azim = -90
ax.dist = 5

#plt.savefig('EPG6L1.png', transparent=True, dpi=400)

## Generate a link to display these neurons on VFB

print(vc.get_vfb_link(list(similar_to_EPG6L1['id']) + ['VFB_001012ay'], template='JRC2018Unisex'))