# Copyright (c) 2015, Ecole Polytechnique Federale de Lausanne, Blue Brain Project
# All rights reserved.
#
# This file is part of NeuroM <https://github.com/BlueBrain/NeuroM>
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#     1. Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#     2. Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#     3. Neither the name of the copyright holder nor the names of
#        its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

'''
Python module of NeuroM to visualize morphologies
'''

from itertools import izip
from neurom.view import common
from neurom.core.types import NeuriteType
from matplotlib.collections import LineCollection
import numpy as np
from neurom.io import COLS
from neurom import _compat
from neurom.analysis.morphmath import segment_radius

DEFAULT_PARAMS = '''        new_fig: boolean \
            Defines if the tree will be plotted \
            in the current figure (False) \
            or in a new figure (True) \
            Default value is True.
        subplot: matplotlib subplot value or False \
            If False the default subplot 111 will be used. \
            For any other value a matplotlib subplot \
            will be generated. \
            Default value is False.
        treecolor: str or None \
            Defines the color of the tree. \
            If None the default values will be used, \
            depending on the type of tree: \
            Basal dendrite: "red" \
            Axon : "blue" \
            Apical dendrite: "purple" \
            Undefined tree: "black" \
            Default value is None.
        linewidth: float \
            Defines the linewidth of the tree, \
            if diameter is set to False. \
            Default value is 1.2.
        alpha: float \
            Defines throughe transparency of the tree. \
            0.0 transparent through 1.0 opaque. \
            Default value is 0.8.

''' + common.PLOT_STYLE_PARAMS


def get_default(variable, **kwargs):
    '''
    Returns default variable or kwargs variable if it exists.
    '''
    default = {'linewidth': 1.2,
               'alpha': 0.8,
               'treecolor': None,
               'diameter': True,
               'diameter_scale': 1.0,
               'white_space': 30.}

    return kwargs.get(variable, default[variable])


def tree(tr, plane='xy', new_fig=True, subplot=False, **kwargs):
    '''
    Generates a 2d figure of the tree's segments. \
    If the tree contains one single point the plot will be empty \
    since no segments can be constructed.

    Parameters:
        tr: Tree \
            neurom.Tree object
        plane: str \
            Accepted values: Any pair of of xyz \
            Default value is 'xy'.treecolor
        diameter: boolean
            If True the diameter, scaled with diameter_scale factor, \
            will define the width of the tree lines. \
            If False use linewidth to select the width of the tree lines. \
            Default value is True.
        diameter_scale: float \
            Defines the scale factor that will be multiplied \
            with the diameter to define the width of the tree line. \
            Default value is 1.
        white_space: float \
            Defines the white space around \
            the boundary box of the morphology. \
            Default value is 1.
    '''
    if plane not in ('xy', 'yx', 'xz', 'zx', 'yz', 'zy'):
        return None, 'No such plane found! Please select one of: xy, xz, yx, yz, zx, zy.'

    # Initialization of matplotlib figure and axes.
    fig, ax = common.get_figure(new_fig=new_fig, subplot=subplot)

    # Data needed for the viewer: x,y,z,r
    bounding_box = _compat.bounding_box(tr)

    white_space = get_default('white_space', **kwargs)

    def _seg_2d(seg):
        '''2d coordinates needed for the plotting of a segment'''
        horz = getattr(COLS, plane[0].capitalize())
        vert = getattr(COLS, plane[1].capitalize())
        parent_point = seg[0]
        child_point = seg[1]
        horz1 = parent_point[horz]
        horz2 = child_point[horz]
        vert1 = parent_point[vert]
        vert2 = child_point[vert]
        return ((horz1, vert1), (horz2, vert2))

    segs = _compat.map_segments(tr, _seg_2d)

    linewidth = get_default('linewidth', **kwargs)
    # Definition of the linewidth according to diameter, if diameter is True.
    if get_default('diameter', **kwargs):
        scale = get_default('diameter_scale', **kwargs)
        # TODO: This was originally a numpy array. Did it have to be one?
        linewidth = [2 * r * scale for r in _compat.map_segments(tr, segment_radius)]
        if len(linewidth) == 0:
            linewidth = get_default('linewidth', **kwargs)

    # Plot the collection of lines.
    collection = LineCollection(segs,
                                color=common.get_color(get_default('treecolor', **kwargs),
                                                       _compat.neurite_type(tr)),
                                linewidth=linewidth, alpha=get_default('alpha', **kwargs))

    ax.add_collection(collection)

    kwargs['title'] = kwargs.get('title', 'Tree view')
    kwargs['xlabel'] = kwargs.get('xlabel', plane[0])
    kwargs['ylabel'] = kwargs.get('ylabel', plane[1])
    kwargs['xlim'] = kwargs.get('xlim', [bounding_box[0][getattr(COLS, plane[0].capitalize())] -
                                         white_space,
                                         bounding_box[1][getattr(COLS, plane[0].capitalize())] +
                                         white_space])
    kwargs['ylim'] = kwargs.get('ylim', [bounding_box[0][getattr(COLS, plane[1].capitalize())] -
                                         white_space,
                                         bounding_box[1][getattr(COLS, plane[1].capitalize())] +
                                         white_space])

    return common.plot_style(fig=fig, ax=ax, **kwargs)


