import cloudvolume
import fastremap
import numpy as np
import pandas as pd
from annotationframeworkclient import FrameworkClient
from meshparty import skeleton, skeletonize, trimesh_io, meshwork

from . import chunk_tools
from . import skel_utils as sk_utils
from . import utils
from . import pcg_anno

DEFAULT_VOXEL_RESOLUTION = [4, 4, 40]
DEFAULT_COLLAPSE_RADIUS = 7500.0


def build_spatial_graph(lvl2_edge_graph, cv):
    """Extract spatial graph and level 2 id lookups from chunkedgraph "lvl2_graph" endpoint.

    Parameters
    ----------
    lvl2_edge_graph : array
        Nx2 edge list of level 2 ids
    cv : cloudvolume.CloudVolume
        Associated cloudvolume object

    Returns
    -------
    eg_arr_rm : np.array
        Nx2 edge list of indices remapped to integers starting at 0 through M, the number of unique level 2 ids.
    l2dict : dict
        Dict with level 2 ids as keys and vertex index as values
    l2dict_reversed : dict
        Dict with vertex index as keys and level 2 id as values
    x_ch : np.array
        Mx3 array of vertex locations in chunk index space.
    """
    lvl2_edge_graph = np.unique(np.sort(lvl2_edge_graph, axis=1), axis=0)
    lvl2_ids = np.unique(lvl2_edge_graph)
    l2dict = {l2id: ii for ii, l2id in enumerate(lvl2_ids)}
    eg_arr_rm = fastremap.remap(lvl2_edge_graph, l2dict)
    l2dict_reversed = {ii: l2id for l2id, ii in l2dict.items()}

    x_ch = [np.array(cv.mesh.meta.meta.decode_chunk_position(l))
            for l in lvl2_ids]
    return eg_arr_rm, l2dict, l2dict_reversed, x_ch


def chunk_index_skeleton(root_id,
                         client=None,
                         datastack_name=None,
                         cv=None,
                         root_point=None,
                         invalidation_d=3,
                         return_mesh=False,
                         return_l2dict=False,
                         return_mesh_l2dict=False,
                         root_point_resolution=[4, 4, 40],
                         root_point_search_radius=300,
                         n_parallel=1):
    """Generate a basic skeleton with chunked-graph index vertices.

    Parameters
    ----------
    root_id : np.uint64
        Neuron root id
    client : annotationframeworkclient.FrameworkClient, optional
        FrameworkClient for a datastack, by default None. If None, you must specify a datastack name.
    datastack_name : str, optional
        Datastack name to create a FrameworkClient, by default None. Only used if client is None.
    cv : cloudvolume.CloudVolume, optional
        CloudVolume associated with the object, by default None. If None, one is created based on the client info.
    root_point : array, optional
        Point in voxel space to set the root vertex. By default None, which makes a random tip root.
    invalidation_d : int, optional
        TEASAR invalidation radius in chunk space, by default 3
    return_mesh : bool, optional
        If True, returns the pre-skeletonization mesh with vertices in chunk index space, by default False
    return_l2dict : bool, optional
        If True, returns the level 2 id to vertex index dict. By default True
    n_parallel : int, optional
        Sets number of parallel threads for cloudvolume, by default 1

    Returns
    -------
    sk : meshparty.skeleton.Skeleton
        Skeleton object
    mesh : meshparty.trimesh_io.Mesh
        Mesh object, only if return_mesh is True
    level2_dict : dict
        Level 2 id to vertex map, only if return_l2dict is True.
    """
    if client is None:
        client = FrameworkClient(datastack_name)
    if cv is None:
        cv = cloudvolume.CloudVolume(client.info.segmentation_source(), parallel=n_parallel,
                                     use_https=True, progress=False, bounded=False)

    lvl2_eg = client.chunkedgraph.level2_chunk_graph(root_id)

    eg, l2dict_mesh, l2dict_r_mesh, x_ch = build_spatial_graph(
        lvl2_eg, cv)
    mesh_chunk = trimesh_io.Mesh(vertices=x_ch, faces=[], link_edges=eg)

    if root_point is not None:
        lvl2_root_chid, lvl2_root_loc = chunk_tools.get_closest_lvl2_chunk(
            root_point,
            root_id,
            client=client,
            cv=None,
            radius=root_point_search_radius,
            voxel_resolution=root_point_resolution,
            return_point=True)  # Need to have cv=None because of a cloudvolume inconsistency
        root_mesh_index = l2dict_mesh[lvl2_root_chid]
    else:
        root_mesh_index = None

    sk_ch = skeletonize.skeletonize_mesh(
        mesh_chunk, invalidation_d=invalidation_d,
        collapse_soma=False, compute_radius=False, cc_vertex_thresh=0,
        root_index=root_mesh_index, remove_zero_length_edges=False)

    l2dict, l2dict_r = sk_utils.filter_l2dict(sk_ch, l2dict_r_mesh)

    out_list = [sk_ch]
    if return_mesh:
        out_list.append(mesh_chunk)
    if return_l2dict:
        out_list.append((l2dict, l2dict_r))
    if return_mesh_l2dict:
        out_list.append((l2dict_mesh, l2dict_r_mesh))
    if len(out_list) == 1:
        return out_list[0]
    else:
        return tuple(out_list)


