#!/usr/bin/python

import pencil as pc
P=pc.P
N=pc.N
from for_pencil import varmax

dim=pc.read_dim()
param=pc.read_param(quiet=True)
param2=pc.read_param(param2=True,quiet=True)
index=pc.read_index(param=param)
grid=pc.read_grid(param=param,quiet=True,trim=True)
vmax=varmax()

var1=0
nvar=vmax+1

P.ion()

P.figure(figsize=(8,8),dpi=32)
P.axis('off')
P.subplots_adjust(bottom=0,top=1,left=0,right=1)
frame = grid.x.min(), grid.x.max(), grid.y.min(), grid.y.max()

# for superimposing the velocity field
qs1 = N.random.random_integers(0, dim.nx-1, 1000)
qs2 = N.random.random_integers(0, dim.ny-1, 1000)
xx, yy = N.meshgrid(grid.x, grid.y)

for ivar in range(var1,vmax+1): 
  var=pc.read_var(ivar=ivar,trimall=True,run2D=param.lwrite_2d,param=param,dim=dim,index=index,quiet=True)
  print var.t
  if (ivar==var1):
    im=P.imshow(var.ux,origin='lower',extent=frame)
    ux = var.ux ; uy = var.uy
    a=ux**2+uy**2 ; cte=N.sqrt(a).max()
    vel = P.quiver(xx[qs2,qs1], yy[qs2,qs1], \
    ux[qs2,qs1]/cte, uy[qs2,qs1]/cte, scale=7)
  else:
    im.set_data(var.ux)
    ux = var.ux ; uy = var.uy
    a = ux**2+uy**2 ; cte=N.sqrt(a).max()
    vel.set_UVC(ux[qs2,qs1]/cte, uy[qs2,qs1]/cte)
    P.draw()

P.ioff()