#!/usr/bin/env python

import argparse
from mpi4py import MPI

#### import the simple module from the paraview
from paraview.simple import *

def SetupVis(fides, varname):
    # get active view
    renderView1 = GetActiveViewOrCreate('RenderView')
    renderView1.ResetCamera()

    # show data in view
    fidesDisplay = Show(fides, renderView1, 'UnstructuredGridRepresentation')

    # trace defaults for the display properties.
    fidesDisplay.Representation = 'Surface'

    # reset view to fit data
    renderView1.ResetCamera()

    # update the view to ensure updated data information
    renderView1.Update()

    # set scalar coloring
    ColorBy(fidesDisplay, ('POINTS', varname))

    # rescale color and/or opacity maps used to include current data range
    fidesDisplay.RescaleTransferFunctionToDataRange(True, False)

    # show color bar/color legend
    fidesDisplay.SetScalarBarVisibility(renderView1, True)

    # get color transfer function/color map for var
    varLUT = GetColorTransferFunction(varname)
    varLUT.RescaleOnVisibilityChange = 1

    # create a new 'Clip'
    clip1 = Clip(Input=fides)
    clip1.ClipType.Origin = [3.15, 3.15, 3.15]
    clip1.ClipType.Normal = [0.0692860629205188, -0.039817789099083185, -0.996801878587834]

    clip1Display = Show(clip1, renderView1, 'UnstructuredGridRepresentation')
    Hide(fides, renderView1)

    # show color bar/color legend
    clip1Display.SetScalarBarVisibility(renderView1, True)

    # update the view to ensure updated data information
    renderView1.Update()

    # current camera placement for renderView1
    renderView1.CameraPosition = [3.15, 3.15, -22.7944]
    renderView1.CameraViewUp = [0, 1, 0]

    return renderView1


def UpdateVis(renderView, output_path, step):
    # update the view to ensure updated data information
    renderView.Update()
    SaveScreenshot(output_path + '/gray-scott' + str(step) + '.png', renderView, ImageResolution=[1920, 1080])


def ParseArgs():
    parser = argparse.ArgumentParser()
    parser.add_argument("-a", "--bp_file", help="path to bp file", type=str)
    parser.add_argument("-o", "--output_path", help="path to save images rendered by paraview", type=str)
    parser.add_argument("-v", "--varname", help="name of variable to visualize", type=str)
    args = parser.parse_args()

    return args


OK = 0
NotReady = 1
EndOfStream = 2

def Streaming(fides, output_path, varname):
    # Should work for SST as well as streaming by BP file
    # Also works for reading through BP file that is already written,
    # but in that case you could also do random access
    step = 0
    renderView = None
    while True:
        status = NotReady
        while status == NotReady:
            # essentially calls BeginStep on ADIOS engine being used
            fides.PrepareNextStep()
            # let ParaView know we need to update the pipeline info,
            # so we can get the status of this step
            fides.UpdatePipelineInformation()
            status = fides.NextStepStatus
        if status == EndOfStream:
            print("ADIOS StepStatus is EndOfStream")
            return
        if step == 0:
            renderView = SetupVis(fides, varname)
        UpdateVis(renderView, output_path, int(step))
        step += 1


if __name__ == "__main__":
    args = ParseArgs()

    print("bp_file: ", args.bp_file)
    print("output_path: ", args.output_path)
    print("varname: ", args.varname)

    #### disable automatic camera reset on 'Show'
    paraview.simple._DisableFirstRenderCameraReset()

    # create a new 'FidesReader'
    fides = FidesReader(FileName=args.bp_file, ConvertToVTK=1)

    Streaming(fides, args.output_path, args.varname)
