cave_utils

def fetch_neurons(
    x,
    lod,
    with_synapses,
    client,
    parallel,
    max_threads,
    materialization="auto",
    **kwargs,
):
    """Fetch neuron meshes.

    Notes
    -----
    Synapses will be attached to the closest vertex on the mesh.

    Parameters
    ----------
    x :             str | int | list-like
                    Segment ID(s). Multiple Ids can be provided as list-like.
    lod :           int
                    Level of detail. Higher ``lod`` = coarser. This parameter
                    is ignored if the data source does not support multi-level
                    meshes.
    with_synapses : bool, optional
                    If True will also attach synapses as ``.connectors``.
    client :        CAVEclient
                    The CAVEclient with which to interact.
    parallel :      bool
                    If True, will use parallel threads to fetch data.
    max_threads :   int
                    Max number of parallel threads to use.
    materialization : "auto" | int
                    Which materialization version to use to look up somas and synapses
                    (if applicable). If "auto" (default) will try to find the most
                    recent version that contains the given root IDs. If an
                    integer is provided will use that version.
    **kwargs
                    Keyword arguments are passed through to the initialization
                    of the ``navis.MeshNeurons``.

    Returns
    -------
    navis.Neuronlist
                    Containing :class:`navis.MeshNeuron`.

    """
    x = utils.make_iterable(x, force_type=int)

    vol = _get_cloudvol(client.info.segmentation_source())  # this is cached

    try:
        somas = _get_somas(x, client=client, materialization=materialization)
        soma_pos = somas.set_index("pt_root_id").pt_position.to_dict()
    except BaseException as e:
        logger.warning("Failed to fetch somas via nucleus segmentation" f"(){e})")
        soma_pos = {}

    nl = []
    if max_threads > 1 and parallel:
        with ThreadPoolExecutor(max_workers=max_threads) as executor:
            futures = {}
            for id in x:
                f = executor.submit(
                    _fetch_single_neuron,
                    id,
                    vol=vol,
                    lod=lod,
                    client=client,
                    with_synapses=with_synapses,
                    materialization=materialization,
                    **kwargs,
                )
                futures[f] = id

            with config.tqdm(
                desc="Fetching",
                total=len(x),
                leave=config.pbar_leave,
                disable=len(x) == 1 or config.pbar_hide,
            ) as pbar:
                for f in as_completed(futures):
                    id = futures[f]
                    pbar.update(1)
                    try:
                        nl.append(f.result())
                    except Exception as exc:
                        print(f"{id} generated an exception:", exc)
    else:
        for id in config.tqdm(
            x,
            desc="Fetching",
            leave=config.pbar_leave,
            disable=len(x) == 1 or config.pbar_hide,
        ):
            n = _fetch_single_neuron(
                id,
                vol=vol,
                lod=lod,
                client=client,
                with_synapses=with_synapses,
                materialization=materialization,
                **kwargs,
            )
            nl.append(n)

    nl = NeuronList(nl)

    for n in nl:
        if n.id in soma_pos:
            # For VoxelResolution see client.materialize.get_table_metadata('nucleus_detection_v0')
            # (attached to df as 'table_voxel_resolution')
            n.soma_pos = (
                np.array(soma_pos[n.id]) * somas.attrs["table_voxel_resolution"]
            )
        else:
            n.soma_pos = None

    return nl

def get_voxels(x, mip, bounds, client):
    """Fetch voxels making a up given root ID.

    Parameters
    ----------
    x :             int
                    A single root ID.
    mip :           int
                    Scale at which to fetch voxels.
    bounds :        list, optional
                    Bounding box [xmin, xmax, ymin, ymax, zmin, zmax] in voxel
                    space. For example, the voxel resolution for mip 0
                    segmentation is 8 x 8 x 40 nm.
    client :        CAVEclient
                    The CAVEclient with which to interact.

    Returns
    -------
    voxels :        (N, 3) np.ndarray
                    In voxel space according to `mip`.

    """
    # Need to get the graphene (not the precomputed) version of the data
    vol_graphene = cv.CloudVolume(
        client.chunkedgraph.cloudvolume_path, use_https=True, progress=False
    )
    url = client.info.get_datastack_info()["segmentation_source"]
    vol_prec = _get_cloudvol(url)

    # Get L2 chunks making up this neuron
    l2_ids = client.chunkedgraph.get_leaves(x, stop_layer=2)

    # Turn l2_ids into chunk indices
    l2_ix = [
        np.array(vol_graphene.mesh.meta.meta.decode_chunk_position(l2)) for l2 in l2_ids
    ]
    l2_ix = np.unique(l2_ix, axis=0)

    # Convert to nm
    l2_nm = np.asarray(_chunks_to_nm(l2_ix, vol=vol_graphene))

    # Convert back to voxel space (according to mip)
    l2_vxl = l2_nm // vol_prec.meta.scales[mip]["resolution"]

    voxels = []
    ch_size = np.array(vol_graphene.mesh.meta.meta.graph_chunk_size)
    ch_size = ch_size // (vol_prec.mip_resolution(mip) / vol_prec.mip_resolution(0))
    ch_size = np.asarray(ch_size).astype(int)
    old_mip = vol_prec.mip

    if not isinstance(bounds, type(None)):
        bounds = np.asarray(bounds)
        if not bounds.ndim == 1 or len(bounds) != 6:
            raise ValueError("`bounds` must be [xmin, xmax, ymin, ymax, zmin, zmax]")
        l2_vxl = l2_vxl[np.all(l2_vxl >= bounds[::2], axis=1)]
        l2_vxl = l2_vxl[np.all(l2_vxl < bounds[1::2] + ch_size, axis=1)]

    try:
        vol_prec.mip = mip
        for ch in config.tqdm(l2_vxl, desc="Loading"):
            ct = vol_prec[
                ch[0] : ch[0] + ch_size[0],
                ch[1] : ch[1] + ch_size[1],
                ch[2] : ch[2] + ch_size[2],
            ][:, :, :, 0]
            this_vxl = np.dstack(np.where(ct == x))[0]
            this_vxl = this_vxl + ch
            voxels.append(this_vxl)
    except BaseException:
        raise
    finally:
        vol_prec.mip = old_mip
    voxels = np.vstack(voxels)

    if not isinstance(bounds, type(None)):
        voxels = voxels[np.all(voxels >= bounds[::2], axis=1)]
        voxels = voxels[np.all(voxels < bounds[1::2], axis=1)]

    return voxels

