#! /usr/bin/env python

import vtk
import time
import random
import sys
import os
from numpy import *
# from scipy.io import read_array

factGraphics = vtk.vtkGraphicsFactory();
factGraphics.SetUseMesaClasses(1)
del factGraphics

factImage = vtk.vtkImagingFactory;
factImage.SetUseMesaClasses(1)
del factImage

## Read the basic configuration of the mesh. 
## Later, its node coordinates are modified
## with a vtkProgrammableFilter
reader = vtk.vtkUnstructuredGridReader()
reader.SetFileName("beam.vtk")

## The filter to modify the coordinates of the mesh
change_coords_filter = vtk.vtkProgrammableFilter()
change_coords_filter.SetInputConnection(reader.GetOutputPort())

## Coordinate system: elevation is `y' axis. The worm displaces
## in the `z' direction. 

##----------------- BEGIN SETTINGS -------------------------------
## Position of worm node coordinates at undisturbed position
## (i.e. a plain cylinder)
# X0 = read_array(file("elastld.nod.tmp"))
X0 = loadtxt("elastld.nod.tmp")

# Scale displacements factor
xscale = 1.0
## Store frames as TIFF images in ./YUV directory
mkvid = 1
# Cycle or not on computed frames
cycle = 1
## Start frame, total nbr of frames
frame_start = 0
frame_end = 900
## Increment in frame number
finc = 10
## Show deformation data with colors
color_data = 1
## computed: 1.32,1.75
mindefo = 1.10
maxdefo = 0.90
## Directory where steps are found
# steps_dir = "./STEPS"
steps_dir = "./SAVED-STEPS/STEPS-2009-06-19-2ndmode-A08"
## Time step
Dt = 0.01
## Diameter of particles (for viz. only)
Dprtcl = 0.04
## Number of points in trail, number of frames
nstep = 100
nprtclref = 2

## Do not update state, only to test de camera
## and other settings
update_state = 1

## The flying camera is as follows.
## It rotates around the focal point but also the distance to the
## object varies sinusoidally between Rcammin and Rcammax
Ly = 5.0
Rcammin = 1.5*Ly
Rcammax = 2.2*Ly
## The speed and initial position of rotation of the camera
wcam = -1
## Initial position angle for the cam
wcam_phi0 = -65.0
## The speed with which the radius varies
wRcam = 1.33*wcam
## Height of camera
Hcam_fac = 0.75
## height of point where the camera points to
Hto = 3.0

## Close-up at upper view
if True:
    Rcammin = 0.8*Ly
    Rcammax = Rcammin
    wcam_phi0 = 10.0
    Hcam_fac = 2.5
    Hto = 0.9*Ly
    
##----------------- END SETTINGS -------------------------------

if mkvid:
    finc=1

if not update_state:
    color_data = 0

# Initial frame
frame = frame_start
# Dimension of problem
ndim = X0.shape[1]
# Max frame computed
maxframe = -1
def existz(afile):
    if os.path.exists(afile):
        return afile
    afilez = afile + ".gz"
    if os.path.exists(afilez):
        return afilez
    return None

DX = None

