#! /usr/bin/python

# Program to Demonstrate Steady, Conservative Flow with Non-Constant Density

# Reference: Some hints on how to make videos:
#   http://vpython.org/contents/FAQ.html

from __future__ import print_function, division
from visual import *
import os               # needed for catching ^C
import traceback        # needed for catching ^C
import sys
import colorsys
import signal
import time             # for time.sleep
import povexport

def keyInput(evt):
     s = evt.key
     print("[", s, "]")

scene = display(title='Whatever',
     x=0, y=0, width=854, height=480,
     center = (0, 3, 0), fov=.01, range=10)

if 0:
  print("vpython version: ", version)
  try:
      scene.bind('keydown', keyInput)
  except:
      print("Scene.bind", "requires vpython version 6")
      # continue without it

# main:
try:
    realtime = 0
    out_tmax = 15       # seconds
    args = sys.argv
    progname = args.pop(0)
    print(progname)
    myradius = 0.1
    myrate = 30.
    dy = 2.
    heap = []

    nobj = 50
    xmin = -9
    xmax = 9
    zone = (xmax - xmin) / 3.
    x0 = xmin
    x1 = xmin + zone
    x2 = x1 + zone
    x3 = xmax
    dyplus = dy + 2*myradius
    vx = (xmax-xmin)*myrate/200.
    if 0:
      vx /= 2
      myrate /= 2
    dtx = 1. / myrate           # ideal dt, may drift relative to walltime

    floor01 = box(pos=((x0+x1)/2., dy/2, 0), size=(x1-x0, dyplus, .1),
        color=(.2, .2, 0))

    floor12 = box(pos=((x1+x2)/2., dy/2, 0), size=(x2-x1, dyplus, .1),
        color=(0, .2, 0))

    floor23 = box(pos=((x2+x3)/2., dy/2, 0), size=(x3-x2, dyplus, .1),
        color=(.2, .2, 0))

    density = nobj / (xmax - xmin)
    nmake = density * vx * dtx           # creations per frame
    turnover = nobj / nmake * dtx
    print("nmake: ", nmake, "  turnover: ", turnover)
    i0 = 0
    i1 = (0.40 * nobj)
    i2 = (0.60 * nobj)
    i3 = nobj
    for ii in range(0, nobj):
      if ii < i1:
        xpos = x0 + (x1-x0)/(i1-i0)*(ii-i0)     # 40%
      elif ii < i2:
        xpos = x1 + (x2-x1)/(i2-i1)*(ii-i1)     # 20%
      else:
        xpos = x2 + (x3-x2)/(i3-i2)*(ii-i2)     # 20%
        xpos = random.uniform(x2, x3)           # remaining 40%
      ypos = random.uniform(0., dy)
      mycolor = colorsys.hsv_to_rgb(random.uniform(0., 1.), .9, .9)
      heap.append(sphere(pos=(xpos, ypos, 0), radius=myradius,
           color=mycolor))

#           material=materials.emissive,

    invisible = []
    tomake = 0.;

    outputting = 1

    sun1 = distant_light(direction=( .5, .5, 0), color=color.white)
    sun2 = distant_light(direction=(-.5, .5, 0), color=color.white)

# the big loop:
    frameno = 0
    walltime = time.time()
    while 1:
      rate(myrate)
      if realtime:
        newtime = time.time()
        dt = newtime - walltime
        if dt < 0.1/myrate:
          continue
        walltime = newtime
      else:
        dt = dtx

      for id in range(0, len(heap)):
        vxprime = vx
        if heap[id].pos[0] >= x1 and heap[id].pos[0] < x2:
          vxprime *= 2
        heap[id].pos[0] += vxprime*dt
        if heap[id].pos[0] > x3 and heap[id].visible:
          heap[id].visible = 0
          invisible.append(id)
      tomake += nmake;
      while tomake >= 1:
        tomake -= 1
        mycolor = colorsys.hsv_to_rgb(random.uniform(0., 1.), .9, .9)
        if len(invisible):
          #xx print("hi there", len(invisible), len(heap))
          id = invisible.pop()
          heap[id].pos[0] = xmin
          heap[id].pos[1] = random.uniform(0., dy)
          heap[id].color = mycolor
          heap[id].visible = 1
        else:
          print("making more")
          ypos = random.uniform(0., dy)
          heap.append(sphere(pos=(xpos, ypos, 0), radius=myradius,
               color=mycolor))

#                material=materials.emissive,

      if outputting:
        povexport.export(display=scene, filename="tmp-ray/x%05d.pov"%(frameno))
        if frameno >= out_tmax * myrate:
          print("outputting done")
          outputting = 0

      frameno += 1
      0-0       # end big loop


except KeyboardInterrupt as e:
  print ('...Interrupt...')
  os._exit(1)

except:
  #xxx print("Unexpected error:\n", sys.exc_info())
  print("Unexpected error... ", end='')
  sys.stdout.flush()

# do whatever other cleanup is needed ....

# Use print_exc() instead of raise, because we want to retain control,
# so we can force an exit that actually exits,
# rather than hanging around with a graphics window open.
  traceback.print_exc()
  os._exit(1)