#!/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')
    # uncomment following to set a specific view size
    # renderView1.ViewSize = [1578, 1078]

    # current camera placement for renderView1
    renderView1.CameraPosition = [12.870669565603832, -1.504489177014077, 1.2347805813350712]
    renderView1.CameraViewUp = [-0.09492203988135463, 0.00497155403877921, 0.9954722949410507]
    renderView1.CameraParallelScale = 3.369047677463545

    # create a new 'Clip'
    clip1 = Clip(Input=fides)
    Hide(fides, renderView1)

    # show data in view
    fidesDisplay = Show(clip1, 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

    # get opacity transfer function/opacity map for var
    varPWF = GetOpacityTransferFunction(varname)

    return renderView1


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


def ParseArgs():
    parser = argparse.ArgumentParser()
    parser.add_argument("-a", "--fides_attributes", help="path to bp file containing attributes needed by fides", 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
    step = 0
    renderView = None
    while True:
        status = NotReady
        while status == NotReady:
            fides.PrepareNextStep()
            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()

    comm = MPI.COMM_WORLD

    print("fides_attributes: ", args.fides_attributes)
    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.fides_attributes, ConvertToVTK=0)

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