## Trail class
class trail():
    def __init__(self,nstep_a,nprtclref_a,datafile):
        self.ndim = 3
        self.nstep = nstep_a
        self.data = loadtxt(datafile)
        ## Number of dots injected in the trail per time step
        ## Must be an integer
        self.nprtclref = nprtclref_a
        ## This is the actual number of dots
        self.nprtcls = nprtclref * (self.nstep-1)
        ## Stores a random double for colouring
        self.color = []
        ## Initialize the `tprtcl' vector
        ## Stores injection time of particles in the trail.
        ## The time is given in number of time steps
        self.tprtcl = []
        self.nprtcls = (self.nstep-1)*self.nprtclref
        h = 1.0/float(self.nprtclref)
        for k in range(0,self.nprtcls):
            r = random.random()-0.5
            tp = (k+r)*h
            if tp<0.0:
                tp = 0.0
            if tp>self.nstep-1:
                tp = self.nstep-1
            self.tprtcl.append(tp)
            self.color.append(random.random())
        ## This is stores the current coordinates of
        ## `nstep' points in the current trail.
        ## They are kept with the `push()' method.
        self.coords = zeros((self.nstep,3),dtype=double)
        self.reset()

    ## Returne the internal container of
    ## the data needed to compute the next
    ## position of the trail
    def tdata(self):
        return self.data

    ## Set all points in the trail to a far position
    def reset(self):
        ## Initially set all trail pts at 
        ## the same place (far from the scene)
        xt1 = array([100.0,100.0,100.0])
        for k in range(0,self.nstep):
            self.coords[k] = xt1

    ## Insert a new position in the trail
    def push(self,xt):
        ## Update the coords vector. Rotate all
        ## the old positions one place to the back
        ## and put the new position in the front
        for k in range(0,self.nstep-1):
            kk = self.nstep-1-k
            self.coords[kk] = self.coords[kk-1]
        self.coords[0] = xt

    ## Update the coordinates in the VTK array
    ## to the positions stored internally
    def update(self,points):
        trdata = grid.GetPointData().GetArray(0)
        for k in range(0,self.nprtcls):
            ## This is the length from the
            ## injection point. 
            tp = self.tprtcl[k]

            ## Interpolate between the two consecutive
            ## positions, i.e.
            ## x = (1-alpha)*x[step] + alpha*x[step+1]
            step = math.floor(tp)
            alpha = tp-step
            ## Check not to overflow the coords vector
            assert step<self.nstep-1
            ## Take coords
            x0 = self.coords[step]
            x1 = self.coords[step+1]
            ## Interpolate
            xx = (1-alpha)*x0+alpha*x1
            ## Set the data in the VTK structures
            points.SetPoint(k,xx[0],xx[1],xx[2])
            trdata.SetValue(k,self.color[k])

## Construct the trail object
mytrail = trail(nstep,nprtclref,"elastld.trail.tmp")

## ------- THE MAIN UPDATING FUNCTION --------------------
change_coords_fun_frame_computed = -1
def change_coords_fun1():
    global DX,maxframe,points,mytrail
#     print "fun1: frame %d, last computed %d" \
#         % (frame,change_coords_fun_frame_computed)
#     print "in change_coords_fun(), frame %d" % frame
    input = change_coords_filter.GetUnstructuredGridInput()
    numPts = input.GetNumberOfPoints()

    output = change_coords_filter.GetUnstructuredGridOutput()
    output.CopyStructure(input)
    bar_points = output.GetPoints()
    
    ## Displacements for the maximum distorted
    ## position of the worm, i.e. the position
    ## of the coordinates of the worm are X = X0+alpha*DX
    ## with 0<=alpha<=1
    state = existz(steps_dir + "/elastld.state-%d.tmp" % frame)
    dframe = frame
    if cycle and maxframe<0 and state is None:
        maxframe = frame
        print "setting cycle maxframe to %d" % maxframe
    if cycle and maxframe>0:
        dframe = (frame % maxframe)
        if dframe==0:
            mytrail.reset()
        tryfile = steps_dir + "/elastld.state-%d.tmp" % dframe
        state = existz(tryfile)
    if state is None:
        print "not found state file for frame %d" % dframe
        sys.exit()
    if update_state or DX is None or random.randint(0,10)==0:
        DX = loadtxt(state)

    ## Change the coordinates point by point
    for i in range(0, numPts):
        xx = X0[i] + xscale*DX[i,0:ndim]
        if ndim==2:
            xx[2] = 0.0
        bar_points.SetPoint(i, xx[0], xx[1], xx[2])

    if color_data:
        defo = vtk.vtkDoubleArray()
        defo.SetName("deformation")
        dfile = existz(steps_dir + "/elastld.data-%d.tmp" % dframe)
        if dfile is None:
            print "Couldn open file %s" % dfile
            sys.exit()
        d = loadtxt(dfile)
        for i in range(0, numPts):
            # val = 10*(d[i]-1.0)
            val = (d[i]-mindefo)/(maxdefo-mindefo)
            defo.InsertNextValue(val)
        output.GetPointData().SetScalars(defo)

    t = frame*Dt
    omega = 4
    
    ## Compute the coordinates of the actual center of the trail
    trail_data = mytrail.tdata()
    npt = trail_data.shape[0]
    xt = array([0.0,0.0,0.0])
    wsum = 0.0
    for k in range(0,npt):
        node = int(trail_data[k,0])
        w = trail_data[k,1]
        wsum += w
        x1 = X0[node-1,0:3]+xscale*DX[node-1,0:3]
        ## print "k %d, node %d, w %f, x %s" % (k,node,w,x1)
        xt += w*x1
    xt /= wsum
    ## print "wsum %f, xt %s" % (wsum,xt)
    mytrail.push(xt)
    mytrail.update(points)

    ## Update the coords vector in the grid
    grid.Modified()