def refine_chunk_index_skeleton(
    sk_ch,
    l2dict_reversed,
    cv,
    refine_inds='all',
    scale_chunk_index=True,
    root_location=None,
    nan_rounds=20,
    return_missing_ids=False,
    segmentation_fallback=True,
    fallback_mip=2,
    cache=None,
    save_to_cache=False,
):
    """Refine skeletons in chunk index space to Euclidean space.

    Parameters
    ----------
    sk_ch : meshparty.skeleton.Skeleton
        Skeleton in chunk index space
    l2dict_reversed : dict
        Mapping between skeleton vertex index and level 2 id.
    cv : cloudvolume.CloudVolume
        Associated cloudvolume
    refine_inds : str, None or list-like, optional
        Skeleton indices to refine, 'all', or None. If 'all', does all skeleton indices.
        If None, downloads no index but can use other options.
        By default 'all'.
    scale_chunk_index : bool, optional
        If True, maps unrefined chunk index locations to the center of the chunk in
        Euclidean space, by default True
    root_location : list-like, optional
        3-element euclidean space location to which to map the root vertex location, by default None
    nan_rounds : int, optional
        Number of passes to smooth over any missing values by averaging proximate vertex locations.
        Only used if refine_inds is 'all'. Default is 20.
    return_missing_ids : bool, optional
        If True, returns ids of any missing level 2 meshes. Default is False
    cache : str
        Filename for a sqlite database storing locations associated with level 2 ids.

    Returns
    -------
    meshparty.skeleton.Skeleton
        Skeleton with remapped vertex locations
    """
    if nan_rounds is None:
        convert_missing = True
    else:
        convert_missing = False

    refine_out = chunk_tools.refine_vertices(sk_ch.vertices,
                                             l2dict_reversed=l2dict_reversed,
                                             cv=cv,
                                             refine_inds=refine_inds,
                                             scale_chunk_index=scale_chunk_index,
                                             convert_missing=convert_missing,
                                             return_missing_ids=return_missing_ids,
                                             segmentation_fallback=segmentation_fallback,
                                             fallback_mip=fallback_mip,
                                             cache=cache,
                                             save_to_cache=save_to_cache)
    if return_missing_ids:
        new_verts, missing_ids = refine_out
    else:
        new_verts = refine_out

    if root_location is not None:
        new_verts[sk_ch.root] = root_location

    l2_sk = skeleton.Skeleton(
        vertices=new_verts,
        edges=sk_ch.edges,
        root=sk_ch.root,
        remove_zero_length_edges=False,
        mesh_index=sk_ch.mesh_index,
        mesh_to_skel_map=sk_ch.mesh_to_skel_map,
    )

    if refine_inds == 'all':
        sk_utils.fix_nan_verts(l2_sk, num_rounds=nan_rounds)

    if return_missing_ids:
        return l2_sk, missing_ids
    else:
        return l2_sk


