insectbrain_db

def authenticate(username=None, password=None, token=None):
    """Authenticate against Insect Brain DB.

    You can either provide username + password, or a token. Each token is only
    valid for 24h though. The better alternative is to provide your
    username + password as environment variables: `INSECT_BRAIN_DB_USER` and
    `INSECT_BRAIN_DB_PASSWORD`, respectively. If you are using these environment
    you don't need to bother with `authenticate()` at all.

    Parameters
    ----------
    username :      str, optional
                    Your username on Insect Brain DB.
    password :      str, optional
                    Your password on Insect Brain DB.
    token :         str, optional
                    A token. If provided you don't need to provide username +
                    password.

    """
    if not token and (not username and not password):
        raise ValueError('Must provide either username + password, or token '
                         '(or both).')

    if username:
        session.username = username
    if password:
        session.password = password

    if token:
        session.token = token
    else:
        session.fetch_token()

def get_brain_meshes(species: Union[str, int],
                     combine: bool = False,
                     max_threads: int = 4
                     ) -> Optional[List[Volume]]:
    """Fetch brain meshes for given species.

    Parameters
    ----------
    species:        str | int
                    Species for which to fetch brain volumes. Strings are
                    interpreted as names (scientific or common), integers as IDs.
    combine :       bool, optional
                    If True, will combine subvolumes (i.e. neuropils) into
                    a single navis.Volume - else will return list with volumes.
    max_threads :   int
                    Number of parallel threads to use for fetching meshes.

    Returns
    -------
    list of navis.Volume

    Examples
    --------
    >>> import navis
    >>> import navis.interfaces.insectbrain_db as ibdb
    >>> v = ibdb.get_brain_meshes('Desert Locust', combine_vols=True)
    >>> navis.plot3d(v)

    """
    # Get info with all available neuropils
    sp_info = get_species_info(species)

    # Go over all brains
    n_brains = len(sp_info.reconstructions)  # type: ignore
    n_reconstr = len([r for r in sp_info.reconstructions if r.get('viewer_files')])  # type: ignore
    logger.info(f'{n_reconstr} reconstruction(s) from {n_brains} brain(s) found')

    volumes: List[Volume] = []
    for brain in config.tqdm(sp_info.reconstructions,
                             disable=config.pbar_hide,
                             leave=config.pbar_leave,
                             desc='Brains'):  # type: ignore
        this_v = []
        # If no reconstructions, continue
        if not brain.get('viewer_files'):  # type: ignore
            continue

        with ThreadPoolExecutor(max_workers=max_threads) as executor:
            futures = {}
            for file in brain['viewer_files']:
                # If no file UUID, continue
                if not file['p_file']['uuid']:
                    continue
                filename = file['p_file']['file_name']
                f = executor.submit(_get_neuropil_mesh, file,)
                futures[f] = filename

            with config.tqdm(desc='Fetching',
                            total=len(futures),
                            leave=config.pbar_leave,
                            disable=len(futures) == 1 or config.pbar_hide) as pbar:
                for f in as_completed(futures):
                    name = futures[f]
                    pbar.update(1)
                    try:
                        this_v.append(f.result())
                    except Exception as exc:
                        print(f'{name} generated an exception:', exc)

        # Combine all volumes in this brain
        if combine:
            this_v = [Volume.combine(this_v)]
            this_v[0].color = (.85, .85, .85, .5)
            this_v[0].name = sp_info.scientific_name

        volumes += this_v

    return volumes

def get_meshes_experiment(id) -> 'NeuronList':
    """Fetch volumes associated with given experiment.

    Parameters
    ----------
    id :    int
            The experiment ID. See e.g. `list_datasets`.

    Returns
    -------
    list

    Examples
    --------
    >>> import navis.interfaces.insectbrain_db as ibdb
    >>> vols = ibdb.get_meshes_experiment(61)

    """
    # Make sure ID is integer
    id = int(id)

    # Get files associated with experiment
    files = list_experiment_files(id)

    # Figure out which files are skeletons
    me_files = files[files.file_name.str.endswith('.glb')]

    if me_files.empty:
        raise ValueError('Did not find any meshes associated with '
                         f'experiment {id}')

    volumes = []
    for f in config.tqdm(me_files.itertuples(),
                         desc='Downloading',
                         total=me_files.shape[0]):
        # Load the file
        r = requests.get(f.url)
        r.raise_for_status()

        name = '.'.join(f.file_name.split('.')[:-1])
        ext = f.file_name.split('.')[-1]

        file = io.BytesIO(r.content)
        scene = tm.load(file, file_type=ext)

        for obj in scene.geometry.values():
            v = Volume(obj.vertices, obj.faces, name=name)
            volumes.append(v)

    logger.info(f'Done! Found {len(volumes)} meshes.')

    return volumes

def get_skeletons(x, max_threads=4):
    """Fetch skeletons for given neuron(s).

    Parameters
    ----------
    x :             str | int | list thereof
                    Name(s) or ID(s) of neurons you want to fetch.
    max_threads :   int
                    Number of parallel threads to use for fetching skeletons.

    Returns
    -------
    navis.NeuronList

    Examples
    --------
    >>> import navis.interfaces.insectbrain_db as ibdb
    >>> neurons = ibdb.get_skeletons('TUps2-2')

    """
    if isinstance(x, (int, str, np.int32, np.int64)):
        neurons = [x]
    else:
        neurons = x

    # First fetch URLs for all neurons
    meta = []
    for x in neurons:
        if isinstance(x, str):
            q = search_neurons(name=x, partial_match=False)
            if q.empty:
                raise ValueError(f'No neuron with name "{x}" found')
            ids = q.id.values
        else:
            ids = x

        for i, id in enumerate(make_iterable(ids)):
            url = make_url(baseurl, 'api', 'v2', 'neuron', 'reconstruction',
                           neuron=id)
            info = session.get(url)

            if (not info
                or 'viewer_files' not in info[0]
                or not info[0]['viewer_files']):
                raise ValueError(f'Neuron {x} ({id}) has no skeleton.')

            meta.append(info[0])

    return _get_skeletons(meta, max_threads=max_threads)