def soma(sm, plane='xy', new_fig=True, subplot=False, **kwargs):
    '''
    Generates a 2d figure of the soma.

    Parameters:
        soma: Soma \
        neurom.Soma object
        plane: str \
            Accepted values: Any pair of of xyz \
            Default value is 'xy'.treecolor
    '''
    treecolor = kwargs.get('treecolor', None)
    outline = kwargs.get('outline', True)

    if plane not in ('xy', 'yx', 'xz', 'zx', 'yz', 'zy'):
        return None, 'No such plane found! Please select one of: xy, xz, yx, yz, zx, zy.'

    # Initialization of matplotlib figure and axes.
    fig, ax = common.get_figure(new_fig=new_fig, subplot=subplot)

    # Definition of the tree color depending on the tree type.
    treecolor = common.get_color(treecolor, tree_type=NeuriteType.soma)

    # Plot the outline of the soma as a circle, is outline is selected.
    if not outline:
        soma_circle = common.plt.Circle(sm.center, sm.radius, color=treecolor,
                                        alpha=get_default('alpha', **kwargs))
        ax.add_artist(soma_circle)
    else:
        horz = []
        vert = []

        for s_point in sm.iter():
            horz.append(s_point[getattr(COLS, plane[0].capitalize())])
            vert.append(s_point[getattr(COLS, plane[1].capitalize())])

        horz.append(horz[0]) # To close the loop for a soma viewer. This might be modified!
        vert.append(vert[0]) # To close the loop for a soma viewer. This might be modified!

        common.plt.plot(horz, vert, color=treecolor,
                        alpha=get_default('alpha', **kwargs),
                        linewidth=get_default('linewidth', **kwargs))

    kwargs['title'] = kwargs.get('title', 'Soma view')
    kwargs['xlabel'] = kwargs.get('xlabel', plane[0])
    kwargs['ylabel'] = kwargs.get('ylabel', plane[1])

    return common.plot_style(fig=fig, ax=ax, **kwargs)


