"""
A simple VTK widget for wxPython.  Note that wxPython comes
with its own wxVTKRenderWindow in wxPython.core.vtk.  Try both
and see which one works better for you.

Find wxPython info at http://wxPython.org

Based on vtkTkRenderWindget.py
"""

"""
Please see the example at the end of this file.

----------------------------------------
Creation:

wxVTKRenderWindow(parent, ID, stereo=0, [wx keywords]):

You should create a wxPySimpleApp() or some other wx**App
before creating the window.

----------------------------------------
Methods:

Render()
AddRenderer(ren)
GetRenderers()
GetRenderWindow()

----------------------------------------
Methods to override (all take a wxEvent):

OnButtonDown(event)  default: propagate event to Left, Right, Middle
OnLeftDown(event)    default: set _Mode to 'Rotate'
OnRightDown(event)   default: set _Mode to 'Zoom'
OnMiddleDown(event)  default: set _Mode to 'Pan'

OnButtonUp(event)    default: propagate event to L, R, M and unset _Mode
OnLeftUp(event)
OnRightUp(event)
OnMiddleUp(event)

OnMotion(event)      default: call appropriate handler for _Mode

OnEnterWindow(event) default: set focus to this window
OnLeaveWindow(event) default: release focus

OnKeyDown(event)     default: [R]eset, [W]irefreme, [S]olid, [P]ick
OnKeyUp(event)
OnChar(event)

OnSetFocus(event)
OnKillFocus(event)

OnSize(event)
OnMove(event)

OnPaint(event)       default: Render()

----------------------------------------
Protected Members:

_Mode:                 Current mode: 'Rotate', 'Zoom', 'Pan'
_LastX, _LastY:        The (x,y) coordinates of the previous event
_CurrentRenderer:      The renderer that was most recently clicked in
_CurrentCamera:        The camera for the current renderer

"""

# import usual libraries
import math
import os
import sys

import wx

# noinspection PyUnresolvedReferences
from vtk import vtkCellPicker, vtkProperty, vtkRenderWindow

# a few configuration items, see what works best on your system

# Use wxGLCanvas as base class instead of wxWindow.
# This is sometimes necessary under wxGTK or the image is blank.
# (in wxWindows 2.3.1 and earlier, the GLCanvas had scroll bars)
# try:
#    wx.WX_USE_GL_CANVAS
# except NameError:
if True:
    if wx.Platform == "__WXMSW__":
        WX_USE_GLCANVAS = 0
    else:
        WX_USE_GLCANVAS = 1
WX_USE_GLCANVAS = 0
# Keep capturing mouse after mouse is dragged out of window
# (in wxGTK 2.3.2 there is a bug that keeps this from working,
# but it is only relevant in wxGTK if there are multiple windows)
# try:
#    wx.WX_USE_X_CAPTURE
# except NameError:
if True:
    if wx.Platform == "__WXMSW__":
        WX_USE_X_CAPTURE = 1
    else:
        WX_USE_X_CAPTURE = 0

# end of configuration items


if WX_USE_GLCANVAS:
    from wx.glcanvas import GLCanvas

    baseClass = GLCanvas
else:
    baseClass = wx.Window