## This function is the kernel of the `vtkProgrammableFilter'
## It changes the coordinates of the worm points
def change_coords_fun():
    global change_coords_fun_frame_computed 
#     print "fun: frame %d, last computed %d" \
#         % (frame,change_coords_fun_frame_computed)
    if change_coords_fun_frame_computed<0 \
           or frame != change_coords_fun_frame_computed:
        change_coords_fun_frame_computed = frame
        change_coords_fun1()

change_coords_filter.SetExecuteMethod(change_coords_fun) 

# Create mapper
ugmapper = vtk.vtkDataSetMapper()
ugmapper.SetInputConnection(change_coords_filter.GetOutputPort())
if not color_data:
    ugmapper.ScalarVisibilityOff()
else:
    print "using colored data"
    ugmapper.ScalarVisibilityOn()
    ugmapper.SetScalarRange(0.0,1.0)
    ugmapper.SetScalarModeToUsePointFieldData()
    ugmapper.ColorByArrayComponent("deformation", 0)
#     colorTransferFunction = vtk.vtkColorTransferFunction()
#     colorTransferFunction.AddRGBPoint(mindefo, 0.0, 0.0, 1.0)
#     ## colorTransferFunction.AddRGBPoint(, 1.0, 1.0, 0.0)
#     colorTransferFunction.AddRGBPoint(maxdefo, 1.0, 0.0, 0.0)
#     ugmapper.SetLookupTable(colorTransferFunction)

# Create actor
ugactor = vtk.vtkActor()
ugactor.SetMapper(ugmapper)

prop = ugactor.GetProperty()
# prop.SetColor(colorTransferFunction)

if False:
    print "using wireframe visualization"
    prop.SetRepresentationToWireframe()
    prop.SetDiffuseColor(1, 0, 0)
    # ugmapper.ScalarVisibilityOff()

# --------------- DRAWING THE TRAIL ---------
treader = vtk.vtkUnstructuredGridReader()
treader.SetFileName("trail.vtk")

## The filter to modify the coordinates of the mesh
tchange_coords_filter = vtk.vtkProgrammableFilter()
tchange_coords_filter.SetInputConnection(treader.GetOutputPort())

def tchange_coords_fun():
    global change_coords_fun_frame_computed
#     print "tfun: frame %d, last computed %d" \
#           % (frame,change_coords_fun_frame_computed)
    if change_coords_fun_frame_computed<0 \
           or frame != change_coords_fun_frame_computed:
        change_coords_fun_frame_computed = frame
        change_coords_fun1()

tchange_coords_filter.SetExecuteMethod(tchange_coords_fun) 
tugmapper = vtk.vtkDataSetMapper()
tugmapper.SetInputConnection(tchange_coords_filter.GetOutputPort())
tugmapper.ScalarVisibilityOff()

# Create actor
tugactor = vtk.vtkActor()
tugactor.SetMapper(tugmapper)
tugactor.GetProperty().SetDiffuseColor(0, 1, 1)

## ------------ A VERSION OF THE TRAIL WITH PARTICLES ---------
## The grid of positions
points = vtk.vtkPoints()
ids = vtk.vtkIdList()
nprtcls = (nstep-1)*nprtclref
for k in range(0,nprtcls):
    points.InsertPoint(k,0.0,0.0,0.0)
    ids.InsertNextId(k)
grid = vtk.vtkUnstructuredGrid()
grid.Allocate(nprtcls,100)
grid.InsertNextCell(12,ids)
grid.SetPoints(points)
trdata = vtk.vtkDoubleArray()
trdata.SetNumberOfValues(nprtcls)
idx = grid.GetPointData().AddArray(trdata)

# Create a sphere to use as a glyph source for vtkGlyph3D.
Sphere = vtk.vtkSphereSource()
Sphere.SetRadius(Dprtcl/2.0)
## We use low resolution so as to allow a large
## number of particles with little CPU effort
Sphere.SetPhiResolution(4) 
Sphere.SetThetaResolution(4)

# The sphere is cloned at each grid position with vtkGlyph3D
Vertices = vtk.vtkGlyph3D()
Vertices.SetInput(grid)
Vertices.SetSource(Sphere.GetOutput())

# The mapper and actor to display the glyphs.
vmapper = vtk.vtkPolyDataMapper()
vmapper.SetInputConnection(Vertices.GetOutputPort())
# vmapper.ScalarVisibilityOff()

