import os, time

from wslink import register as exportRpc

from paraview import simple, servermanager
from paraview.web import protocols as pv_protocols

from vtkmodules.vtkCommonCore import vtkUnsignedCharArray, vtkCollection
from vtkmodules.vtkCommonDataModel import vtkImageData
from vtkmodules.vtkWebCore import vtkDataEncoder

try:
    # PV 5.6
    from vtkmodules.vtkPVClientServerCoreRendering import vtkPVRenderView
    from vtkmodules.vtkPVServerManagerRendering import vtkSMPVRepresentationProxy, vtkSMTransferFunctionProxy, vtkSMTransferFunctionManager
except:
    pass

try:
    # PV 5.7
    from paraview.modules.vtkPVClientServerCoreRendering import vtkPVRenderView
    from paraview.modules.vtkPVServerManagerRendering import vtkSMPVRepresentationProxy, vtkSMTransferFunctionProxy, vtkSMTransferFunctionManager
except:
    pass

try:
    # PV 5.8
    from paraview.servermanager import vtkPVRenderView
    from paraview.servermanager import vtkSMPVRepresentationProxy
    from paraview.servermanager import vtkSMTransferFunctionManager
    from paraview.servermanager import vtkSMTransferFunctionProxy
except:
    pass

class ParaViewLite(pv_protocols.ParaViewWebProtocol):
    def __init__(self, **kwargs):
      super(pv_protocols.ParaViewWebProtocol, self).__init__()
      self.lineContext = None

    @exportRpc("paraview.lite.proxy.name")
    def getProxyName(self, pid):
      proxy = self.mapIdToProxy(pid)
      if not proxy:
        return {
          'id': pid,
          'error': 'No proxy for id %s' % pid,
        }

      return {
        'id': pid,
        'group': proxy.GetXMLGroup(),
        'name': proxy.GetXMLName(),
        'label': proxy.GetXMLLabel(),
      }

    @exportRpc("paraview.lite.camera.get")
    def getCamera(self, viewId):
      view = self.getView(viewId)
      bounds = [-1, 1, -1, 1, -1, 1]

      if view and view.GetClientSideView().GetClassName() == 'vtkPVRenderView':
        rr = view.GetClientSideView().GetRenderer()
        bounds = rr.ComputeVisiblePropBounds()

      return {
        'id': viewId,
        'bounds': bounds,
        'position': tuple(view.CameraPosition),
        'viewUp': tuple(view.CameraViewUp),
        'focalPoint': tuple(view.CameraFocalPoint),
        'centerOfRotation': tuple(view.CenterOfRotation),
      }

    @exportRpc("paraview.lite.lut.get")
    def getLookupTableForArrayName(self, name, numSamples = 255):
      lutProxy = simple.GetColorTransferFunction(name)
      lut = lutProxy.GetClientSideObject()
      dataRange = lut.GetRange()
      delta = (dataRange[1] - dataRange[0]) / float(numSamples)

      colorArray = vtkUnsignedCharArray()
      colorArray.SetNumberOfComponents(3)
      colorArray.SetNumberOfTuples(numSamples)

      rgb = [ 0, 0, 0 ]
      for i in range(numSamples):
          lut.GetColor(dataRange[0] + float(i) * delta, rgb)
          r = int(round(rgb[0] * 255))
          g = int(round(rgb[1] * 255))
          b = int(round(rgb[2] * 255))
          colorArray.SetTuple3(i, r, g, b)

      # Add the color array to an image data
      imgData = vtkImageData()
      imgData.SetDimensions(numSamples, 1, 1)
      aIdx = imgData.GetPointData().SetScalars(colorArray)

      # Use the vtk data encoder to base-64 encode the image as png, using no compression
      encoder = vtkDataEncoder()
      # two calls in a row crash on Windows - bald timing hack to avoid the crash.
      time.sleep(0.01);
      b64Str = encoder.EncodeAsBase64Jpg(imgData, 100)

      return { 'image': 'data:image/jpg;base64,' + b64Str, 'range': dataRange, 'name': name }


    @exportRpc("paraview.lite.lut.range.update")
    def updateLookupTableRange(self, arrayName, dataRange):
      lutProxy = simple.GetColorTransferFunction(arrayName)
      vtkSMTransferFunctionProxy.RescaleTransferFunction(lutProxy.SMProxy, dataRange[0], dataRange[1], False)
      self.getApplication().InvokeEvent('UpdateEvent')


    @exportRpc("paraview.lite.lut.preset")
    def getLookupTablePreset(self, presetName, numSamples = 512):
      lutProxy = simple.GetColorTransferFunction('__PRESET__')
      lutProxy.ApplyPreset(presetName, True)
      lut = lutProxy.GetClientSideObject()
      dataRange = lut.GetRange()
      delta = (dataRange[1] - dataRange[0]) / float(numSamples)

      colorArray = vtkUnsignedCharArray()
      colorArray.SetNumberOfComponents(3)
      colorArray.SetNumberOfTuples(numSamples)

      rgb = [ 0, 0, 0 ]
      for i in range(numSamples):
          lut.GetColor(dataRange[0] + float(i) * delta, rgb)
          r = int(round(rgb[0] * 255))
          g = int(round(rgb[1] * 255))
          b = int(round(rgb[2] * 255))
          colorArray.SetTuple3(i, r, g, b)

      # Add the color array to an image data
      imgData = vtkImageData()
      imgData.SetDimensions(numSamples, 1, 1)
      aIdx = imgData.GetPointData().SetScalars(colorArray)

      # Use the vtk data encoder to base-64 encode the image as png, using no compression
      encoder = vtkDataEncoder()
      # two calls in a row crash on Windows - bald timing hack to avoid the crash.
      time.sleep(0.01);
      b64Str = encoder.EncodeAsBase64Jpg(imgData, 100)

      return { 'name': presetName, 'image': 'data:image/jpg;base64,' + b64Str }


    @exportRpc("paraview.lite.lut.set.preset")
    def applyPreset(self, arrayName, presetName):
      lutProxy = simple.GetColorTransferFunction(arrayName)
      lutProxy.ApplyPreset(presetName, True)
      self.getApplication().InvokeEvent('UpdateEvent')


    @exportRpc("paraview.lite.context.line.set")
    def updateLineContext(self, visible = False, p1 = [0, 0, 0], p2 = [1, 1, 1]):
      if not self.lineContext:
        self.lineContext = servermanager.extended_sources.LineSource(Resolution=2, Point1=p1, Point2=p2)
        self.lineRepresentation = simple.Show(self.lineContext)

      self.lineRepresentation.Visibility = 1 if visible else 0
      self.lineContext.Point1 = p1
      self.lineContext.Point2 = p2

      self.getApplication().InvokeEvent('UpdateEvent')

      return self.lineContext.GetGlobalIDAsString()