templates

class TemplateBrain:
    """Generic base class for template brains.

    Minimally, a template should have a `name` and `label` property. For
    mirroring, it also needs a `boundingbox`.

    See [flybrains](https://github.com/navis-org/navis-flybrains) for
    an example of how to use template brains.

    """
    def __init__(self, **properties):
        """Initialize class."""
        for k, v in properties.items():
            setattr(self, k, v)

    @property
    def mesh(self):
        """Mesh represenation of this brain."""
        if not hasattr(self, '_mesh'):
            name = getattr(self, 'regName', getattr(self, 'name', None))
            raise ValueError(f'{name} does not appear to have a mesh')
        return self._mesh

class TemplateRegistry:
    """Tracks template brains, available transforms and produces bridging sequences.

    Parameters
    ----------
    scan_paths :    bool
                    If True will scan paths on initialization.

    """
    def __init__(self, scan_paths: bool = True):
        # Paths to scan for transforms
        self._transpaths = _OS_TRANSPATHS.copy()
        # Transforms
        self._transforms = []
        # Template brains
        self._templates = []

        if scan_paths:
            self.scan_paths()

    def __contains__(self, other) -> bool:
        """Check if transform is in registry.

        Parameters
        ----------
        other :     transform, filepath, tuple
                    Either a transform (e.g. CMTKtransform), a filepath (e.g.
                    to a .list file) or a tuple of `(source, target, transform)`
                    where `transform` can be a transform or a filepath.

        """
        if isinstance(other, (tuple, list)):
            return any([t == other for t in self.transforms])
        else:
            return other in [t.transform for t in self.transforms]

    def __len__(self) -> int:
        return len(self.transforms)

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return f'TemplateRegistry with {len(self)} transforms'

    @property
    def transpaths(self) -> list:
        """Paths searched for transforms.

        Use `.scan_paths` to trigger a scan. Use `.register_path` to add
        more path(s).
        """
        return self._transpaths

    @property
    def templates(self) -> list:
        """Registered template (brains)."""
        return self._templates

    @property
    def transforms(self) -> list:
        """Registered transforms (bridging + mirror)."""
        return self._transforms

    @property
    def bridges(self) -> list:
        """Registered bridging transforms."""
        return [t for t in self.transforms if t.type == 'bridging']

    @property
    def mirrors(self) ->list:
        """Registered mirror transforms."""
        return [t for t in self.transforms if t.type == 'mirror']

    def clear_caches(self):
        """Clear caches of all cached functions."""
        self.bridging_graph.cache_clear()
        self.shortest_bridging_seq.cache_clear()

    def summary(self) -> pd.DataFrame:
        """Generate summary of available transforms."""
        return pd.DataFrame(self.transforms)

    def register_path(self, paths: str, trigger_scan: bool = True):
        """Register path(s) to scan for transforms.

        Parameters
        ----------
        paths :         str | list thereof
                        Paths (or list thereof) to scans for transforms. This
                        is not permanent. For permanent additions set path(s)
                        via the `NAVIS_TRANSFORMS` environment variable.
        trigger_scan :  bool
                        If True, a re-scan of all paths will be triggered.

        """
        paths = utils.make_iterable(paths)

        for p in paths:
            # Try not to duplicate paths
            if p not in self.transpaths:
                self._transpaths.append(p)

        if trigger_scan:
            self.scan_paths()

    def register_templatebrain(self, template: 'TemplateBrain',
                               skip_existing=True):
        """Register a template brain.

        This is used, for example, by navis.mirror_brain.

        Parameters
        ----------
        template :      TemplateBrain
                        TemplateBrain to register.
        skip_existing : bool
                        If True, will skip existing template brains.

        """
        utils.eval_param(template,
                         name='template',
                         allowed_types=(TemplateBrain, ))

        if template not in self._templates or not skip_existing:
            self._templates.append(template)

    def register_transform(self, transform: BaseTransform, source: str,
                           target: str, transform_type: str,
                           skip_existing: bool = True,
                           weight: int = 1):
        """Register a transform.

        Parameters
        ----------
        transform :         subclass of BaseTransform | TransformSequence
                            A transform (AffineTransform, CMTKtransform, etc.)
                            or a TransformSequence.
        source :            str
                            Source for forward transform.
        target :            str
                            Target for forward transform. Ignored for mirror
                            transforms.
        transform_type :    "bridging" | "mirror"
                            Type of transform.
        skip_existing :     bool
                            If True will skip if transform is already in registry.
        weight :            int
                            Giving a transform a lower weight will make it
                            preferable when plotting bridging sequences.

        See Also
        --------
        register_transformfile
                            If you want to register a file instead of an
                            already constructed transform.

        """
        assert transform_type in ('bridging', 'mirror')
        assert isinstance(transform, (BaseTransform, TransformSequence))

        # Translate into edge
        edge = transform_reg(source=source, target=target, transform=transform,
                             type=transform_type,
                             invertible=hasattr(transform, '__neg__'),
                             weight=weight)

        # Don't add if already exists
        if not skip_existing or edge not in self:
            self.transforms.append(edge)

        # Clear cached functions
        self.clear_caches()

    def register_transformfile(self, path: str, **kwargs):
        """Parse and register a transform file.

        File/Directory name must follow the a `{TARGET}_{SOURCE}.{ext}`
        convention (e.g. `JRC2013_FCWB.list`).

        Parameters
        ----------
        path :          str
                        Path to transform.
        **kwargs
                        Keyword arguments are passed to the constructor of the
                        Transform (e.g. CMTKtransform for `.list` directory).

        See Also
        --------
        register_transform
                        If you want to register an already constructed transform
                        instead of a transform file that still needs to be
                        parsed.

        """
        assert isinstance(path, (str, pathlib.Path))

        path = pathlib.Path(path).expanduser()

        if not path.is_dir() and not path.is_file():
            raise ValueError(f'File/directory "{path}" does not exist')

        # Parse properties
        try:
            if 'mirror' in path.name or 'imgflip' in path.name:
                transform_type = 'mirror'
                source = path.name.split('_')[0]
                target = None
            else:
                transform_type = 'bridging'
                target = path.name.split('_')[0]
                source = path.name.split('_')[1].split('.')[0]

            # Initialize the transform
            transform = factory.factory_methods[path.suffix](path, **kwargs)

            self.register_transform(transform=transform,
                                    source=source,
                                    target=target,
                                    transform_type=transform_type)
        except BaseException as e:
            logger.error(f'Error registering {path} as transform: {str(e)}')

    def scan_paths(self, extra_paths: List[str] = None):
        """Scan registered paths for transforms and add to registry.

        Will skip transforms that already exist in this registry.

        Parameters
        ----------
        extra_paths :   list of str
                        Any Extra paths to search.

        """
        search_paths = self.transpaths

        if isinstance(extra_paths, str):
            extra_paths = [i for i in extra_paths.split(';') if len(i) > 0]
            search_paths = np.append(search_paths, extra_paths)

        for path in search_paths:
            path = pathlib.Path(path).expanduser()
            # Skip if path does not exist
            if not path.is_dir():
                continue

            # Go over the file extensions we can work with (.h5, .list, .json)
            # These file extensions are registered in the
            # `navis.transforms.factory` module
            for ext in factory.factory_methods:
                for hit in path.rglob(f'*{ext}'):
                    if hit.is_dir() or hit.is_file():
                        # Register this file
                        self.register_transformfile(hit)

        # Clear cached functions
        self.clear_caches()

    @functools.lru_cache()
    def bridging_graph(self,
                       reciprocal: Union[Literal[False], int, float] = True) -> nx.DiGraph:
        """Generate networkx Graph describing the bridging paths.

        Parameters
        ----------
        reciprocal :    bool | float
                        If True or float, will add forward and inverse edges for
                        transforms that are invertible. If float, the inverse
                        edges' weights will be scaled by that factor.

        Returns
        -------
        networkx.MultiDiGraph

        """
        # Drop mirror transforms
        bridge = [t for t in self.transforms if t.type == 'bridging']
        bridge_inv = [t for t in bridge if t.invertible]

        # Generate graph
        # Note we are using MultiDi graph here because we might
        # have multiple edges between nodes. For example, there
        # is a JFRC2013DS_JFRC2013 and a JFRC2013_JFRC2013DS
        # bridging registration. If we include the inverse, there
        # will be two edges connecting JFRC2013DS and JFRC2013 in
        # both directions
        G = nx.MultiDiGraph()
        edges = [(t.source, t.target,
                  {'transform': t.transform,
                   'type': type(t.transform).__name__,
                   'weight': t.weight}) for t in bridge]

        if reciprocal:
            if isinstance(reciprocal, numbers.Number):
                rv_edges = [(t.target, t.source,
                             {'transform': -t.transform,  # note inverse transform!
                              'type': str(type(t.transform)).split('.')[-1],
                              'weight': t.weight * reciprocal}) for t in bridge_inv]
            else:
                rv_edges = [(t.target, t.source,
                             {'transform': -t.transform,  # note inverse transform!
                              'type': str(type(t.transform)).split('.')[-1],
                              'weight': t.weight}) for t in bridge_inv]
            edges += rv_edges

        G.add_edges_from(edges)

        return G

    def find_bridging_path(self, source: str,
                           target: str,
                           via: Optional[str] = None,
                           avoid: Optional[str] = None,
                           reciprocal=True) -> tuple:
        """Find bridging path from source to target.

        Parameters
        ----------
        source :        str
                        Source from which to transform to `target`.
        target :        str
                        Target to which to transform to.
        via :           str | list thereof, optional
                        Force specific intermediate template(s).
        avoid :         str | list thereof, optional
                        Avoid going through specific intermediate template(s).
        reciprocal :    bool | float
                        If True or float, will add forward and inverse edges for
                        transforms that are invertible. If float, the inverse
                        edges' weights will be scaled by that factor.

        Returns
        -------
        path :          list
                        Path from source to target: [source, ..., target]
        transforms :    list
                        Transforms as [[path_to_transform, inverse], ...]

        """
        # Generate (or get cached) bridging graph
        G = self.bridging_graph(reciprocal=reciprocal)

        if len(G) == 0:
            raise ValueError('No bridging registrations available')

        # Do not remove the conversion to list - fuzzy matching does act up
        # otherwise
        nodes = list(G.nodes)
        if source not in nodes:
            best_match = fw.process.extractOne(source, nodes,
                                               scorer=fw.fuzz.token_sort_ratio)
            raise ValueError(f'Source "{source}" has no known bridging '
                             f'registrations. Did you mean "{best_match[0]}" '
                             'instead?')
        if target not in G.nodes:
            best_match = fw.process.extractOne(target, nodes,
                                               scorer=fw.fuzz.token_sort_ratio)
            raise ValueError(f'Target "{target}" has no known bridging '
                             f'registrations. Did you mean "{best_match[0]}" '
                             'instead?')

        if via:
            via = list(utils.make_iterable(via))  # do not remove the list() here
            for v in via:
                if v not in G.nodes:
                    best_match = fw.process.extractOne(v, nodes,
                                                       scorer=fw.fuzz.token_sort_ratio)
                    raise ValueError(f'Via "{v}" has no known bridging '
                                    f'registrations. Did you mean "{best_match[0]}" '
                                    'instead?')

        if avoid:
            avoid = list(utils.make_iterable(avoid))

        # This will raise a error message if no path is found
        if not via and not avoid:
            try:
                path = nx.shortest_path(G, source, target, weight='weight')
            except nx.NetworkXNoPath:
                raise nx.NetworkXNoPath(f'No bridging path connecting {source} '
                                        f'and {target} found.')
        else:
            # Go through all possible paths and find one that...
            found_any = False  # track if we found any path
            found_good = False  # track if we found a path matching the criteria
            for path in nx.all_simple_paths(G, source, target):
                found_any = True
                # ... has all `via`s...
                if via and all([v in path for v in via]):
                    # ... and none of the `avoid`
                    if avoid:
                        if not any([v in path for v in avoid]):
                            found_good = True
                            break
                    else:
                        found_good = True
                        break
                # If we only have `avoid` but no `via`
                elif avoid and not any([v in path for v in avoid]):
                    found_good = True
                    break

            if not found_any:
                raise nx.NetworkXNoPath(f'No bridging path connecting {source} '
                                        f'and {target} found.')
            elif not found_good:
                if via and avoid:
                    raise nx.NetworkXNoPath(f'No bridging path connecting {source}'
                                            f'and {target} via "{via}" and '
                                            f'avoiding "{avoid}" found')
                elif via:
                    raise nx.NetworkXNoPath(f'No bridging path connecting {source}'
                                            f'and {target} via "{via}" found.')
                else:
                    raise nx.NetworkXNoPath(f'No bridging path connecting {source}'
                                            f'and {target} avoiding "{avoid}" found.')

        # `path` holds the sequence of nodes we are traversing but not which
        # transforms (i.e. edges) to use
        transforms = []
        for n1, n2 in zip(path[:-1], path[1:]):
            this_edges = []
            i = 0
            # First collect all edges between those two nodes
            # - this is annoyingly complicated with MultiDiGraphs
            while True:
                try:
                    e = G.edges[(n1, n2, i)]
                except KeyError:
                    break
                this_edges.append([e['transform'], e['weight']])
                i += 1

            # Now find the edge with the highest weight
            # (inverse transforms might have a lower weight)
            this_edges = sorted(this_edges, key=lambda x: x[-1])
            transforms.append(this_edges[-1][0])

        return path, transforms

    def find_all_bridging_paths(self, source: str,
                                target: str,
                                via: Optional[str] = None,
                                avoid: Optional[str] = None,
                                reciprocal: bool = True,
                                cutoff: int = None) -> tuple:
        """Find all bridging paths from source to target.

        Parameters
        ----------
        source :        str
                        Source from which to transform to `target`.
        target :        str
                        Target to which to transform to.
        via :           str | list thereof, optional
                        Force specific intermediate template(s).
        avoid :         str | list thereof, optional
                        Avoid specific intermediate template(s).
        reciprocal :    bool | float
                        If True or float, will add forward and inverse edges for
                        transforms that are invertible. If float, the inverse
                        edges' weights will be scaled by that factor.
        cutoff :        int, optional
                        Depth to stop the search. Only paths of length
                        <= cutoff are returned.

        Returns
        -------

        path :          list
                        Path from source to target: [source, ..., target]
        transforms :    list
                        Transforms as [[path_to_transform, inverse], ...]

        """
        # Generate (or get cached) bridging graph
        G = self.bridging_graph(reciprocal=reciprocal)

        if len(G) == 0:
            raise ValueError('No bridging registrations available')

        # Do not remove the conversion to list - fuzzy matching does act up
        # otherwise
        nodes = list(G.nodes)
        if source not in nodes:
            best_match = fw.process.extractOne(source, nodes,
                                               scorer=fw.fuzz.token_sort_ratio)
            raise ValueError(f'Source "{source}" has no known bridging '
                             f'registrations. Did you mean "{best_match[0]}" '
                             'instead?')
        if target not in G.nodes:
            best_match = fw.process.extractOne(target, nodes,
                                               scorer=fw.fuzz.token_sort_ratio)
            raise ValueError(f'Target "{target}" has no known bridging '
                             f'registrations. Did you mean "{best_match[0]}" '
                             'instead?')

        if via and via not in G.nodes:
            best_match = fw.process.extractOne(via, nodes,
                                               scorer=fw.fuzz.token_sort_ratio)
            raise ValueError(f'Via "{via}" has no known bridging '
                             f'registrations. Did you mean "{best_match[0]}" '
                             'instead?')

        # This will raise a error message if no path is found
        for path in nx.all_simple_paths(G, source, target, cutoff=cutoff):
            # Skip paths that don't contain `via`
            if isinstance(via, str) and (via not in path):
                continue
            elif isinstance(via, (list, tuple, np.ndarray)) and not all([v in path for v in via]):
                continue

            # Skip paths that contain `avoid`
            if isinstance(avoid, str) and (avoid in path):
                continue
            elif isinstance(avoid, (list, tuple, np.ndarray)) and any([v in path for v in avoid]):
                continue

            # `path` holds the sequence of nodes we are traversing but not which
            # transforms (i.e. edges) to use
            transforms = []
            for n1, n2 in zip(path[:-1], path[1:]):
                this_edges = []
                i = 0
                # First collect all edges between those two nodes
                # - this is annoyingly complicated with MultiDiGraphs
                while True:
                    try:
                        e = G.edges[(n1, n2, i)]
                    except KeyError:
                        break
                    this_edges.append([e['transform'], e['weight']])
                    i += 1

                # Now find the edge with the highest weight
                # (inverse transforms might have a lower weight)
                this_edges = sorted(this_edges, key=lambda x: x[-1])
                transforms.append(this_edges[-1][0])

            yield path, transforms

    @functools.lru_cache()
    def shortest_bridging_seq(self, source: str, target: str,
                              via: Optional[str] = None,
                              inverse_weight: float = .5) -> tuple:
        """Find shortest bridging sequence to get from source to target.

        Parameters
        ----------
        source :            str
                            Source from which to transform to `target`.
        target :            str
                            Target to which to transform to.
        via :               str | list of str
                            Waystations to traverse on the way from source to
                            target.
        inverse_weight :    float
                            Weight for inverse transforms. If < 1 will prefer
                            forward transforms.

        Returns
        -------
        sequence :          (N, ) array
                            Sequence of registrations that will be traversed.
        transform_seq :     TransformSequence
                            Class that collates the required transforms to get
                            from source to target.

        """
        # Generate sequence of nodes we need to find a path for
        # Minimally it's just from source to target
        nodes = np.array([source, target])

        if via:
            nodes = np.insert(nodes, 1, via)

        seq = [nodes[0]]
        transforms = []
        for n1, n2 in zip(nodes[:-1], nodes[1:]):
            path, tr = self.find_bridging_path(n1, n2, reciprocal=inverse_weight)
            seq = np.append(seq, path[1:])
            transforms = np.append(transforms, tr)

        if any(np.unique(seq, return_counts=True)[1] > 1):
            logger.warning('Bridging sequence contains loop: '
                           f'{"->".join(seq)}')

        # Generate the transform sequence
        transform_seq = TransformSequence(*transforms)

        return seq, transform_seq

    def find_mirror_reg(self, template: str, non_found: str = 'raise') -> tuple:
        """Search for a mirror transformation for given template.

        Typically a mirror transformation specifies a non-rigid transformation
        to correct asymmetries in an image.

        Parameters
        ----------
        template :  str
                    Name of the template to find a mirror transformation for.
        non_found : "raise" | "ignore"
                    What to do if no mirror transformation is found. If "ignore"
                    and no mirror transformation found, will silently return
                    `None`.

        Returns
        -------
        tuple
                    Named tuple containing a mirror transformation. Will only
                    ever return one - even if multiple are available.

        """
        for tr in self.mirrors:
            if tr.source == template:
                return tr

        if non_found == 'raise':
            raise ValueError(f'No mirror transformation found for {template}')
        return None

    def find_closest_mirror_reg(self, template: str, non_found: str = 'raise') -> str:
        """Search for the closest mirror transformation for given template.

        Typically a mirror transformation specifies a non-rigid transformation
        to correct asymmetries in an image.

        Parameters
        ----------
        template :  str
                    Name of the template to find a mirror transformation for.
        non_found : "raise" | "ignore"
                    What to do if there is no path to a mirror transformation.
                    If "ignore" and no path is found, will silently return
                    `None`.

        Returns
        -------
        str
                    Name of the closest template with a mirror transform.

        """
        # Templates with mirror registrations
        temps_w_mirrors = [t.source for t in self.mirrors]

        # Add symmetrical template brains
        temps_w_mirrors += [t.label for t in self.templates if getattr(t, 'symmetrical', False) == True]

        if not temps_w_mirrors:
            raise ValueError('No mirror transformations registered')

        # If this template has a mirror registration:
        if template in temps_w_mirrors:
            return template

        # Get bridging graph
        G = self.bridging_graph()

        if template not in G.nodes:
            raise ValueError(f'"{template}" does not appear to be a registered '
                             'template')

        # Get path lengths from template to all other nodes
        pl = nx.single_source_dijkstra_path_length(G, template)

        # Subset to targets that have a mirror reg
        pl = {k: v for k, v in pl.items() if k in temps_w_mirrors}

        # Find the closest mirror
        cl = sorted(pl.keys(), key=lambda x: pl[x])

        # If any, return the closests
        if cl:
            return cl[0]

        if non_found == 'raise':
            raise ValueError(f'No path to a mirror transformation found for "{template}"')

        return None

    def find_template(self, name: str, non_found: str = 'raise') -> 'TemplateBrain':
        """Search for a given template (brain).

        Parameters
        ----------
        name :      str
                    Name of the template to find a mirror transformation for.
                    Searches against `name` and `label` (short name) properties
                    of registered templates.
        non_found : "raise" | "ignore"
                    What to do if no mirror transformation is found. If "ignore"
                    and no mirror transformation found, will silently return
                    `None`.

        Returns
        -------
        TemplateBrain

        """
        for tmp in self.templates:
            if getattr(tmp, 'label', None) == name:
                return tmp
            if getattr(tmp, 'name', None) == name:
                return tmp

        if non_found == 'raise':
            raise ValueError(f'No template brain registered that matches "{name}"')
        return None

    def plot_bridging_graph(self, **kwargs):
        """Draw bridging graph using networkX.

        Parameters
        ----------
        **kwargs
                    Keyword arguments are passed to `networkx.draw_networkx`.

        Returns
        -------
        None

        """
        # Get graph
        G = self.bridging_graph(reciprocal=False)

        # Draw nodes and edges
        node_labels = {n: n for n in G.nodes}
        pos = nx.kamada_kawai_layout(G)

        # Draw all nodes
        nx.draw_networkx_nodes(G, pos=pos, node_color='lightgrey',
                               node_shape='o', node_size=300)
        nx.draw_networkx_labels(G, pos=pos, labels=node_labels,
                                font_color='k', font_size=10)

        # Draw edges by type of transform
        edge_types = set([e[2]['type'] for e in G.edges(data=True)])

        lines = []
        labels = []
        for t, c in zip(edge_types,
                        sns.color_palette('muted', len(edge_types))):
            subset = [e for e in G.edges(data=True) if e[2]['type'] == t]
            nx.draw_networkx_edges(G, pos=pos, edgelist=subset,
                                   edge_color=mcl.to_hex(c), width=1.5)
            lines.append(Line2D([0], [0], color=c, linewidth=2, linestyle='-'))
            labels.append(t)

        plt.legend(lines, labels)