from vtk import *
from vtk.web import getJSArrayType
from vtk.web.camera import *
from vtk.web.query_data_model import *
from vtk.web import iteritems, buffer
import json, os, math, gzip, shutil

# Global helper variables
encode_codes = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'

# -----------------------------------------------------------------------------
# Capture image from render window
# -----------------------------------------------------------------------------

class CaptureRenderWindow(object):
    def __init__(self, magnification=1):
        self.windowToImage = vtkWindowToImageFilter()
        self.windowToImage.SetMagnification(magnification)
        self.windowToImage.SetInputBufferTypeToRGB()
        self.windowToImage.ReadFrontBufferOn()
        self.writer = None

    def SetRenderWindow(self, renderWindow):
        self.windowToImage.SetInput(renderWindow)

    def SetFormat(self, mimeType):
        if mimeType == 'image/png':
            self.writer = vtkPNGWriter()
            self.writer.SetInputConnection(self.windowToImage.GetOutputPort())
        elif mimeType == 'image/jpg':
            self.writer = vtkJPEGWriter()
            self.writer.SetInputConnection(self.windowToImage.GetOutputPort())

    def writeImage(self, path):
        if self.writer:
            self.windowToImage.Modified()
            self.windowToImage.Update()
            self.writer.SetFileName(path)
            self.writer.Write()

# -----------------------------------------------------------------------------
# Basic Dataset Builder
# -----------------------------------------------------------------------------

class DataSetBuilder(object):
    def __init__(self, location, camera_data, metadata={}, sections={}):
        self.dataHandler = DataHandler(location)
        self.cameraDescription = camera_data
        self.camera = None
        self.imageCapture = CaptureRenderWindow()

        for key, value in iteritems(metadata):
            self.dataHandler.addMetaData(key, value)

        for key, value in iteritems(sections):
            self.dataHandler.addSection(key, value)

    def getDataHandler(self):
        return self.dataHandler

    def getCamera(self):
        return self.camera

    def updateCamera(self, camera):
        update_camera(self.renderer, camera)
        self.renderWindow.Render()

    def start(self, renderWindow=None, renderer=None):
        if renderWindow:
            # Keep track of renderWindow and renderer
            self.renderWindow = renderWindow
            self.renderer = renderer

            # Initialize image capture
            self.imageCapture.SetRenderWindow(renderWindow)

            # Handle camera if any
            if self.cameraDescription:
                if self.cameraDescription['type'] == 'spherical':
                    self.camera = create_spherical_camera(renderer, self.dataHandler, self.cameraDescription['phi'], self.cameraDescription['theta'])
                elif self.cameraDescription['type'] == 'cylindrical':
                    self.camera = create_cylindrical_camera(renderer, self.dataHandler, self.cameraDescription['phi'], self.cameraDescription['translation'])

            # Update background color
            bgColor = renderer.GetBackground()
            bgColorString = 'rgb(%d, %d, %d)' % tuple(int(bgColor[i]*255) for i in range(3))
            self.dataHandler.addMetaData('backgroundColor', bgColorString)

        # Update file patterns
        self.dataHandler.updateBasePattern()

    def stop(self):
        self.dataHandler.writeDataDescriptor()

# -----------------------------------------------------------------------------
# Image Dataset Builder
# -----------------------------------------------------------------------------

class ImageDataSetBuilder(DataSetBuilder):
    def __init__(self, location, imageMimeType, cameraInfo, metadata={}, sections={}):
        DataSetBuilder.__init__(self, location, cameraInfo, metadata, sections)
        imageExtenstion = '.' + imageMimeType.split('/')[1]
        self.dataHandler.registerData(name='image', type='blob', mimeType=imageMimeType, fileName=imageExtenstion)
        self.imageCapture.SetFormat(imageMimeType)

    def writeImage(self):
        self.imageCapture.writeImage(self.dataHandler.getDataAbsoluteFilePath('image'))

    def writeImages(self):
        for cam in self.camera:
            update_camera(self.renderer, cam)
            self.renderWindow.Render()
            self.imageCapture.writeImage(self.dataHandler.getDataAbsoluteFilePath('image'))