def roots_to_mat(
    ids,
    client,
    verbose=True,
    allow_multiple=False,
    raise_missing=True,
):
    """Find a materialization version (or live) for given root ID(s).

    Parameters
    ----------
    ids :           int | iterable
                    Root ID(s) to check.
    client :        CAVEclient
                    The CAVEclient with which to interact.
    verbose :       bool
                    Whether to print results of search.
    allow_multiple : bool
                    If True, will track if IDs can be found spread across multiple
                    materialization versions if there is no single one containing
                    all.
    raise_missing : bool
                    Only relevant if `allow_multiple=True`. If False, will return
                    versions even if some IDs could not be found.

    Returns
    -------
    version :       int | "live"
                    A single version (including "live") that contains all given
                    root IDs.
    versions :      np.ndarray
                    If no single version was found and `allow_multiple=True` will
                    return a vector of `len(ids)` with the latest version at which
                    the respective ID can be found.
                    Important: "live" version will be return as -1!
                    If `raise_missing=False` and one or more root IDs could not
                    be found in any of the available materialization versions
                    these IDs will be return as version 0.

    """
    ids = utils.make_iterable(ids)

    # For each ID track the most recent valid version
    latest_valid = np.zeros(len(ids), dtype=np.int32)

    # Get the meta data for the available materialization versions
    # This is a list of dicts where each dict has a "time_stamp" key
    vmeta = client.materialize.get_versions_metadata()

    # Sort by "time_stamp"
    vmeta = sorted(vmeta, key=lambda x: x["time_stamp"], reverse=True)

    # Go over each version (start with the most recent)
    for i, mat in enumerate(vmeta):
        ts_m = mat["time_stamp"]
        version = mat["version"]

        # Check which root IDs were valid at the time
        is_valid = client.chunkedgraph.is_latest_roots(ids, timestamp=ts_m)

        # Update latest valid versions
        latest_valid[(latest_valid == 0) & is_valid] = version

        if all(is_valid):
            if verbose and not SILENCE_FIND_MAT_VERSION:
                print(f"Using materialization version {version}.")
            return version

    # If no single materialized version can be found, see if we can get
    # by with the live materialization
    is_latest = client.chunkedgraph.is_latest_roots(ids, timestamp=None)
    latest_valid[(latest_valid == 0) & is_latest] = -1  # track "live" as -1
    if all(is_latest) and dataset != "public":  # public does not have live
        if verbose:
            print("Using live materialization")
        return "live"

    if allow_multiple and any(latest_valid != 0):
        if all(latest_valid != 0):
            if verbose and not SILENCE_FIND_MAT_VERSION:
                print(
                    f"Found root IDs spread across {len(np.unique(latest_valid))} "
                    "materialization versions."
                )
            return latest_valid

        msg = (
            f"Found root IDs spread across {len(np.unique(latest_valid)) - 1} "
            f"materialization versions but {(latest_valid == 0).sum()} IDs "
            "do not exist in any of the materialized tables."
        )

        if not raise_missing:
            if verbose and not SILENCE_FIND_MAT_VERSION:
                print(msg)
            return latest_valid
        else:
            raise MaterializationMatchError(msg)

    if dataset not in ("public, "):
        raise MaterializationMatchError(
            "Given root IDs do not (co-)exist in any of the available "
            "materialization versions (including live). Try updating "
            "root IDs and rerun your query."
        )
    else:
        raise MaterializationMatchError(
            "Given root IDs do not (co-)exist in any of the available "
            "public materialization versions. Please make sure that "
            "the root IDs do exist and rerun your query."
        )