#!/usr/bin/python

# Written by Andrew Brown
# This is just a quick and dirty script to iterate through
# files plotting them in gnuplot.  Have fun!

import os,sys,curses,pickle,time,md5
(gp,gpoutput) = os.popen4("gnuplot","w")
gpoutput.close()

defaults = {'params':'','filename':'','step':1,'zeros':1}

if len(sys.argv) > 1:
  try:
    f = open(sys.argv[1],'r')
    options = pickle.load(f)
    f.close()
  except:
    options = defaults
  # Make sure all default keys are there
  for key in defaults.iterkeys():
    if not key in options.keys():
      options[key] = defaults[key]
else:
  options = defaults

def plot(stdscr, number):
  if len(options['filename']) > 0:
    # Check that the file exists
    try:
      open(options['filename'].replace("%n",str(number).zfill(options['zeros'])),'r')
    except IOError:
      curses.beep()
      return 1
  params = options['params'].replace('%f',options['filename'])
  params = params.replace('%n',str(number).zfill(options['zeros']))
  stdscr.move(1,0)
  stdscr.clrtoeol()
  stdscr.move(2,0)
  stdscr.clrtoeol()
  stdscr.addstr(1,0,params)
  gp.write(params + "\n")
  gp.flush()

def go(stdscr):
  curses.curs_set(0)
  #stdscr.keypad(1)
  startplot = 1

  title = "Welcome.  Press Enter to re-plot, arrow keys move filenumber."
  plotnumber = 0
  while 1:
    # Menu
    stdscr.addstr(0,5,title)
    stdscr.move(3,0)
    stdscr.clrtobot()
    stdscr.addstr(3,1,"Menu",curses.A_UNDERLINE)
    stdscr.addstr(4,1,"1) Set Filename (" + options['filename'] + ")")
    stdscr.addstr(5,1,"2) Set Parameters")
    if options['params'] != '':
      stdscr.addstr(6,3,options['params'][:stdscr.getmaxyx()[1]-3])
    else:
      stdscr.addstr(6,5,"[none set]  (this must be set before you can plot)")
    stdscr.addstr(7,1,"3) Jump To File number. (currently plotting #%s)" % str(plotnumber))
    stdscr.addstr(8,1,"4) Set File Step (currently at " + str(options['step']) + ")")
    if len(sys.argv) > 1:
      stdscr.addstr(9,1,"5) Save parameters to " + sys.argv[1])
    else:
      stdscr.addstr(9,1,"[Save parameters disabled, start program with a file for saving options]")
    stdscr.addstr(10,1,"6) Length of number (padding) (currently at " + str(options['zeros']) + ")")
    stdscr.addstr(11,1,"7) Export sequence as MPEG movie (requires ImageMagick and mpeg2vidcodec)")
    stdscr.addstr(12,1,"0) Quit")
    stdscr.addstr(13,1,"[Enter] Plots the current step")
    stdscr.refresh()
    
    input = stdscr.getch()
    
    if input == ord('7'):
      stdscr.move(3,0)
      stdscr.clrtobot()
      stdscr.addstr(3,0,"This will compile an mpeg movie using the current gnuplot parameters.")
      stdscr.addstr(4,0,"Please enter the starting file number for the sequence: ")
      curses.echo()
      curses.curs_set(1)
      start = stdscr.getstr()
      if not start.isdigit():
        curses.noecho()
        curses.curs_set(0)
        continue
      stdscr.addstr(5,0,"Please enter the ending number: ")
      end = stdscr.getstr()
      if not end.isdigit():
        curses.noecho()
        curses.curs_set(0)
        continue
      stdscr.addstr(6,0,"Please enter the movie name: ")
      movname = stdscr.getstr()
      if movname[-4:].lower() != ".mpg":
        movname += ".mpg"
      
      curses.noecho()
      curses.curs_set(0)
      start = int(start)
      end = int(end)
      stdscr.addstr(7,0,"Working...this will take a while")
      
      # Grab a temporary directory to work with
      #tmppath = os.tempnam() # Is a security risk
      tmppath = "/tmp/movefiles_" + md5.new(str(time.time())).hexdigest() # Always unique
      os.mkdir(tmppath)
      filesfile = open(tmppath + "/files","w")
      todelete = [tmppath + "/files"]
      stdscr.addstr(8,0,"Exporting frames...")
      stdscr.refresh()
      gp.write("set term png\n")
      frame = 0
      for x in range(start,end+1,options['step']):
        gp.write('set output "' + tmppath + '/output' + str(frame) + '.png"\n')
        params = options['params'].replace('%f',options['filename'])
        params = params.replace('%n',str(x).zfill(options['zeros']))
        gp.write(params + "\n")
        gp.flush()
        todelete.append(tmppath + "/output" + str(frame) + ".png")
        filesfile.write(tmppath + "/output" + str(frame) + ".png\n")
        frame += 1
      filesfile.close()
      stdscr.addstr(8,19,"Done!")
      gp.write("set term x11\n")
      gp.flush()
      
      # Now merge everything into a mov
      stdscr.move(9,0)
      stdscr.clrtoeol()
      stdscr.addstr("Converting to MPG movie...")
      stdscr.refresh()
      time.sleep(1)
      imagemagick = os.popen4('convert @' + tmppath + "/files " + movname,'r')[1]
      while 1:
        line = imagemagick.readline()
        if not line:
          break
        stdscr.move(10,0)
        stdscr.clrtoeol()
        stdscr.addstr(10,0,line.strip())
        stdscr.refresh()
      status = imagemagick.close()
      
      # Clean up files
      for x in todelete:
        os.unlink(x)
      os.rmdir(tmppath)
      
      if status:
        stdscr.addstr(11,0,"Imagemagick exited with abnormal status.  Movie creation failed")
      else:
        stdscr.addstr(11,0,"All done!")
      stdscr.addstr(12,0,"Press any key to continue")
      stdscr.getch()
      
      
      
      
      
    elif input == ord('1') or input == curses.KEY_C1:
      stdscr.move(2,0)
      stdscr.clrtoeol()
      stdscr.move(1,0)
      stdscr.clrtoeol()
      stdscr.addstr(1,0,"Enter filename. (%n) represents the file number.  This is the name of the file")
      stdscr.addstr(2,0,"to check for existance every step and is optional: ")
      curses.echo()
      curses.curs_set(1)
      a = stdscr.getstr()
      curses.noecho()
      curses.curs_set(0)
      options['filename'] = a
      stdscr.move(1,0)
      stdscr.clrtoeol()
      stdscr.move(2,0)
      stdscr.clrtoeol()
      
    elif input == ord('2'):
      #request parameters
      #curses.echo()
      #stdscr.keypad(1)
      xlength = stdscr.getmaxyx()[1] # constant (unless window changes)
      pos = len(options['params'])
      length = pos
      stdscr.move(1,0)
      stdscr.clrtobot()
      stdscr.addstr(3,0,"Edit GNUplot Parameters Here.  '%f' represents the filename, '%n' represents\nthe file number.")
      for q in range(0,length,xlength):
        stdscr.addstr(6+(q/xlength),0,options['params'][q:q+xlength])
      stdscr.move(6+pos/xlength,pos-(pos/xlength)*xlength)
      curses.curs_set(1)
      keys = ""
      while 1:
        c = stdscr.getch()
        
        #xlength = stdscr.getmaxyx()[1] # constant (unless window changes)
        if False:
          stdscr.move(8,5)
          stdscr.clrtoeol()
          keys += " " + str(c)
          stdscr.addstr(8,0,keys)
          stdscr.move(6,pos)
        
        
        if c == 10:
          break
        elif c == curses.KEY_LEFT:
          if pos != 0:
            pos -= 1
            stdscr.move(6+pos/xlength,pos-(pos/xlength)*xlength)
        elif c == curses.KEY_RIGHT:
          if pos != length:
            pos += 1
            stdscr.move(6+pos/xlength,pos-(pos/xlength)*xlength)
        elif c == curses.KEY_BACKSPACE or c == 127:
          if pos != 0:
            pos -= 1
            length -= 1
            chrs = ""
            tryng = 0
            while 1:
              newchrs = stdscr.instr(6+tryng,0)
              if newchrs.strip() == "":
                chrs = chrs.strip()
                break
              chrs += newchrs
              stdscr.move(6+tryng,0)
              tryng += 1
              stdscr.clrtoeol()
            chrs = chrs[:pos]+chrs[pos+1:]
            for q in range(0,length,xlength):
              stdscr.addstr(6+(q/xlength),0,chrs[q:q+xlength])
            stdscr.move(6+pos/xlength,pos-(pos/xlength)*xlength)
        elif (c == curses.KEY_DC):
          if pos != length:
            length -= 1
            chrs = ""
            tryng = 0
            while 1:
              newchrs = stdscr.instr(6+tryng,0)
              if newchrs.strip() == "":
                chrs = chrs.strip()
                break
              chrs += newchrs
              stdscr.move(6+tryng,0)
              tryng += 1
              stdscr.clrtoeol()
            chrs = chrs[:pos]+chrs[pos+1:]
            for q in range(0,length,xlength):
              stdscr.addstr(6+(q/xlength),0,chrs[q:q+xlength])
            stdscr.move(6+pos/xlength,pos-(pos/xlength)*xlength)
        elif c == curses.KEY_HOME:
          pos = 0
          stdscr.move(6,0)
        elif c == curses.KEY_END:
          pos = length
          stdscr.move(6+pos/xlength,pos-(pos/xlength)*xlength)
        elif c >= 0 and c <= 255:
          chrs = ""
          tryng = 0
          while 1:
            newchrs = stdscr.instr(6+tryng,0)
            if newchrs.strip() == "":
              chrs = chrs.strip()
              break
            chrs += newchrs
            stdscr.move(6+tryng,0)
            tryng += 1
            stdscr.clrtoeol()
          if pos > len(chrs):
            # If inputting something in the middle of a blank line
            chrs = chrs + " " * (pos - len(chrs)) + chr(c)
          else:
            chrs = chrs[:pos] + chr(c) + chrs[pos:]
          length += 1
          pos += 1
          for q in range(0,length,xlength):
            stdscr.addstr(6+(q/xlength),0,chrs[q:q+xlength])
          stdscr.move(6+pos/xlength,pos-(pos/xlength)*xlength)
        
      chrs = ""
      tryng = 0
      while 1:
        newchrs = stdscr.instr(6+tryng,0)
        if newchrs.strip() == "":
          chrs = chrs.strip()
          break
        chrs += newchrs
        tryng += 1
      options['params'] = chrs
      #curses.noecho()
      curses.curs_set(0)
      stdscr.clear()
      
    elif input == ord('0'):
      #quit
      break
      
      
      
    elif input == curses.KEY_RIGHT:
      if len(options['params']) > 0:
        if not plot(stdscr,plotnumber + options['step']):
          plotnumber += options['step']
    elif input == curses.KEY_LEFT:
      if len(options['params']) > 0:
        if not plot(stdscr,plotnumber - options['step']):
          plotnumber -= options['step']
    elif input == 10:
      if len(options['params']) > 0:
        plot(stdscr,plotnumber)
      
      
      
    elif input == ord('3'):
      stdscr.move(2,0)
      stdscr.clrtoeol()
      stdscr.addstr(2,0,"Type file number: ")
      curses.echo()
      curses.curs_set(1)
      a = stdscr.getstr(2,18)
      curses.noecho()
      curses.curs_set(0)
      if a.isdigit():
        plotnumber = int(a)
        if startplot == 1: plot(stdscr,plotnumber)
        stdscr.move(2,0)
        stdscr.clrtoeol()
      else:
        stdscr.addstr(2,0,"     Non-integer given.  Number was not changed")
        
    elif input == ord('4'):
      stdscr.move(2,0)
      stdscr.clrtoeol()
      stdscr.addstr(2,0,"Type new file step: ")
      curses.echo()
      curses.curs_set(1)
      a = stdscr.getstr(2,20)
      curses.noecho()
      curses.curs_set(0)
      if a.isdigit():
        options['step'] = int(a)
        stdscr.move(2,0)
        stdscr.clrtoeol()
      else:
        stdscr.addstr(2,0,"     Non-integer given.  Step was not changed")
    elif input == ord('6'):
      stdscr.move(2,0)
      stdscr.clrtoeol()
      stdscr.addstr(2,0,"Length of number (0 for no padding): ")
      curses.echo()
      curses.curs_set(1)
      a = stdscr.getstr()
      curses.noecho()
      curses.curs_set(0)
      if a.isdigit():
        options['zeros'] = int(a)
        stdscr.move(2,0)
        stdscr.clrtoeol()
      else:
        stdscr.addstr(2,0,"     Non-integer given.  No changes were made")
    elif input == ord('5'):
      if len(sys.argv) > 1:
        f = open(sys.argv[1],'w')
        pickle.dump(options,f)
        f.close()
        stdscr.move(2,0)
        stdscr.clrtoeol()
        stdscr.addstr(2,0,"     Data options saved")
        
    
    else:
      stdscr.move(2,0)
      stdscr.clrtoeol()
      #stdscr.addstr(2,0,str(input))
  # Loop back to main menu

curses.wrapper(go)
gp.close()