# -----------------------------------------------------------------------------
# Volume Composite Dataset Builder
# -----------------------------------------------------------------------------
class VolumeCompositeDataSetBuilder(DataSetBuilder):
    def __init__(self, location, imageMimeType, cameraInfo, metadata={}, sections={}):
        DataSetBuilder.__init__(self, location, cameraInfo, metadata, sections)

        self.dataHandler.addTypes('volume-composite', 'rgba+depth')

        self.imageMimeType = imageMimeType
        self.imageExtenstion = '.' + imageMimeType.split('/')[1]

        if imageMimeType == 'image/png':
            self.imageWriter = vtkPNGWriter()
        if imageMimeType == 'image/jpg':
            self.imageWriter = vtkJPEGWriter()

        self.imageDataColor = vtkImageData()
        self.imageWriter.SetInputData(self.imageDataColor)

        self.imageDataDepth = vtkImageData()
        self.depthToWrite = None

        self.layerInfo = {}
        self.colorByMapping = {}
        self.compositePipeline = {'layers': [], 'dimensions': [], 'fields': {}, 'layer_fields': {}, 'pipeline': []}
        self.activeDepthKey = ''
        self.activeRGBKey = ''
        self.nodeWithChildren = {}

    def _getColorCode(self, colorBy):
        if colorBy in self.colorByMapping:
            # The color code exist
            return self.colorByMapping[colorBy]
        else:
            # No color code assigned yet
            colorCode = encode_codes[len(self.colorByMapping)]
            # Assign color code
            self.colorByMapping[colorBy] = colorCode
            # Register color code with color by value
            self.compositePipeline['fields'][colorCode] = colorBy
            # Return the color code
            return colorCode

    def _getLayerCode(self, parent, layerName):
        if layerName in self.layerInfo:
            # Layer already exist
            return (self.layerInfo[layerName]['code'], False)
        else:
            layerCode = encode_codes[len(self.layerInfo)]
            self.layerInfo[layerName] = { 'code': layerCode, 'name': layerName, 'parent': parent }
            self.compositePipeline['layers'].append(layerCode)
            self.compositePipeline['layer_fields'][layerCode] = []

            # Let's register it in the pipeline
            if parent:
                if parent not in self.nodeWithChildren:
                    # Need to create parent
                    rootNode = {'name': parent, 'ids': [], 'children':[]}
                    self.nodeWithChildren[parent] = rootNode
                    self.compositePipeline['pipeline'].append(rootNode)

                # Add node to its parent
                self.nodeWithChildren[parent]['children'].append({'name': layerName, 'ids': [layerCode]})
                self.nodeWithChildren[parent]['ids'].append(layerCode)

            else:
                self.compositePipeline['pipeline'].append({'name': layerName, 'ids': [layerCode]})

            return (layerCode, True)

    def _needToRegisterColor(self, layerCode, colorCode):
        if colorCode in self.compositePipeline['layer_fields'][layerCode]:
            return False
        else:
            self.compositePipeline['layer_fields'][layerCode].append(colorCode)
            return True

    def activateLayer(self, parent, name, colorBy):
        layerCode, needToRegisterDepth = self._getLayerCode(parent, name)
        colorCode = self._getColorCode(colorBy)
        needToRegisterColor = self._needToRegisterColor(layerCode, colorCode)

        # Update active keys
        self.activeDepthKey = '%s_depth' % layerCode
        self.activeRGBKey   = '%s%s_rgb' % (layerCode, colorCode)

        # Need to register data
        if needToRegisterDepth:
            self.dataHandler.registerData(name=self.activeDepthKey, type='array', fileName='/%s_depth.uint8' % layerCode, categories=[ layerCode ])

        if needToRegisterColor:
            self.dataHandler.registerData(name=self.activeRGBKey, type='blob', fileName='/%s%s_rgb%s' % (layerCode, colorCode, self.imageExtenstion), categories=[ '%s%s' % (layerCode, colorCode) ], mimeType=self.imageMimeType)

    def writeData(self, mapper):
        width = self.renderWindow.GetSize()[0]
        height = self.renderWindow.GetSize()[1]

        if not self.depthToWrite:
            self.depthToWrite = bytearray(width * height)

        for cam in self.camera:
            self.updateCamera(cam)
            imagePath = self.dataHandler.getDataAbsoluteFilePath(self.activeRGBKey)
            depthPath = self.dataHandler.getDataAbsoluteFilePath(self.activeDepthKey)

            # -----------------------------------------------------------------
            # Write Image
            # -----------------------------------------------------------------
            mapper.GetColorImage(self.imageDataColor)
            self.imageWriter.SetFileName(imagePath)
            self.imageWriter.Write()

            # -----------------------------------------------------------------
            # Write Depth
            # -----------------------------------------------------------------
            mapper.GetDepthImage(self.imageDataDepth)
            inputArray = self.imageDataDepth.GetPointData().GetArray(0)
            size = inputArray.GetNumberOfTuples()
            for idx in range(size):
                self.depthToWrite[idx] = int(inputArray.GetValue(idx))

            with open(depthPath, 'wb') as f:
                f.write(self.depthToWrite)

    def start(self, renderWindow, renderer):
        DataSetBuilder.start(self, renderWindow, renderer)
        self.camera.updatePriority([2,1])

    def stop(self, compress=True):
        # Push metadata
        self.compositePipeline['dimensions'] = self.renderWindow.GetSize()
        self.compositePipeline['default_pipeline'] = 'A'.join(self.compositePipeline['layers']) + 'A'
        self.dataHandler.addSection('CompositePipeline', self.compositePipeline)

        # Write metadata
        DataSetBuilder.stop(self)

        if compress:
            for root, dirs, files in os.walk(self.dataHandler.getBasePath()):
                print('Compress', root)
                for name in files:
                    if '.uint8' in name and '.gz' not in name:
                        with open(os.path.join(root, name), 'rb') as f_in:
                            with gzip.open(os.path.join(root, name + '.gz'), 'wb') as f_out:
                                shutil.copyfileobj(f_in, f_out)
                        os.remove(os.path.join(root, name))