def pcg_skeleton(root_id,
                 client=None,
                 datastack_name=None,
                 cv=None,
                 refine='all',
                 root_point=None,
                 root_point_resolution=[4, 4, 40],
                 root_point_search_radius=300,
                 collapse_soma=False,
                 collapse_radius=10_000.0,
                 invalidation_d=3,
                 return_mesh=False,
                 return_l2dict=False,
                 return_l2dict_mesh=False,
                 return_missing_ids=False,
                 nan_rounds=20,
                 segmentation_fallback=True,
                 fallback_mip=2,
                 cache=None,
                 save_to_cache=False,
                 n_parallel=1):
    """Create a euclidean-space skeleton from the pychunkedgraph

    Parameters
    ----------
    root_id : uint64
        Root id of the neuron to skeletonize
    client : annotationframeworkclient.FrameworkClientFull or None, optional
        Pre-specified framework client for the pcg. If this is not set, datastack_name must be provided. By default None
    datastack_name : str or None, optional
        If no client is specified, a FrameworkClient is created with this datastack name, by default None
    cv : cloudvolume.CloudVolume or None, optional
        Prespecified cloudvolume instance. If None, uses the client info to make one, by default None
    refine : 'all', 'ep', 'bp', 'epbp', 'bpep', or None, optional
        Selects how to refine vertex locations by downloading mesh chunks. Unrefined vertices are placed in the
        center of their chunk in euclidean space.
        * 'all' refines all vertex locations. (Default)
        * 'ep' refines end points only
        * 'bp' refines branch points only
        * 'bpep' or 'epbp' refines both branch and end points.
        * None refines no points.
    root_point : array-like or None, optional
        3 element xyz location for the location to set the root in units set by root_point_resolution,
        by default None. If None, a distal tip is selected.
    root_point_resolution : array-like, optional
        Resolution in euclidean space of the root_point, by default [4, 4, 40]
    root_point_search_radius : int, optional
        Distance in euclidean space to look for segmentation when finding the root vertex, by default 300
    collapse_soma : bool, optional,
        If True, collapses vertices within a given radius of the root point into the root vertex, typically to better
        represent primary neurite branches. Requires a specified root_point. Default if False.
    collapse_radius : float, optional
        Max distance in euclidean space for soma collapse. Default is 10,000 nm (10 microns).
    invalidation_d : int, optional
        Invalidation radius in hops for the mesh skeletonization along the chunk adjacency graph, by default 3
    return_mesh : bool, optional
        If True, returns the mesh in chunk index space, by default False
    return_l2dict : bool, optional
        If True, returns the tuple (l2dict, l2dict_r), by default False.
        l2dict maps all neuron level2 ids to skeleton vertices. l2dict_r maps skeleton indices to their direct level 2 id.
    return_l2dict_mesh : bool, optional
        If True, returns the tuple (l2dict_mesh, l2dict_mesh_r), by default False.
        l2dict_mesh maps neuron level 2 ids to mesh vertices, l2dict_r maps mesh indices to level 2 ids.
    return_missing_ids : bool, optional
        If True, returns level 2 ids that were missing in the chunkedgraph, by default False. This can be useful
        for submitting remesh requests in case of errors.
    nan_rounds : int, optional
        Maximum number of rounds of smoothing to eliminate missing vertex locations in the event of a
        missing level 2 mesh, by default 20. This is only used when refine=='all'.
    segmentation_fallback : bool, optional
        If True, uses the segmentation in cases of missing level 2 meshes. This is slower but more robust.
        Default is True.
    cache : str or None, optional
        Filename to a sqlite database with cached lookups for l2 ids. Optional, default is None.
    n_parallel : int, optional
        Number of parallel downloads passed to cloudvolume, by default 1

    Returns
    -------
    sk_l2 : meshparty.skeleton.Skeleton
        Skeleton with vertices in euclidean space
    mesh_l2 : meshparty.mesh.Mesh, optional
        Mesh with vertices in chunk index space. Only if return_mesh is True.
    (l2dict, l2dict_r) : (dict, dict), optional
        Mappings between level 2 ids and skeleton indices. Only if return_l2dict is True.
    (l2dict_mesh, l2dict_mesh_r) : (dict, dict), optional
        Mappings between level 2 ids and mesh indices. Only if return_l2dict_mesh is True.
    missing_ids : np.array, optional
        List of level 2 ids with missing mesh fragments. Only if return_missing_ids is True.
   """
    if client is None:
        client = FrameworkClient(datastack_name)

    if cv is None:
        cv = cloudvolume.CloudVolume(client.info.segmentation_source(), parallel=n_parallel,
                                     use_https=True, progress=False, bounded=False)

    sk_ch, mesh_ch, (l2dict, l2dict_r), (l2dict_mesh, l2dict_mesh_r) = chunk_index_skeleton(root_id,
                                                                                            client=client,
                                                                                            datastack_name=datastack_name,
                                                                                            cv=cv,
                                                                                            root_point=root_point,
                                                                                            root_point_resolution=root_point_resolution,
                                                                                            root_point_search_radius=root_point_search_radius,
                                                                                            invalidation_d=invalidation_d,
                                                                                            return_mesh=True,
                                                                                            return_mesh_l2dict=True,
                                                                                            return_l2dict=True,
                                                                                            n_parallel=n_parallel)
    if refine == 'all':
        refine_inds = 'all'
    elif refine == 'bp':
        refine_inds = sk_ch.branch_points_undirected
    elif refine == 'ep':
        refine_inds = sk_ch.end_points_undirected
    elif refine == 'epbp' or refine == 'bpep':
        refine_inds = np.concatenate(
            (sk_ch.end_points_undirected, sk_ch.branch_points_undirected))
    elif refine is None:
        refine_inds = None
    else:
        raise ValueError(
            '"refine" must be one of "all", "bp", "ep", "epbp"/"bpep", or None')

    if root_point is not None:
        root_point_euc = root_point * np.array([root_point_resolution])
    else:
        root_point_euc = None
    sk_l2, missing_ids = refine_chunk_index_skeleton(sk_ch,
                                                     l2dict_r,
                                                     cv=cv,
                                                     refine_inds=refine_inds,
                                                     scale_chunk_index=True,
                                                     root_location=root_point_euc,
                                                     nan_rounds=nan_rounds,
                                                     return_missing_ids=True,
                                                     segmentation_fallback=segmentation_fallback,
                                                     fallback_mip=fallback_mip,
                                                     cache=cache,
                                                     save_to_cache=save_to_cache)

    if collapse_soma and root_point is not None:
        sk_l2 = collapse_pcg_skeleton(sk_l2.vertices[sk_l2.root],
                                      sk_l2,
                                      collapse_radius)

    output = [sk_l2]
    if return_mesh:
        output.append(mesh_ch)
    if return_l2dict:
        output.append((sk_utils.propagate_l2dict(
            sk_l2, l2dict_mesh), l2dict_r))
    if return_l2dict_mesh:
        output.append((l2dict_mesh, l2dict_mesh_r))
    if return_missing_ids:
        output.append(missing_ids)
    if len(output) == 1:
        return output[0]
    else:
        return tuple(output)