vmapper.ScalarVisibilityOn()
vmapper.SetScalarRange(0.0,1.0)
vmapper.SetScalarModeToUsePointFieldData()
vmapper.SetColorModeToMapScalars()
vmapper.SelectColorArray(0)

particles = vtk.vtkActor()
particles.SetMapper(vmapper)
# particles.GetProperty().SetDiffuseColor(1,0,0)

# ----------- DRAW AXES ---------------
# A cylinder to model the reactor
cyl = vtk.vtkCylinderSource()
cyl.SetHeight(4*Ly)
cyl.SetRadius(0.005*Ly)
cyl.SetResolution(30)

cylMapper = vtk.vtkPolyDataMapper()
cylMapper.SetInputConnection(cyl.GetOutputPort() )

## Create the cylinder actor
cylY = vtk.vtkActor()
cylY.SetMapper(cylMapper)
## cylY.RotateX(90.0)
## cylY.SetPosition(0.0,yc,0.0)

cylX = vtk.vtkActor()
cylX.SetMapper(cylMapper)
cylX.RotateZ(90.0)

cylZ = vtk.vtkActor()
cylZ.SetMapper(cylMapper)
cylZ.RotateX(90.0)

cylX.GetProperty().SetColor(1.0,0.0,0.0)
cylY.GetProperty().SetColor(0.0,1.0,0.0)
cylZ.GetProperty().SetColor(0.0,0.0,1.0)

## ----------- CREATE THE USUAL RENDERING STUFF ---------------
ren = vtk.vtkRenderer()
renWin = vtk.vtkRenderWindow()
renWin.AddRenderer(ren)
renWin.SetSize(640,480)
renWin.OffScreenRenderingOn();
print "trace 0"
# renWin.SetSize(1280,720)

ren.SetBackground(.4, .4, .4)
ren.AddActor(ugactor)
#ren.AddActor(tugactor)
ren.AddActor(particles)
ren.AddActor(cylX)
ren.AddActor(cylY)
ren.AddActor(cylZ)

## ----------- WRITING TO FILE -----------------
## This is used to store the frames
## for creating a movie
print "trace 1"
w2i = vtk.vtkWindowToImageFilter()
print "trace 2"
w2i.SetInput(renWin)
print "trace 3"
w2i.Modified()
w2i.Update()

## The TIFF writer
writer = vtk.vtkTIFFWriter()
writer.SetInputConnection(w2i.GetOutputPort())
writer.SetCompressionToJPEG()
    
## ----------- MORE CAMERA SETTINGS -----------------
## Initialize camera
cam = ren.GetActiveCamera()
cam.SetFocalPoint(0.,Hto,0.)
cam.SetViewUp(0.,1.,0.);

## This is added so that it gives time to set
## no border in the OpenGL window and other stuff
## like minimizing other windows. 
if mkvid:
    renWin.Render()
    raw_input("Enter something to continue: ")

## --------  THE MAIN FRAME LOOP ----------------------
# Loop while: rotating the camera and modify
# node coordinates
t = 0.0
while 1:
    t += finc*Dt
    if not mkvid:
        time.sleep(0.3)
    if frame % 10 == 0:
        print "frame: %d [%s]" % (frame,time.asctime())

    ## This recomputes the clipping plane.
    ## Otherwise, since the camera
    ## is rotating, some objects may disappear
    ren.ResetCameraClippingRange()
    change_coords_filter.Modified()
    tchange_coords_filter.Modified()

    ## Update camera
    phi = wcam_phi0*pi/180.0 + wcam*t
    Rcamav =(Rcammin+Rcammax)/2.0
    DRcam = (Rcammin-Rcammax)/2.0
    Rcam = Rcamav + DRcam*sin(wRcam*t)
    Hcam = Hcam_fac*Rcam
    cam.SetPosition(Rcam*cos(phi),Hcam,Rcam*sin(phi))
#     phielev = 80.0*pi/180.0
#     cam.SetPosition(0.,Rcam*sin(phielev),Rcam*cos(phielev))
    renWin.Render()

    ## Save current frame
    if mkvid:
        assert os.path.isdir("./YUV")
        w2i.Modified()
        tiff = "./frame.tiff"
        yuv = "./YUV/frame." + str(frame) +".yuv"
        writer.SetFileName(tiff)
        writer.Write()
        os.system("convert %s %s ; gzip -f %s" % (tiff,yuv,yuv))
        os.unlink(tiff)

    if frame_end>0 and frame >= frame_end:
        break
    
    ## Update frame counter
    frame += finc