# -----------------------------------------------------------------------------
# Data Prober Dataset Builder
# -----------------------------------------------------------------------------
class DataProberDataSetBuilder(DataSetBuilder):
    def __init__(self, location, sampling_dimesions, fields_to_keep, custom_probing_bounds = None, metadata={}):
        DataSetBuilder.__init__(self, location, None, metadata)
        self.fieldsToWrite = fields_to_keep
        self.resamplerFilter = vtkPResampleFilter()
        self.resamplerFilter.SetSamplingDimension(sampling_dimesions)
        if custom_probing_bounds:
            self.resamplerFilter.SetUseInputBounds(0)
            self.resamplerFilter.SetCustomSamplingBounds(custom_probing_bounds)
        else:
            self.resamplerFilter.SetUseInputBounds(1)

        # Register all fields
        self.dataHandler.addTypes('data-prober', 'binary')
        self.DataProber = { 'types': {}, 'dimensions': sampling_dimesions, 'ranges': {}, 'spacing': [1,1,1] }
        for field in self.fieldsToWrite:
            self.dataHandler.registerData(name=field, type='array', fileName='/%s.array' % field)

    def setDataToProbe(self, dataset):
        self.resamplerFilter.SetInputData(dataset)

    def setSourceToProbe(self, source):
        self.resamplerFilter.SetInputConnection(source.GetOutputPort())

    def writeData(self):
        self.resamplerFilter.Update()
        arrays = self.resamplerFilter.GetOutput().GetPointData()
        for field in self.fieldsToWrite:
            array = arrays.GetArray(field)
            if array:
                b = buffer(array)
                with open(self.dataHandler.getDataAbsoluteFilePath(field), 'wb') as f:
                    f.write(b)

                self.DataProber['types'][field] = getJSArrayType(array)
                if field in self.DataProber['ranges']:
                    dataRange = array.GetRange()
                    if dataRange[0] < self.DataProber['ranges'][field][0]:
                        self.DataProber['ranges'][field][0] = dataRange[0]
                    if dataRange[1] > self.DataProber['ranges'][field][1]:
                        self.DataProber['ranges'][field][1] = dataRange[1]
                else:
                    self.DataProber['ranges'][field] = [array.GetRange()[0], array.GetRange()[1]]
            else:
                print('No array for', field)
                print(self.resamplerFilter.GetOutput())

    def stop(self, compress=True):
        # Push metadata
        self.dataHandler.addSection('DataProber', self.DataProber)

        # Write metadata
        DataSetBuilder.stop(self)

        if compress:
            for root, dirs, files in os.walk(self.dataHandler.getBasePath()):
                print('Compress', root)
                for name in files:
                    if '.array' in name and '.gz' not in name:
                        with open(os.path.join(root, name), 'rb') as f_in:
                            with gzip.open(os.path.join(root, name + '.gz'), 'wb') as f_out:
                                shutil.copyfileobj(f_in, f_out)
                        os.remove(os.path.join(root, name))

