from __future__ import print_function

import os

import ipyvolume as ipv
import nibabel as nib
import numpy as np
from ipywidgets import fixed, interact, widgets


def length(x):
    """Returns the sum of euclidean distances between neighboring points"""
    return np.sum(
        np.sqrt(
            np.sum((x[:-1, :] - x[1:, :]) * (x[:-1, :] - x[1:, :]), axis=1)
        )
    )


def color(x):
    """Returns an approximation for color of line based on its endpoints"""
    dirvec = x[0, :] - x[-1, :]
    return (dirvec / (np.sqrt(np.sum(dirvec * dirvec, axis=-1)))).dot(
        np.eye(3)
    )


class StreamlineWidget:
    """
    Turns nibabel track files into interactive plots using ipyvolume.

    Color for each line is rendered as a function of the endpoints of the
    streamline. Currently, this resolves to red for left-right, green for
    anterior-posterior, and blue for inferior-superior.

    Args
    ----
        filename : str, pathlib.Path
                The path to your ``.trk`` file. Can be a string, or a
                ``PosixPath`` from python3's pathlib.
        streamlines : a nibabel streamline object
                An streamlines attribute of an object loaded by
                nibabel.streamlines.load
    """

    def __init__(self, filename=None, streamlines=None):

        if filename:
            filename = str(filename)
            if not os.path.isfile(filename):
                # for Python3 should have FileNotFoundError here
                raise IOError("file {} not found".format(filename))

            if not nib.streamlines.is_supported(filename):
                raise ValueError(
                    (
                        "File {0} is not a streamline file supported"
                        " by nibabel"
                    ).format(filename)
                )

            # load data in advance
            self.streamlines = nib.streamlines.load(filename).streamlines
        elif streamlines:
            self.streamlines = streamlines
        else:
            raise ValueError(
                "One of filename or streamlines must be specified"
            )
        self.lines2use_ = None

    def plot(self, display_fraction=0.1, **kwargs):
        """
        This is the main method for this widget.

        Args
        ----
            display_fraction : float
                    The fraction of streamlines to show
            percentile : int
                    The initial number of streamlines to show using a
                    percentile of length distribution
            width : int
                    The width of the figure
            height : int
                    The height of the figure
        """

        if display_fraction is not None and (
            display_fraction > 1 or display_fraction <= 0
        ):
            raise ValueError(
                "proportion_to_display is a float between 0 and 1"
                " (0 excluded) or None"
            )

        N = len(self.streamlines)
        num_streamlines = int(display_fraction * N)
        indices = np.random.permutation(N)[:num_streamlines]
        self.lines2use_ = self.streamlines[indices]
        self._default_plotter(**kwargs)

    def _create_mesh(self, indices2use=None):
        if indices2use is None:
            lines2use = self.lines2use_
            local_colors = self.colors
        else:
            lines2use = self.lines2use_[indices2use]
            local_colors = self.colors[indices2use]
        x, y, z = np.concatenate(lines2use).T

        # will contain indices to the verties, [0, 1, 1, 2, 2, 3, 3, 4, 4, 5..]
        indices = np.zeros(
            np.sum((len(line) - 1) * 2 for line in lines2use), dtype=np.uint32
        )
        colors = np.zeros((len(x), 3), dtype=np.float32)

        vertex_offset = 0
        line_offset = 0
        line_pointers = []
        # if we have a line of 4 vertices, we need to add the indices:
        #  offset + [0, 1, 1, 2, 2, 3]
        # so we have approx 2x the number of indices compared to vertices
        for idx, line in enumerate(lines2use):
            line_length = len(line)
            # repeat all but the start and end vertex
            line_indices = np.repeat(
                np.arange(
                    vertex_offset,
                    vertex_offset + line_length,
                    dtype=indices.dtype,
                ),
                2,
            )[1:-1]
            indices[
                line_offset : line_offset + line_length * 2 - 2
            ] = line_indices
            line_pointers.append([line_offset, line_length, line_indices])
            colors[vertex_offset : vertex_offset + line_length] = local_colors[
                idx
            ]
            line_offset += line_length * 2 - 2
            vertex_offset += line_length
        return x, y, z, indices, colors, line_pointers

    def _default_plotter(self, **kwargs):
        """
        Basic plot function to be used if no custom function is specified.

        This is called by plot, you shouldn't call it directly.
        """

        self.lengths = np.array([length(x) for x in self.lines2use_])
        if not ("grayscale" in kwargs and kwargs["grayscale"]):
            self.colors = np.array([color(x) for x in self.lines2use_])
        else:
            self.colors = np.zeros((len(self.lines2use_), 3), dtype=np.float16)
            self.colors[:] = [0.5, 0.5, 0.5]
        self.state = {"threshold": 0, "indices": []}

        width = 600
        height = 600
        perc = 80
        if "width" in kwargs:
            width = kwargs["width"]
        if "height" in kwargs:
            height = kwargs["height"]
        if "percentile" in kwargs:
            perc = kwargs["percentile"]

        ipv.clear()
        fig = ipv.figure(width=width, height=height)
        self.state["fig"] = fig

        with fig.hold_sync():
            x, y, z, indices, colors, self.line_pointers = self._create_mesh()
            limits = np.array(
                [
                    min([x.min(), y.min(), z.min()]),
                    max([x.max(), y.max(), z.max()]),
                ]
            )
            mesh = ipv.Mesh(x=x, y=y, z=z, lines=indices, color=colors)
            fig.meshes = [mesh]
            if "style" not in kwargs:
                fig.style = {
                    "axes": {
                        "color": "black",
                        "label": {"color": "black"},
                        "ticklabel": {"color": "black"},
                        "visible": False,
                    },
                    "background-color": "white",
                    "box": {"visible": False},
                }
            else:
                fig.style = kwargs["style"]
            ipv.pylab._grow_limits(limits, limits, limits)
            fig.camera_fov = 1
        ipv.show()

        interact(
            self._plot_lines,
            state=fixed(self.state),
            threshold=widgets.FloatSlider(
                value=np.percentile(self.lengths, perc),
                min=self.lengths.min() - 1,
                max=self.lengths.max() - 1,
                continuous_update=False,
            ),
        )

    def _plot_lines(self, state, threshold):
        """
        Plots streamlines

        This function is called by _default_plotter
        """
        if threshold < state["threshold"]:
            # when threshold is reduced, increase the number of lines
            state["indices"] = np.where(self.lengths > threshold)[0]
            with state["fig"].hold_sync():
                mesh = state["fig"].meshes[0]
                copy = mesh.lines.copy()
                for idx in state["indices"]:
                    (
                        line_offset,
                        line_length,
                        line_indices,
                    ) = self.line_pointers[idx]
                    copy[
                        line_offset : line_offset + line_length * 2 - 2
                    ] = line_indices
                mesh.lines = copy
                mesh.send_state("lines")
        else:
            # when threshold is increased, decrease the number of lines
            indices = np.where(self.lengths <= threshold)[0]
            with state["fig"].hold_sync():
                mesh = state["fig"].meshes[0]
                copy = mesh.lines.copy()
                for idx in indices:
                    (
                        line_offset,
                        line_length,
                        line_indices,
                    ) = self.line_pointers[idx]
                    copy[line_offset : line_offset + line_length * 2 - 2] = 0
                mesh.lines = copy
                mesh.send_state("lines")
            state["indices"] = np.where(self.lengths > threshold)[0]
        state["threshold"] = threshold