class wxVTKRenderWindow(baseClass):
    """
    A wxRenderWindow for wxPython.
    Use GetRenderWindow() to get the vtkRenderWindow.
    Create with the keyword stereo=1 in order to
    generate a stereo-capable window.
    """

    def __init__(self, parent, ID, *args, **kw):
        # miscellaneous protected variables
        self._CurrentRenderer = None
        self._CurrentCamera = None
        self._CurrentZoom = 1.0
        self._CurrentLight = None

        self._ViewportCenterX = 0
        self._ViewportCenterY = 0

        self._Picker = vtkCellPicker()
        self._PickedActor = None
        self._PickedProperty = vtkProperty()
        self._PickedProperty.SetColor(1, 0, 0)
        self._PrePickedProperty = None

        # these record the previous mouse position
        self._LastX = 0
        self._LastY = 0

        # the current interaction mode (Rotate, Pan, Zoom, etc)
        self._Mode = None
        self._ActiveButton = None

        # private attributes
        self.__OldFocus = None

        # used by the LOD actors
        self._DesiredUpdateRate = 15
        self._StillUpdateRate = 0.0001

        # First do special handling of some keywords:
        # stereo, position, size, width, height, style

        stereo = 0

        if "stereo" in kw:
            if kw["stereo"]:
                stereo = 1
            del kw["stereo"]

        position = wx.DefaultPosition

        if "position" in kw:
            position = kw["position"]
            del kw["position"]

        size = wx.DefaultSize

        if "size" in kw:
            size = kw["size"]
            del kw["size"]

        if "width" in kw and "height" in kw:
            size = (kw["width"], kw["height"])
            del kw["width"]
            del kw["height"]

        # wxWANTS_CHARS says to give us e.g. TAB
        # wxNO_FULL_REPAINT_ON_RESIZE cuts down resize flicker under GTK
        style = wx.WANTS_CHARS | wx.NO_FULL_REPAINT_ON_RESIZE

        if "style" in kw:
            style = style | kw["style"]
            del kw["style"]

        # the enclosing frame must be shown under GTK or the windows
        #  don't connect together properly
        l = []
        p = parent
        while p:  # make a list of all parents
            l.append(p)
            p = p.GetParent()
        l.reverse()  # sort list into descending order
        for p in l:
            p.Show()

        # initialize the wxWindow
        baseClass.__init__(self, parent, ID, position, size, style)

        # create the RenderWindow and initialize it
        self._RenderWindow = vtkRenderWindow()
        try:
            self._RenderWindow.SetSize(size.width, size.height)
        except AttributeError:
            self._RenderWindow.SetSize(size[0], size[1])
        if stereo:
            self._RenderWindow.StereoCapableWindowOn()
            self._RenderWindow.SetStereoTypeToCrystalEyes()

        self.__Created = 0
        # Tell the RenderWindow to render inside the wxWindow.
        if self.GetHandle():
            self.__Created = 1
            self._RenderWindow.SetWindowInfo(str(self.GetHandle()))

        # refresh window by doing a Render
        self.Bind(wx.EVT_PAINT, self.OnPaint)
        # turn off background erase to reduce flicker
        self.Bind(wx.EVT_ERASE_BACKGROUND, lambda e: None)

        # Bind the events to the event converters
        self.Bind(wx.EVT_RIGHT_DOWN, self._OnButtonDown)
        self.Bind(wx.EVT_LEFT_DOWN, self._OnButtonDown)
        self.Bind(wx.EVT_MIDDLE_DOWN, self._OnButtonDown)
        self.Bind(wx.EVT_RIGHT_UP, self._OnButtonUp)
        self.Bind(wx.EVT_LEFT_UP, self._OnButtonUp)
        self.Bind(wx.EVT_MIDDLE_UP, self._OnButtonUp)
        self.Bind(wx.EVT_MOTION, self.OnMotion)

        self.Bind(wx.EVT_ENTER_WINDOW, self._OnEnterWindow)
        self.Bind(wx.EVT_LEAVE_WINDOW, self._OnLeaveWindow)

        self.Bind(wx.EVT_CHAR, self.OnChar)
        self.Bind(wx.EVT_KEY_DOWN, self.OnKeyDown)
        self.Bind(wx.EVT_KEY_UP, self.OnKeyUp)

        self.Bind(wx.EVT_SIZE, self._OnSize)
        self.Bind(wx.EVT_MOVE, self.OnMove)

        self.Bind(wx.EVT_SET_FOCUS, self.OnSetFocus)
        self.Bind(wx.EVT_KILL_FOCUS, self.OnKillFocus)

    def SetDesiredUpdateRate(self, rate):
        """Mirrors the method with the same name in
        vtkRenderWindowInteractor."""
        self._DesiredUpdateRate = rate

    def GetDesiredUpdateRate(self):
        """Mirrors the method with the same name in
        vtkRenderWindowInteractor."""
        return self._DesiredUpdateRate

    def SetStillUpdateRate(self, rate):
        """Mirrors the method with the same name in
        vtkRenderWindowInteractor."""
        self._StillUpdateRate = rate

    def GetStillUpdateRate(self):
        """Mirrors the method with the same name in
        vtkRenderWindowInteractor."""
        return self._StillUpdateRate

    def OnPaint(self, event):
        dc = wx.PaintDC(self)
        self.Render()

    def _OnSize(self, event):
        if wx.Platform != "__WXMSW__":
            try:
                width, height = event.GetSize()
            except:
                width = event.GetSize().width
                height = event.GetSize().height
            self._RenderWindow.SetSize(width, height)
        self.OnSize(event)
        if self.__Created:
            self._RenderWindow.Render()

    def OnSize(self, event):
        pass

    def OnMove(self, event):
        pass

    def _OnEnterWindow(self, event):
        self.UpdateRenderer(event)
        self.OnEnterWindow(event)

    def OnEnterWindow(self, event):
        if self.__OldFocus is None:
            self.__OldFocus = wx.Window.FindFocus()
            self.SetFocus()

    def _OnLeaveWindow(self, event):
        self.OnLeaveWindow(event)

    def OnLeaveWindow(self, event):
        if self.__OldFocus:
            self.__OldFocus.SetFocus()
            self.__OldFocus = None

    def OnSetFocus(self, event):
        pass

    def OnKillFocus(self, event):
        pass

    def _OnButtonDown(self, event):
        # helper function for capturing mouse until button released
        self._RenderWindow.SetDesiredUpdateRate(self._DesiredUpdateRate)

        if event.RightDown():
            button = "Right"
        elif event.LeftDown():
            button = "Left"
        elif event.MiddleDown():
            button = "Middle"
        else:
            button = None

        # save the button and capture mouse until the button is released
        if button and not self._ActiveButton:
            self._ActiveButton = button
            if WX_USE_X_CAPTURE:
                self.CaptureMouse()

        self.OnButtonDown(event)

    def OnButtonDown(self, event):
        if not self._Mode:
            # figure out what renderer the mouse is over
            self.UpdateRenderer(event)

        if event.LeftDown():
            self.OnLeftDown(event)
        elif event.RightDown():
            self.OnRightDown(event)
        elif event.MiddleDown():
            self.OnMiddleDown(event)

    def OnLeftDown(self, event):
        if not self._Mode:
            if event.ControlDown():
                self._Mode = "Zoom"
            elif event.ShiftDown():
                self._Mode = "Pan"
            else:
                self._Mode = "Rotate"

    def OnRightDown(self, event):
        if not self._Mode:
            self._Mode = "Zoom"

    def OnMiddleDown(self, event):
        if not self._Mode:
            self._Mode = "Pan"

    def _OnButtonUp(self, event):
        # helper function for releasing mouse capture
        self._RenderWindow.SetDesiredUpdateRate(self._StillUpdateRate)

        if event.RightUp():
            button = "Right"
        elif event.LeftUp():
            button = "Left"
        elif event.MiddleUp():
            button = "Middle"
        else:
            button = None

        # if the ActiveButton is realeased, then release mouse capture
        if self._ActiveButton and button == self._ActiveButton:
            if WX_USE_X_CAPTURE:
                self.ReleaseMouse()
            self._ActiveButton = None

        self.OnButtonUp(event)

    def OnButtonUp(self, event):
        if event.LeftUp():
            self.OnLeftUp(event)
        elif event.RightUp():
            self.OnRightUp(event)
        elif event.MiddleUp():
            self.OnMiddleUp(event)

        # if not interacting, then do nothing more
        if self._Mode:
            if self._CurrentRenderer:
                self.Render()

        self._Mode = None

    def OnLeftUp(self, event):
        pass

    def OnRightUp(self, event):
        pass

    def OnMiddleUp(self, event):
        pass

    def OnMotion(self, event):
        if self._Mode == "Pan":
            self.Pan(event)
        elif self._Mode == "Rotate":
            self.Rotate(event)
        elif self._Mode == "Zoom":
            self.Zoom(event)

    def OnChar(self, event):
        pass

    def OnKeyDown(self, event):
        if event.GetKeyCode() == ord("R"):
            self.Reset(event)
        if event.GetKeyCode() == ord("W"):
            self.Wireframe()
        if event.GetKeyCode() == ord("S"):
            self.Surface()
        if event.GetKeyCode() == ord("P"):
            self.PickActor(event)

        if event.GetKeyCode() < 256:
            self.OnChar(event)

    def OnKeyUp(self, event):
        pass

    def GetZoomFactor(self):
        return self._CurrentZoom

    def GetRenderWindow(self):
        return self._RenderWindow

    def GetPicker(self):
        return self._Picker

    def Render(self):
        if self._CurrentLight:
            light = self._CurrentLight
            light.SetPosition(self._CurrentCamera.GetPosition())
            light.SetFocalPoint(self._CurrentCamera.GetFocalPoint())

        if self.__Created:
            self._RenderWindow.Render()
        elif self.GetHandle():
            self._RenderWindow.SetWindowInfo(str(self.GetHandle()))
            self.__Created = 1
            self._RenderWindow.Render()

    def UpdateRenderer(self, event):
        """
        UpdateRenderer will identify the renderer under the mouse and set
        up _CurrentRenderer, _CurrentCamera, and _CurrentLight.
        """
        x = event.GetX()
        y = event.GetY()
        windowX, windowY = self._RenderWindow.GetSize()

        renderers = self._RenderWindow.GetRenderers()
        numRenderers = renderers.GetNumberOfItems()

        self._CurrentRenderer = None
        renderers.InitTraversal()
        for i in range(0, numRenderers):
            renderer = renderers.GetNextItem()
            vx, vy = (0, 0)
            if windowX > 1:
                vx = float(x) / (windowX - 1)
            if windowY > 1:
                vy = (windowY - float(y) - 1) / (windowY - 1)
            (vpxmin, vpymin, vpxmax, vpymax) = renderer.GetViewport()

            if vpxmin <= vx <= vpxmax and vpymin <= vy <= vpymax:
                self._CurrentRenderer = renderer
                self._ViewportCenterX = float(windowX) * (vpxmax - vpxmin) / 2.0 + vpxmin
                self._ViewportCenterY = float(windowY) * (vpymax - vpymin) / 2.0 + vpymin
                self._CurrentCamera = self._CurrentRenderer.GetActiveCamera()
                lights = self._CurrentRenderer.GetLights()
                lights.InitTraversal()
                self._CurrentLight = lights.GetNextItem()
                break

        self._LastX = x
        self._LastY = y

    def GetCurrentRenderer(self):
        return self._CurrentRenderer

    def Rotate(self, event):
        if self._CurrentRenderer:
            x = event.GetX()
            y = event.GetY()

            self._CurrentCamera.Azimuth(self._LastX - x)
            self._CurrentCamera.Elevation(y - self._LastY)
            self._CurrentCamera.OrthogonalizeViewUp()

            self._LastX = x
            self._LastY = y

            self._CurrentRenderer.ResetCameraClippingRange()
            self.Render()

    def Pan(self, event):
        if self._CurrentRenderer:
            x = event.GetX()
            y = event.GetY()

            renderer = self._CurrentRenderer
            camera = self._CurrentCamera
            (pPoint0, pPoint1, pPoint2) = camera.GetPosition()
            (fPoint0, fPoint1, fPoint2) = camera.GetFocalPoint()

            if camera.GetParallelProjection():
                renderer.SetWorldPoint(fPoint0, fPoint1, fPoint2, 1.0)
                renderer.WorldToDisplay()
                fx, fy, fz = renderer.GetDisplayPoint()
                renderer.SetDisplayPoint(fx - x + self._LastX, fy + y - self._LastY, fz)
                renderer.DisplayToWorld()
                fx, fy, fz, fw = renderer.GetWorldPoint()
                camera.SetFocalPoint(fx, fy, fz)

                renderer.SetWorldPoint(pPoint0, pPoint1, pPoint2, 1.0)
                renderer.WorldToDisplay()
                fx, fy, fz = renderer.GetDisplayPoint()
                renderer.SetDisplayPoint(fx - x + self._LastX, fy + y - self._LastY, fz)
                renderer.DisplayToWorld()
                fx, fy, fz, fw = renderer.GetWorldPoint()
                camera.SetPosition(fx, fy, fz)

            else:
                (fPoint0, fPoint1, fPoint2) = camera.GetFocalPoint()
                # Specify a point location in world coordinates
                renderer.SetWorldPoint(fPoint0, fPoint1, fPoint2, 1.0)
                renderer.WorldToDisplay()
                # Convert world point coordinates to display coordinates
                dPoint = renderer.GetDisplayPoint()
                focalDepth = dPoint[2]

                aPoint0 = self._ViewportCenterX + (x - self._LastX)
                aPoint1 = self._ViewportCenterY - (y - self._LastY)

                renderer.SetDisplayPoint(aPoint0, aPoint1, focalDepth)
                renderer.DisplayToWorld()

                (rPoint0, rPoint1, rPoint2, rPoint3) = renderer.GetWorldPoint()
                if rPoint3 != 0.0:
                    rPoint0 = rPoint0 / rPoint3
                    rPoint1 = rPoint1 / rPoint3
                    rPoint2 = rPoint2 / rPoint3

                camera.SetFocalPoint(
                    (fPoint0 - rPoint0) + fPoint0, (fPoint1 - rPoint1) + fPoint1, (fPoint2 - rPoint2) + fPoint2
                )

                camera.SetPosition(
                    (fPoint0 - rPoint0) + pPoint0, (fPoint1 - rPoint1) + pPoint1, (fPoint2 - rPoint2) + pPoint2
                )

            self._LastX = x
            self._LastY = y

            self.Render()

    def Zoom(self, event):
        if self._CurrentRenderer:
            x = event.GetX()
            y = event.GetY()

            renderer = self._CurrentRenderer
            camera = self._CurrentCamera

            zoomFactor = math.pow(1.02, (0.5 * (self._LastY - y)))
            self._CurrentZoom = self._CurrentZoom * zoomFactor

            if camera.GetParallelProjection():
                parallelScale = camera.GetParallelScale() / zoomFactor
                camera.SetParallelScale(parallelScale)
            else:
                camera.Dolly(zoomFactor)
                renderer.ResetCameraClippingRange()

            self._LastX = x
            self._LastY = y

            self.Render()

    def Reset(self, event=None):
        if self._CurrentRenderer:
            self._CurrentRenderer.ResetCamera()

        self.Render()

    def Wireframe(self):
        actors = self._CurrentRenderer.GetActors()
        numActors = actors.GetNumberOfItems()
        actors.InitTraversal()
        for i in range(0, numActors):
            actor = actors.GetNextItem()
            actor.GetProperty().SetRepresentationToWireframe()

        self.Render()

    def Surface(self):
        actors = self._CurrentRenderer.GetActors()
        numActors = actors.GetNumberOfItems()
        actors.InitTraversal()
        for i in range(0, numActors):
            actor = actors.GetNextItem()
            actor.GetProperty().SetRepresentationToSurface()

        self.Render()

    def PickActor(self, event):
        if self._CurrentRenderer:
            x = event.GetX()
            y = event.GetY()

            renderer = self._CurrentRenderer
            picker = self._Picker

            windowX, windowY = self._RenderWindow.GetSize()
            picker.Pick(x, (windowY - y - 1), 0.0, renderer)
            actor = picker.GetActor()

            if self._PickedActor is not None and self._PrePickedProperty is not None:
                self._PickedActor.SetProperty(self._PrePickedProperty)
                # release hold of the property
                self._PrePickedProperty.UnRegister(self._PrePickedProperty)
                self._PrePickedProperty = None

            if actor is not None:
                self._PickedActor = actor
                self._PrePickedProperty = self._PickedActor.GetProperty()
                # hold onto the property
                self._PrePickedProperty.Register(self._PrePickedProperty)
                self._PickedActor.SetProperty(self._PickedProperty)

            self.Render()