# -----------------------------------------------------------------------------
# Sorted Composite Dataset Builder
# -----------------------------------------------------------------------------
class ConvertVolumeStackToSortedStack(object):
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.layers = 0

    def convert(self, directory):
        imagePaths = {}
        depthPaths = {}
        layerNames = []
        for fileName in os.listdir(directory):
            if '_rgb' in fileName or '_depth' in fileName:
                fileId = fileName.split('_')[0][0]
                if '_rgb' in fileName:
                    imagePaths[fileId] = os.path.join(directory, fileName)
                else:
                    layerNames.append(fileId)
                    depthPaths[fileId] = os.path.join(directory, fileName)

        layerNames.sort()

        if len(layerNames) == 0:
            return

        # Load data in Memory
        depthArrays = []
        imageReader = vtkPNGReader()
        numberOfValues = self.width * self.height * len(layerNames)
        imageSize = self.width * self.height
        self.layers = len(layerNames)

        # Write all images as single buffer
        opacity = vtkUnsignedCharArray()
        opacity.SetNumberOfComponents(1)
        opacity.SetNumberOfTuples(numberOfValues)

        intensity = vtkUnsignedCharArray()
        intensity.SetNumberOfComponents(1)
        intensity.SetNumberOfTuples(numberOfValues)

        for layer in range(self.layers):
            imageReader.SetFileName(imagePaths[layerNames[layer]])
            imageReader.Update()

            rgbaArray = imageReader.GetOutput().GetPointData().GetArray(0)

            for idx in range(imageSize):
                intensity.SetValue((layer * imageSize) + idx, rgbaArray.GetValue(idx*4))
                opacity.SetValue((layer * imageSize) + idx, rgbaArray.GetValue(idx*4 + 3))

            with open(depthPaths[layerNames[layer]], 'rb') as depthFile:
                depthArrays.append(depthFile.read())

        # Apply pixel sorting
        destOrder = vtkUnsignedCharArray()
        destOrder.SetNumberOfComponents(1)
        destOrder.SetNumberOfTuples(numberOfValues)

        opacityOrder = vtkUnsignedCharArray()
        opacityOrder.SetNumberOfComponents(1)
        opacityOrder.SetNumberOfTuples(numberOfValues)

        intensityOrder = vtkUnsignedCharArray()
        intensityOrder.SetNumberOfComponents(1)
        intensityOrder.SetNumberOfTuples(numberOfValues)

        for pixelIdx in range(imageSize):
            depthStack = []
            for depthArray in depthArrays:
                depthStack.append((depthArray[pixelIdx], len(depthStack)))
            depthStack.sort(key=lambda tup: tup[0])

            for destLayerIdx in range(len(depthStack)):
                sourceLayerIdx = depthStack[destLayerIdx][1]

                # Copy Idx
                destOrder.SetValue((imageSize * destLayerIdx) + pixelIdx, sourceLayerIdx)
                opacityOrder.SetValue((imageSize * destLayerIdx) + pixelIdx, opacity.GetValue((imageSize * sourceLayerIdx) + pixelIdx))
                intensityOrder.SetValue((imageSize * destLayerIdx) + pixelIdx, intensity.GetValue((imageSize * sourceLayerIdx) + pixelIdx))

        with open(os.path.join(directory, 'alpha.uint8'), 'wb') as f:
            f.write(buffer(opacityOrder))

        with open(os.path.join(directory, 'intensity.uint8'), 'wb') as f:
            f.write(buffer(intensityOrder))

        with open(os.path.join(directory, 'order.uint8'), 'wb') as f:
            f.write(buffer(destOrder))