def neuron(nrn, plane='xy', new_fig=True, subplot=False, **kwargs):
    '''
    Generates a 2d figure of the neuron, \
    that contains a soma and a list of trees.

    Parameters:
        neuron: Neuron \
        neurom.Neuron object
        plane: str \
            Accepted values: Any pair of of xyz \
            Default value is 'xy'.treecolor
        diameter: boolean
            If True the diameter, scaled with diameter_scale factor, \
            will define the width of the tree lines. \
            If False use linewidth to select the width of the tree lines. \
            Default value is True.
        diameter_scale: float \
            Defines the scale factor that will be multiplied \
            with the diameter to define the width of the tree line. \
            Default value is 1.
    '''
    if plane not in ('xy', 'yx', 'xz', 'zx', 'yz', 'zy'):
        return None, 'No such plane found! Please select one of: xy, xz, yx, yz, zx, zy.'

    # Initialization of matplotlib figure and axes.
    fig, ax = common.get_figure(new_fig=new_fig, subplot=subplot)

    kwargs['new_fig'] = False
    kwargs['subplot'] = subplot

    kwargs['final'] = False

    white_space = get_default('white_space', **kwargs)

    soma(nrn.soma, plane=plane, **kwargs)

    kwargs['title'] = kwargs.get('title', nrn.name)
    kwargs['xlabel'] = kwargs.get('xlabel', plane[0])
    kwargs['ylabel'] = kwargs.get('ylabel', plane[1])

    h = []
    v = []

    for temp_tree in nrn.neurites:

        bounding_box = _compat.bounding_box(temp_tree)

        h.append([bounding_box[0][getattr(COLS, plane[0].capitalize())],
                  bounding_box[1][getattr(COLS, plane[0].capitalize())]])
        v.append([bounding_box[0][getattr(COLS, plane[1].capitalize())],
                  bounding_box[1][getattr(COLS, plane[1].capitalize())]])

        tree(temp_tree, plane=plane, **kwargs)

    if h:
        kwargs['xlim'] = kwargs.get('xlim', [np.min(h) - white_space,
                                             np.max(h) + white_space])
    if v:
        kwargs['ylim'] = kwargs.get('ylim', [np.min(v) - white_space,
                                             np.max(v) + white_space])

    kwargs['final'] = True

    return common.plot_style(fig=fig, ax=ax, **kwargs)


def tree3d(tr, new_fig=True, new_axes=True, subplot=False, **kwargs):
    '''
    Generates a figure of the tree in 3d.
    If the tree contains one single point the plot will be empty \
    since no segments can be constructed.

    Parameters:
        tr: Tree \
        neurom.Tree object
        diameter: boolean \
            If True the diameter, scaled with diameter_scale factor, \
            will define the width of the tree lines. \
            If False use linewidth to select the width of the tree lines. \
            Default value is True.
        diameter_scale: float \
            Defines the scale factor that will be multiplied \
            with the diameter to define the width of the tree line. \
            Default value is 1.
        white_space: float \
            Defines the white space around \
            the boundary box of the morphology. \
            Default value is 1.
    '''
    from mpl_toolkits.mplot3d.art3d import Line3DCollection

    # Initialization of matplotlib figure and axes.
    fig, ax = common.get_figure(new_fig=new_fig, new_axes=new_axes,
                                subplot=subplot, params={'projection': '3d'})

    # Data needed for the viewer: x,y,z,r
    bounding_box = _compat.bounding_box(tr)

    def _seg_3d(seg):
        '''2d coordinates needed for the plotting of a segment'''
        horz = getattr(COLS, 'X')
        vert = getattr(COLS, 'Y')
        depth = getattr(COLS, 'Z')
        parent_point = seg[0]
        child_point = seg[1]
        horz1 = parent_point[horz]
        horz2 = child_point[horz]
        vert1 = parent_point[vert]
        vert2 = child_point[vert]
        depth1 = parent_point[depth]
        depth2 = child_point[depth]
        return ((horz1, vert1, depth1), (horz2, vert2, depth2))

    segs = _compat.map_segments(tr, _seg_3d)

    linewidth = get_default('linewidth', **kwargs)

    # Definition of the linewidth according to diameter, if diameter is True.
    if get_default('diameter', **kwargs):
        # TODO: This was originally a numpy array. Did it have to be one?
        scale = get_default('diameter_scale', **kwargs)
        linewidth = [2 * r * scale for r in _compat.map_segments(tr, segment_radius)]
        if len(linewidth) == 0:
            linewidth = get_default('linewidth', **kwargs)

    # Plot the collection of lines.
    collection = Line3DCollection(segs,
                                  color=common.get_color(get_default('treecolor', **kwargs),
                                                         _compat.neurite_type(tr)),
                                  linewidth=linewidth, alpha=get_default('alpha', **kwargs))

    ax.add_collection3d(collection)

    kwargs['title'] = kwargs.get('title', 'Tree 3d-view')
    kwargs['xlabel'] = kwargs.get('xlabel', 'X')
    kwargs['ylabel'] = kwargs.get('ylabel', 'Y')
    kwargs['zlabel'] = kwargs.get('zlabel', 'Z')

    kwargs['xlim'] = kwargs.get('xlim', [bounding_box[0][0] - get_default('white_space', **kwargs),
                                         bounding_box[1][0] + get_default('white_space', **kwargs)])
    kwargs['ylim'] = kwargs.get('ylim', [bounding_box[0][1] - get_default('white_space', **kwargs),
                                         bounding_box[1][1] + get_default('white_space', **kwargs)])
    kwargs['zlim'] = kwargs.get('zlim', [bounding_box[0][2] - get_default('white_space', **kwargs),
                                         bounding_box[1][2] + get_default('white_space', **kwargs)])

    return common.plot_style(fig=fig, ax=ax, **kwargs)


