# -*- coding: utf-8 -*-
"""Mayavi/traits GUI visualization elements."""

# Authors: Christian Brodbeck <christianbrodbeck@nyu.edu>
#
# License: BSD (3-clause)

import numpy as np

from mayavi.mlab import pipeline, text3d
from mayavi.modules.glyph import Glyph
from mayavi.modules.surface import Surface
from mayavi.sources.vtk_data_source import VTKDataSource
from mayavi.tools.mlab_scene_model import MlabSceneModel
from traits.api import (HasTraits, HasPrivateTraits, on_trait_change,
                        Instance, Array, Bool, Button, Enum, Float, Int, List,
                        Range, Str, Property, cached_property, ArrayOrNone)
from traitsui.api import (View, Item, HGroup, VGrid, VGroup, Spring,
                          TextEditor)
from tvtk.api import tvtk

from ..defaults import DEFAULTS
from ..surface import _CheckInside, _DistanceQuery
from ..transforms import apply_trans
from ..utils import SilenceStdout
from ..viz.backends._pysurfer_mayavi import (_create_mesh_surf,
                                             _toggle_mlab_render)

try:
    from traitsui.api import RGBColor
except ImportError:
    from traits.api import RGBColor

headview_borders = VGroup(Item('headview', style='custom', show_label=False),
                          show_border=True, label='View')


def _mm_fmt(x):
    """Format data in units of mm."""
    return '%0.1f' % x


laggy_float_editor_mm = TextEditor(auto_set=False, enter_set=True,
                                   evaluate=float,
                                   format_func=lambda x: '%0.1f' % x)

laggy_float_editor_scale = TextEditor(auto_set=False, enter_set=True,
                                      evaluate=float,
                                      format_func=lambda x: '%0.1f' % x)

laggy_float_editor_headscale = TextEditor(auto_set=False, enter_set=True,
                                          evaluate=float,
                                          format_func=lambda x: '%0.3f' % x)

laggy_float_editor_weight = TextEditor(auto_set=False, enter_set=True,
                                       evaluate=float,
                                       format_func=lambda x: '%0.2f' % x)

laggy_float_editor_deg = TextEditor(auto_set=False, enter_set=True,
                                    evaluate=float,
                                    format_func=lambda x: '%0.1f' % x)

_BUTTON_WIDTH = -80
_DEG_WIDTH = -50  # radian floats
_MM_WIDTH = _DEG_WIDTH  # mm floats
_SCALE_WIDTH = _DEG_WIDTH  # scale floats
_INC_BUTTON_WIDTH = -25  # inc/dec buttons
_DEG_STEP_WIDTH = -50
_MM_STEP_WIDTH = _DEG_STEP_WIDTH
_SCALE_STEP_WIDTH = _DEG_STEP_WIDTH
_WEIGHT_WIDTH = -60  # weight floats
_VIEW_BUTTON_WIDTH = -60
# width is optimized for macOS and Linux avoid a horizontal scroll-bar
# even when a vertical one is present
_COREG_WIDTH = -290
_TEXT_WIDTH = -260
_REDUCED_TEXT_WIDTH = _TEXT_WIDTH - 40 * np.sign(_TEXT_WIDTH)
_DIG_SOURCE_WIDTH = _TEXT_WIDTH - 50 * np.sign(_
TEXT_WIDTH)
_MRI_FIDUCIALS_WIDTH = _TEXT_WIDTH - 60 * np.sign(_TEXT_WIDTH)
_SHOW_BORDER = True
_RESET_LABEL = u"↻"
_RESET_WIDTH = _INC_BUTTON_WIDTH