class SortedCompositeDataSetBuilder(VolumeCompositeDataSetBuilder):
    def __init__(self, location, cameraInfo, metadata={}, sections={}):
        VolumeCompositeDataSetBuilder.__init__(self, location, 'image/png', cameraInfo, metadata, sections)
        self.dataHandler.addTypes('sorted-composite', 'rgba')

        # Register order and color textures
        self.layerScalars = []
        self.dataHandler.registerData(name='order',     type='array', fileName='/order.uint8')
        self.dataHandler.registerData(name='alpha',     type='array', fileName='/alpha.uint8')
        self.dataHandler.registerData(name='intensity', type='array', fileName='/intensity.uint8', categories=['intensity'])

    def start(self, renderWindow, renderer):
        VolumeCompositeDataSetBuilder.start(self, renderWindow, renderer)
        imageSize = self.renderWindow.GetSize()
        self.dataConverter = ConvertVolumeStackToSortedStack(imageSize[0], imageSize[1])

    def activateLayer(self, colorBy, scalar):
        VolumeCompositeDataSetBuilder.activateLayer(self, 'root', '%s' % scalar, colorBy)
        self.layerScalars.append(scalar)

    def writeData(self, mapper):
        VolumeCompositeDataSetBuilder.writeData(self, mapper)

        # Fill data pattern
        self.dataHandler.getDataAbsoluteFilePath('order')
        self.dataHandler.getDataAbsoluteFilePath('alpha')
        self.dataHandler.getDataAbsoluteFilePath('intensity')

    def stop(self, clean=True, compress=True):
        VolumeCompositeDataSetBuilder.stop(self, compress=False)

        # Go through all directories and convert them
        for root, dirs, files in os.walk(self.dataHandler.getBasePath()):
            for name in dirs:
                print('Process', os.path.join(root, name))
                self.dataConverter.convert(os.path.join(root, name))

        # Rename index.json to info_origin.json
        os.rename(os.path.join(self.dataHandler.getBasePath(), "index.json"), os.path.join(self.dataHandler.getBasePath(), "index_origin.json"))

        # Update index.json
        with open(os.path.join(self.dataHandler.getBasePath(), "index_origin.json"), "r") as infoFile:
            metadata = json.load(infoFile)
            metadata['SortedComposite'] = {
                'dimensions': metadata['CompositePipeline']['dimensions'],
                'layers': self.dataConverter.layers,
                'scalars': self.layerScalars[0:self.dataConverter.layers]
            }

            # Clean metadata
            dataToKeep = []
            del metadata['CompositePipeline']
            for item in metadata['data']:
                if item['name'] in ['order', 'alpha', 'intensity']:
                    dataToKeep.append(item)
            metadata['data'] = dataToKeep
            metadata['type'] = [ "tonic-query-data-model", "sorted-composite", "alpha" ]

            # Override index.json
            with open(os.path.join(self.dataHandler.getBasePath(), "index.json"), 'w') as newMetaFile:
                newMetaFile.write(json.dumps(metadata))

        # Clean temporary data
        if clean:
            for root, dirs, files in os.walk(self.dataHandler.getBasePath()):
                print('Clean', root)
                for name in files:
                    if '_rgb.png' in name or '_depth.uint8' in name or name == "index_origin.json":
                        os.remove(os.path.join(root, name))

        if compress:
            for root, dirs, files in os.walk(self.dataHandler.getBasePath()):
                print('Compress', root)
                for name in files:
                    if '.uint8' in name and '.gz' not in name:
                        with open(os.path.join(root, name), 'rb') as f_in:
                            with gzip.open(os.path.join(root, name + '.gz'), 'wb') as f_out:
                                shutil.copyfileobj(f_in, f_out)
                        os.remove(os.path.join(root, name))