def soma3d(sm, new_fig=True, new_axes=True, subplot=False, **kwargs):
    '''
    Generates a 3d figure of the soma.

    Parameters:
        soma: Soma \
        neurom.Soma object
    '''
    treecolor = kwargs.get('treecolor', None)

    # Initialization of matplotlib figure and axes.
    fig, ax = common.get_figure(new_fig=new_fig, new_axes=new_axes,
                                subplot=subplot, params={'projection': '3d'})

    # Definition of the tree color depending on the tree type.
    treecolor = common.get_color(treecolor, tree_type=NeuriteType.soma)

    xs = sm.center[0]
    ys = sm.center[1]
    zs = sm.center[2]

    # Plot the soma as a circle.
    fig, ax = common.plot_sphere(fig, ax, center=[xs, ys, zs], radius=sm.radius, color=treecolor,
                                 alpha=get_default('alpha', **kwargs))

    kwargs['title'] = kwargs.get('title', 'Soma view')
    kwargs['xlabel'] = kwargs.get('xlabel', 'X')
    kwargs['ylabel'] = kwargs.get('ylabel', 'Y')
    kwargs['zlabel'] = kwargs.get('zlabel', 'Z')

    return common.plot_style(fig=fig, ax=ax, **kwargs)


def neuron3d(nrn, new_fig=True, new_axes=True, subplot=False, **kwargs):
    '''
    Generates a figure of the neuron,
    that contains a soma and a list of trees.

    Parameters:
        neuron: Neuron \
        neurom.Neuron object
        diameter: boolean
            If True the diameter, scaled with diameter_scale factor, \
            will define the width of the tree lines. \
            If False use linewidth to select the width of the tree lines. \
            Default value is True.
        diameter_scale: float \
            Defines the scale factor that will be multiplied \
            with the diameter to define the width of the tree line. \
            Default value is 1.
    '''
    # Initialization of matplotlib figure and axes.
    fig, ax = common.get_figure(new_fig=new_fig, new_axes=new_axes,
                                subplot=subplot, params={'projection': '3d'})

    white_space = get_default('white_space', **kwargs)

    kwargs['new_fig'] = False
    kwargs['subplot'] = subplot
    kwargs['new_axes'] = False
    kwargs['title'] = kwargs.get('title', nrn.name)

    kwargs['final'] = False

    soma3d(nrn.soma, **kwargs)

    boundaries = [[], [], []]

    for temp_tree in nrn.neurites:

        bounding_box = _compat.bounding_box(temp_tree)

        boundaries[0].append([bounding_box[0][getattr(COLS, 'X')],
                              bounding_box[1][getattr(COLS, 'X')]])
        boundaries[1].append([bounding_box[0][getattr(COLS, 'Y')],
                              bounding_box[1][getattr(COLS, 'Y')]])
        boundaries[2].append([bounding_box[0][getattr(COLS, 'Z')],
                              bounding_box[1][getattr(COLS, 'Z')]])

        tree3d(temp_tree, **kwargs)

    if len(boundaries[0]) > 0:
        kwargs['xlim'] = kwargs.get('xlim', [np.min(boundaries[0]) - white_space,
                                             np.max(boundaries[0]) + white_space])
    if len(boundaries[1]) > 0:
        kwargs['ylim'] = kwargs.get('ylim', [np.min(boundaries[1]) - white_space,
                                             np.max(boundaries[1]) + white_space])
    if len(boundaries[2]) > 0:
        kwargs['zlim'] = kwargs.get('zlim', [np.min(boundaries[2]) - white_space,
                                             np.max(boundaries[2]) + white_space])

    kwargs['final'] = True

    return common.plot_style(fig=fig, ax=ax, **kwargs)