class HeadViewController(HasTraits):
    """Set head views for the given coordinate system.

    Para
meters
    ----------
    system : 'RAS' | 'ALS' | 'ARI'
        Coordinate system described as initials for directions associated with
        the x, y, and z axes. Relevant terms are: Anterior, Right, Left,
        Superior, Inferior.
    """

    system = Enum("RAS", "ALS", "ARI", desc="Coordinate system: directions of "
                  "the x, y, and z axis.")

    right = Button()
    front = Button()
    left = Button()
    top = Button()
    interaction = Enum('trackball', 'terrain')

    scale = Float(0.16)

    scene = Instance(MlabSceneModel)

    view = View(VGroup(
        VGrid('0', Item('top', width=_VIEW_BUTTON_WIDTH), '0',
              Item('right', width=_VIEW_BUTTON_WIDTH),
              Item('front', width=_VIEW_BUTTON_WIDTH),
              Item('left', width=_VIEW_BUTTON_WIDTH),
              columns=3, show_labels=False),
        '_',
        HGroup(Item('scale', label='Scale',
                    editor=laggy_float_editor_headscale,
                    width=_SCALE_WIDTH, show_label=True),
               Item('interaction', tooltip='Mouse interaction mode',
                    show_label=False), Spring()),
        show_labels=False))

    @on_trait_change('scene.activated')
    def _init_view(self):
        self.scene.parallel_projection = True

        # apparently scene,activated happens several times
        if self.scene.renderer:
            self.sync_trait('scale', self.scene.camera, 'parallel_scale')
            # and apparently this does not happen by default:
            self.on_trait_change(self.scene.render, 'scale')
            self.interaction = self.interaction  # could be delayed

    @on_trait_change('interaction')
    def on_set_interaction(self, _, interaction):
        if self.scene is None or self.scene.interactor is None:
            return
        # Ensure we're in the correct orientation for the
        # InteractorStyleTerrain to have the correct "up"
        self.on_set_view('front', '')
        self.scene.mlab.draw()
        self.scene.interactor.interactor_style = \
            tvtk.InteractorStyleTerrain() if interaction == 'terrain' else \
            tvtk.InteractorStyleTrackballCamera()
        # self.scene.interactor.interactor_style.
        self.on_set_view('front', '')
        self.scene.mlab.draw()

    @on_trait_change('top,left,right,front')
    def on_set_view(self, view, _):
        if self.scene is None:
            return

        system = self.system
        kwargs = dict(ALS=dict(front=(0, 90, -90),
                               left=(90, 90, 180),
                               right=(-90, 90, 0),
                               top=(0, 0, -90)),
                      RAS=dict(front=(90., 90., 180),
                               left=(180, 90, 90),
                               right=(0., 90, 270),
                               top=(90, 0, 180)),
                      ARI=dict(front=(0, 90, 90),
                               left=(-90, 90, 180),
                               right=(90, 90, 0),
                               top=(0, 180, 90)))
        if system not in kwargs:
            raise ValueError("Invalid system: %r" % system)
        if view not in kwargs[system]:
            raise ValueError("Invalid view: %r" % view)
        kwargs = dict(zip(('azimuth', 'elevation', 'roll'),
                          kwargs[system][view]))
        kwargs['focalpoint'] = (0., 0., 0.)
        with SilenceStdout():
            self.scene.mlab.view(distance=None, reset_roll=True,
                                 figure=self.scene.mayavi_scene, **kwargs)


class Object(HasPrivateTraits):
    """Represent a 3d object in a mayavi scene."""

    points = Array(float, shape=(None, 3))
    nn = Array(float, shape=(None, 3))
    name = Str

    scene = Instance(MlabSceneModel, ())
    src = Instance(VTKDataSource)

    # This should be Tuple, but it is broken on Anaconda as of 2016/12/16
    color = RGBColor((1., 1., 1.))
    # Due to a MESA bug, we use 0.99 opacity to force alpha blending
    opacity = Range(low=0., high=1., value=0.99)
    visible = Bool(True)

    def _update_points(self):
        """Update the location of the plotted points."""
        if hasattr(self.src, 'data'):
            self.src.data.points = self.points
            return True


class PointObject(Object):
    """Represent a group of individual points in a mayavi scene."""

    label = Bool(False)
    label_scale = Float(0.01)
    projectable = Bool(False)  # set based on type of points
    orientable = Property(depends_on=['nearest'])
    text3d = List
    point_scale = Float(10, label='Point Scale')

    # projection onto a surface
    nearest = Instance(_DistanceQuery)
    check_inside = Instance(_CheckInside)
    project_to_trans = ArrayOrNone(float, shape=(4, 4))
    project_to_surface = Bool(False, label='Project', desc='project points '
                              'onto the surface')
    orient_to_surface = Bool(False, label='Orient', desc