#!/usr/bin/env python

# Testing vtkWindowToImageFilter, vtkPNGWriter and vtkPNGReader.
import sys
import vtk
import vtk.test.Testing
from vtk.util.misc import vtkGetDataRoot
VTK_DATA_ROOT = vtkGetDataRoot()

class TestWIFPNGWriteRead(vtk.test.Testing.vtkTest):
    '''
    Test class for testing write/read of a PNG file.
    '''

    # Uncomment these lines to write a log file.
    # fow = vtk.vtkFileOutputWindow()
    # fow.SetFileName('TestWIFPNGWriteRead.log')
    # ow = vtk.vtkOutputWindow()
    # ow.SetInstance(fow)

    fn = ['TestWIFPNGRGB.png', 'TestWIFPNGRGBA.png']


    def WritePNG(self, idx):
        '''
            1. Create a source.
            2. Take a (partial) screenshot using the WindowToImageFilter.
            3. Write the screenshot to a file.

            :param: idx - If 0, write a RGB file, if 1 write a RGBA file.
        '''
        # try:
        #     # Can we write to the file in the directory?
        #     channel = open(self.fn[idx], 'wb')
        #     channel.close()
        # except IOError:
        #     print('Error: Unable to write file.')
        #     sys.exit(1)

        sphereSource = vtk.vtkSphereSource()
        sphereSource.SetCenter(0.0, 0.0, 0.0)
        sphereSource.SetRadius(5.0)
        sphereSource.SetThetaResolution(12)
        sphereSource.SetPhiResolution(12)
        sphereSource.Update()

        # Render the sphere
        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputConnection(sphereSource.GetOutputPort())

        actor = vtk.vtkActor()
        darkTurquoise = list(map(lambda x: x / 255.0, [0.0, 206.0, 209.0]))
        actor.GetProperty().SetDiffuseColor(darkTurquoise)
        actor.GetProperty().SetSpecular(.5)
        actor.GetProperty().SetSpecularPower(20)
        actor.SetMapper(mapper)

        renderer = vtk.vtkRenderer()
        renderWindow = vtk.vtkRenderWindow()
        renderWindow.AddRenderer(renderer)
        renderWindow.SetSize(300, 300)

        renderWindowInteractor = vtk.vtkRenderWindowInteractor()
        renderWindowInteractor.SetRenderWindow(renderWindow)

        renderer.AddActor(actor)
        royalBlue = list(map(lambda x: x / 255.0, [65.0, 105.0, 225.0]))
        renderer.SetBackground(royalBlue)

        renderWindow.Render()
        renderWindow.SetWindowName('Sphere')

        # This needs to be set in Linux
        # On Linux VTK may not be defaulting to asking
        # for alpha bitplanes and  on Windows it does.
        # renderWindow.SetAlphaBitPlanes(1)

        # Take the partial screenshot
        windowToImageFilter = vtk.vtkWindowToImageFilter()
        windowToImageFilter.SetInput(renderWindow)
        windowToImageFilter.SetMagnification(1)
        if idx == 0:
            windowToImageFilter.SetInputBufferTypeToRGB()
        elif idx == 1:
            windowToImageFilter.SetInputBufferTypeToRGBA()

        windowToImageFilter.SetViewport(0.5, 0.5, 0.8, 1.0)
        windowToImageFilter.ReadFrontBufferOff()  # Read from the back buffer.
        windowToImageFilter.Update()

        # interact with the data.
        # renderWindowInteractor.Start()

        # -------------------------------
        # Save the image
        # -------------------------------
        imageWriter = vtk.vtkPNGWriter()
        imageWriter.SetInputConnection(windowToImageFilter.GetOutputPort())
        imageWriter.SetFileName(self.fn[idx])
        imageWriter.Write()

    def testWIFRGB(self):
        self.WritePNG(0)

        '''
            Read the png file and test it.
        '''
        # -------------------------------
        # Load the image
        # -------------------------------
        img_file = self.fn[0]
        reader = vtk.vtkPNGReader()
        reader.SetFileName(img_file)

        renderWindowInteractor = vtk.vtkRenderWindowInteractor()

        imageViewer = vtk.vtkImageViewer2()
        imageViewer.SetInputConnection(reader.GetOutputPort())
        imageViewer.SetupInteractor(renderWindowInteractor)
        imageViewer.Render()
        imageViewer.GetRenderWindow().SetWindowName('PNG Sphere')
        # interact with the data.
        # renderWindowInteractor.Start()
        vtk.test.Testing.compareImage(
            renderWindowInteractor.GetRenderWindow(),
            vtk.test.Testing.getAbsImagePath(img_file),
            threshold = 10)
        vtk.test.Testing.interact()

    def testWIFRGBA(self):
        self.WritePNG(1)

        '''
            Read the png file and test it.
        '''
        # -------------------------------
        # Load the image
        # -------------------------------
        img_file = self.fn[1]
        reader = vtk.vtkPNGReader()
        reader.SetFileName(img_file)

        renderWindowInteractor = vtk.vtkRenderWindowInteractor()

        imageViewer = vtk.vtkImageViewer2()
        imageViewer.SetInputConnection(reader.GetOutputPort())
        imageViewer.SetupInteractor(renderWindowInteractor)
        imageViewer.Render()
        imageViewer.GetRenderWindow().SetWindowName('PNG Sphere')
        # interact with the data.
        # renderWindowInteractor.Start()
        vtk.test.Testing.compareImage(
            renderWindowInteractor.GetRenderWindow(),
            vtk.test.Testing.getAbsImagePath(img_file),
            threshold = 10)
        vtk.test.Testing.interact()

if __name__ == '__main__':
    vtk.test.Testing.main([(TestWIFPNGWriteRead, 'test')])