def _format_str(string):
    ''' String formatting
    '''
    return string.replace('NeuriteType.', '').replace('_', ' ').capitalize()


def _generate_collection(group, ax, ctype, colors):
    ''' Render rectangle collection
    '''
    from matplotlib.collections import PolyCollection

    color = common.TREE_COLOR[ctype]

    # generate segment collection
    collection = PolyCollection(group, closed=False, antialiaseds=True,
                                edgecolors='face', facecolors=color)

    # add it to the axes
    ax.add_collection(collection)

    # dummy plot for the legend
    if color not in colors:
        ax.plot((0., 0.), (0., 0.), c=color, label=_format_str(str(ctype)))
        colors.add(color)


def _render_dendrogram(dnd, ax, displacement):
    '''Renders dendrogram
    '''
    # set of unique colors that reflect the set of types of the neurites
    colors = set()

    for n, (indices, ctype) in enumerate(izip(dnd.groups, dnd.types)):

        # slice rectangles array for the current neurite
        group = dnd.data[indices[0]: indices[1]]

        if n > 0:
            # displace the neurites by half of their maximum x dimension
            # plus half of the previous neurite's maxmimum x dimension
            displacement += 0.5 * (dnd.dims[n - 1][0] + dnd.dims[n][0])

        # arrange the trees without overlapping with each other
        group += (displacement, 0.)

        # create the polygonal collection of the dendrogram
        # segments
        _generate_collection(group, ax, ctype, colors)

    soma_square = dnd.soma

    if soma_square is not None:

        _generate_collection((soma_square + (displacement / 2., 0.),), ax, NeuriteType.soma, colors)
        ax.plot((displacement / 2., displacement), (0., 0.), color='k')
        ax.plot((0., displacement / 2.), (0., 0.), color='k')

    return displacement


def dendrogram(obj, show_diameters=True, new_fig=True, new_axes=True, subplot=False, **kwargs):
    '''
    Generates a figure of the neuron,
    that contains a soma and a list of trees.

    Parameters:
        obj: Neuron or tree \
        neurom.Neuron, neurom.Tree
        show_diameters : boolean \
            Determines if node diameters will \
            be show or not.
        new_fig: boolean \
            Defines if the tree will be plotted \
            in the current figure (False) \
            or in a new figure (True) \
            Default value is True.
        subplot: matplotlib subplot value or False \
            If False the default subplot 111 will be used. \
            For any other value a matplotlib subplot \
            will be generated. \
            Default value is False.

    Returns:
        A 2D matplotlib figure with a dendrogram view.
    '''
    if _compat.is_new_style(obj):
        from neurom.fst._dendrogram import Dendrogram
    else:
        from neurom.analysis.dendrogram import Dendrogram

    # create dendrogram and generate rectangle collection
    dnd = Dendrogram(obj, show_diameters=show_diameters)
    dnd.generate()

    fig, ax = common.get_figure(new_fig=new_fig, new_axes=new_axes, subplot=subplot)

    # render dendrogram and take into account neurite displacement which
    # starts as zero. It is important to avoid overlapping of neurites
    # and to determine tha limits of the figure.

    displacement = _render_dendrogram(dnd, ax, 0.)

    # customization settings
    kwargs['xlim'] = [- dnd.dims[0][0] * 0.5, dnd.dims[-1][0] * 0.5 + displacement]

    kwargs['title'] = kwargs.get('title', 'Morphology Dendrogram')
    kwargs['xlabel'] = kwargs.get('xlabel', 'micrometers (um)')
    kwargs['ylabel'] = kwargs.get('ylabel', 'micrometers (um)')
    kwargs['no_legend'] = False

    return common.plot_style(fig=fig, ax=ax, **kwargs)

neuron.__doc__ += DEFAULT_PARAMS # pylint: disable=no-member
tree.__doc__ += DEFAULT_PARAMS # pylint: disable=no-member
soma.__doc__ += DEFAULT_PARAMS # pylint: disable=no-member
neuron3d.__doc__ += DEFAULT_PARAMS # pylint: disable=no-member
tree3d.__doc__ += DEFAULT_PARAMS # pylint: disable=no-member
soma3d.__doc__ += DEFAULT_PARAMS # pylint: disable=no-member