def pcg_meshwork(root_id,
                 datastack_name=None,
                 client=None,
                 cv=None,
                 refine='all',
                 root_point=None,
                 root_point_resolution=[4, 4, 40],
                 root_point_search_radius=300,
                 collapse_soma=False,
                 collapse_radius=DEFAULT_COLLAPSE_RADIUS,
                 synapses=None,
                 synapse_table=None,
                 remove_self_synapse=True,
                 invalidation_d=3,
                 segmentation_fallback=True,
                 fallback_mip=2,
                 cache=None,
                 n_parallel=None,
                 ):
    """Generate a meshwork file based on the level 2 graph.

    Parameters
    ----------
    root_id : int
        Root id of an object in the pychunkedgraph. 
    datastack_name : str or None, optional
        Datastack name to use to initialize a client, if none is provided. By default None.
    client : annotationframeworkclient.FrameworkClientFull or None, optional
        Initialized annotationframeworkclient. If None is given, will use the datastack_name to create one. By default None
    cv : cloudvolume.CloudVolume or None, optional
        Initialized cloudvolume. If none is given, the client info will be used to create one. By default None
    refine : 'all', 'ep', 'bp', 'epbp'/'bpep', or None, optional
        Selects how to refine vertex locations by downloading mesh chunks.
        Unrefined vertices are placed in the center of their chunk in euclidean space.
        * 'all' refines all vertex locations. (Default)
        * 'ep' refines end points only
        * 'bp' refines branch points only
        * 'bpep' or 'epbp' refines both branch and end points.
        * None refines no points.
    root_point : array-like or None, optional
        3 element xyz location for the location to set the root in units set by root_point_resolution,
        by default None. If None, a distal tip is selected.
    root_point_resolution : array-like, optional
        Resolution in euclidean space of the root_point, by default [4, 4, 40]
    root_point_search_radius : int, optional
        Distance in euclidean space to look for segmentation when finding the root vertex, by default 300
    collapse_soma : bool, optional,
        If True, collapses vertices within a given radius of the root point into the root vertex, typically to better
        represent primary neurite branches. Requires a specified root_point. Default if False.
    collapse_radius : float, optional
        Max distance in euclidean space for soma collapse. Default is 10,000 nm (10 microns).
    synapses : 'pre', 'post', 'all', or None, optional
        If not None, queries the synapse_table for presynaptic synapses (if 'pre'),  postsynaptic sites (if 'post'), or both (if 'all'). By default None
    synapse_table : str, optional
        Name of the synapse table to query if synapses are requested, by default None
    remove_self_synapse : bool, optional
        If True, filters out synapses whose pre- and postsynaptic root ids are the same neuron, by default True
    invalidation_d : int, optional
        Invalidation radius in hops for the mesh skeletonization along the chunk adjacency graph, by default 3
    cache : str or None, optional
        Filename to a sqlite database with cached lookups for l2 ids. Optional, default is None.
    n_parallel : int, optional
        Number of parallel downloads passed to cloudvolume, by default 1

    Returns
    -------
    meshparty.meshwork.Meshwork
        Meshwork object with skeleton based on the level 2 graph. See documentation for details.
    """
    if client is None:
        client = FrameworkClient(datastack_name)
    if cv is None:
        cv = cloudvolume.CloudVolume(client.info.segmentation_source(), parallel=n_parallel,
                                     use_https=True, progress=False, bounded=False)

    sk_l2, mesh_chunk, (l2dict_mesh, l2dict_mesh_r) = pcg_skeleton(root_id,
                                                                   client=client,
                                                                   cv=cv,
                                                                   root_point=root_point,
                                                                   root_point_resolution=root_point_resolution,
                                                                   root_point_search_radius=root_point_search_radius,
                                                                   collapse_soma=collapse_soma,
                                                                   collapse_radius=collapse_radius,
                                                                   refine=refine,
                                                                   invalidation_d=invalidation_d,
                                                                   n_parallel=n_parallel,
                                                                   return_mesh=True,
                                                                   return_l2dict_mesh=True,
                                                                   segmentation_fallback=segmentation_fallback,
                                                                   fallback_mip=fallback_mip,
                                                                   cache=cache)

    nrn = meshwork.Meshwork(mesh_chunk, seg_id=root_id, skeleton=sk_l2)

    if synapses is not None and synapse_table is not None:
        if synapses == 'pre':
            pre, post = True, False
        elif synapses == 'post':
            pre, post = False, True
        elif synapses == 'all':
            pre, post = True, True
        else:
            raise ValueError(
                'Synapses must be one of "pre", "post", or "all".')

        pre_syn_df, post_syn_df = pcg_anno.get_level2_synapses(root_id,
                                                               l2dict_mesh,
                                                               client,
                                                               synapse_table,
                                                               remove_self=remove_self_synapse,
                                                               pre=pre,
                                                               post=post,
                                                               )
        if pre_syn_df is not None:
            nrn.anno.add_annotations(
                'pre_syn', pre_syn_df, index_column='pre_pt_mesh_ind', point_column = 'ctr_pt_position')
        if post_syn_df is not None:
            nrn.anno.add_annotations(
                'post_syn', post_syn_df, index_column='post_pt_mesh_ind', point_column = 'ctr_pt_position')

    lvl2_df = pd.DataFrame({'lvl2_id': list(l2dict_mesh.keys()),
                            'mesh_ind': list(l2dict_mesh.values())})
    nrn.anno.add_annotations('lvl2_ids', lvl2_df, index_column='mesh_ind')

    _adjust_meshwork(nrn, cv)

    return nrn