def get_skeletons_experiment(id) -> 'NeuronList':
    """Fetch all skeletons for given experiment.

    Parameters
    ----------
    id :    int
            The experiment ID. See e.g. `list_datasets`.

    Returns
    -------
    NeuronList

    Examples
    --------
    >>> import navis.interfaces.insectbrain_db as ibdb
    >>> nl = ibdb.get_skeletons_experiment(61)

    """
    # Make sure ID is integer
    id = int(id)

    # Get files associated with experiment
    files = list_experiment_files(id)

    # Figure out which files are skeletons
    sk_files = files[files.file_name.str.contains('skeleton') | files.file_name.str.endswith('.gz')]

    if sk_files.empty:
        raise ValueError('Did not find any skeleton files associated with '
                         f'experiment {id}')

    skeletons = []
    for f in sk_files.itertuples():
        logger.info(f'Downloading {f.file_name}')
        # Load the file
        r = requests.get(f.url)
        r.raise_for_status()

        # Files appear to be json-formatted and not compressed
        data = r.json()

        for i, neuron in enumerate(data['data']):
            for sk in neuron['skeletons']:
                # Load SWC table
                swc = pd.DataFrame(sk['data'],
                                   columns=['node_id', 'skeleton_id',
                                            'x', 'y', 'z', 'radius',
                                            'parent_id'])
                # Some cleaning up
                swc.drop('skeleton_id', axis=1, inplace=True)
                swc['parent_id'] = swc.parent_id.fillna(-1).astype(int)
                # Create neuron
                tn = TreeNeuron(swc,
                                id=sk.get('id', 1),
                                name=neuron.get('name', 'NA'),
                                annotations=neuron.get('annotations', []),
                                soma=None)
                skeletons.append(tn)
    logger.info(f'Done! Found {len(skeletons)} skeletons.')

    return NeuronList(skeletons)

def get_skeletons_species(species, max_threads=4):
    """Fetch all skeletons for given species.

    Note that some neurons might have multiple reconstructions. They will
    show up with the same ID with different names.

    Parameters
    ----------
    species :       str | int
                    Name or ID of a species to fetch skeletons for.
    max_threads :   int
                    Number of parallel threads to use for fetching skeletons.

    Returns
    -------
    navis.NeuronList

    Examples
    --------
    >>> import navis.interfaces.insectbrain_db as ibdb
    >>> neurons = ibdb.get_skeletons_species('Desert Locust')

    """
    if isinstance(species, str):
        species = _get_species_id(species)

    # First fetch URLs for all neurons
    url = make_url(baseurl, 'api', 'v2', 'neuron', 'reconstruction',
                   neuron__species=species)
    meta = session.get(url)

    meta = [e for e in meta if e['viewer_files']]

    return _get_skeletons(meta, max_threads=max_threads)

def list_datasets() -> pd.DataFrame:
    """List publication datasets and associated experiments.

    Returns
    -------
    pandas.DataFrame
            DataFrame with available datasets.

    Examples
    --------
    >>> import navis.interfaces.insectbrain_db as ibdb
    >>> datasets = ibdb.list_datasets()

    """
    url = make_url(baseurl, 'api', 'publications', 'experiments?offset=0&limit=500')

    return _sort_columns(pd.DataFrame.from_records(session.get(url)['results']))

def list_experiment_files(id) -> pd.DataFrame:
    """List files associated with given experiment.

    Parameters
    ----------
    id :    int
            The experiment ID. See e.g. `list_datasets`.

    Returns
    -------
    pandas.DataFrame
            DataFrame with files.

    Examples
    --------
    >>> import navis.interfaces.insectbrain_db as ibdb
    >>> files = ibdb.list_experiment_files(61)

    """
    url = make_url(baseurl, 'api', 'v2', 'experiment', id, 'file')

    return _sort_columns(pd.DataFrame.from_records(session.get(url)))

def search_neurons(name=None, short_name=None, species=None, sex=None,
                   arborization=None, partial_match=True) -> pd.DataFrame:
    """Search for neurons matching given parameters.

    Parameters
    ----------
    name :          str, optional
                    Name of the neuron.
    short_name :    str, optional
                    Short name of the neuron.
    species :       str | int, optional
                    Name or ID of the species. Can be common or scientific name.
    sex :           "FEMALE" | "MALE" | "UNKNOWN", optional
                    Sex of the neuron.
    arborization :  str, optional
                    Restrict to neurons having arborizations in given neuropil.
    partial_match : bool
                    Whether to allow partial matches (does not apply for species).

    Returns
    -------
    pandas.DataFrame

    Examples
    --------
    >>> import navis.interfaces.insectbrain_db as ibdb
    >>> neurons = ibdb.search_neurons(species='Desert Locust')

    """
    # Construct query
    options = {}
    if species:
        if not isinstance(species, int):
            species = _get_species_id(species)
        options['species'] = species

    for key, value in zip(['name', 'short_name', 'sex',
                           'arborization_region__structure'],
                          [name, short_name, sex, arborization]):
        if not value:
            continue
        if partial_match:
            key += '__icontains'
        options[key] = value

    url = make_url(baseurl, 'api', 'v2', 'neuron', **options)

    resp = requests.get(url)

    resp.raise_for_status()

    return _sort_columns(pd.DataFrame.from_records(resp.json()))