def _adjust_meshwork(nrn, cv):
    """Transform vertices in chunk index space to euclidean"""
    nrn._mesh.vertices = utils.chunk_to_nm(nrn._mesh.vertices, cv)


def collapse_pcg_skeleton(soma_pt, sk, soma_r):
    """Use soma point vertex and collapse soma as sphere
    Parameters
    ----------
    soma_pt : array
        3-element location of soma center (in nm)
    sk: skeleton.Skeleton
        Coarse skeleton
    soma_r : float
        Soma collapse radius (in nm)
    Returns
    -------
    skeleton
        New skeleton with updated properties
    """
    soma_verts, _ = skeletonize.soma_via_sphere(
        soma_pt, sk.vertices, sk.edges, soma_r)
    min_soma_vert = np.argmin(np.linalg.norm(
        sk.vertices[soma_verts] - soma_pt, axis=1))
    root_vert = soma_verts[min_soma_vert]

    new_v, new_e, new_skel_map, vert_filter, root_ind = skeletonize.collapse_soma_skeleton(soma_verts[soma_verts != root_vert],
                                                                                           soma_pt,
                                                                                           sk.vertices,
                                                                                           sk.edges,
                                                                                           sk.mesh_to_skel_map,
                                                                                           collapse_index=root_vert,
                                                                                           return_soma_ind=True,
                                                                                           return_filter=True)

    new_mesh_index = sk.mesh_index[vert_filter]
    new_skeleton = skeleton.Skeleton(
        new_v, new_e, root=root_ind, mesh_to_skel_map=new_skel_map, mesh_index=new_mesh_index, remove_zero_length_edges=False)
    return new_skeleton
