#!/usr/bin/env python3

import argparse
import warnings

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

def main():
	parser = argparse.ArgumentParser(description='Plot columnated data')
	parser.add_argument('-F', default=' ', type=str, help='set field delimiter')
	parser.add_argument('INFILE', nargs='+', type=str, help='input files')
	args = parser.parse_args()

	# FIXME currently really don't care about handling multiple lines properly
	# but this should be x, *ys
	frames = []
	print(args.INFILE)
	for path in args.INFILE:
		with open(path) as f:
			x, y = read_columns(f, args.F, diefunc=parser.error)
			frames.append((x, y))

	x0, y0 = frames[0]
	fig, ax = plt.subplots()
	line, = ax.plot(x0, y0)

	def animate(i):
		x,y = frames[i]
		line.set_xdata(x) # update the data
		line.set_ydata(y) # update the data
		return line,

	# Init only required for blitting to give a clean slate.
	def init():
		line.set_ydata(np.ma.array(x, mask=True))
		return line,

	animation.FuncAnimation(fig, animate, np.arange(len(frames)), init_func=init,
	                              interval=250, blit=True)
	plt.show()

def read_columns(f, fieldsep, diefunc):
	# Handle most of the parsing with a simple generator that
	# produces lists of floats for each row of data
	def inner():
		for line in f:
			line = line.strip()
			if not line:
				continue
			strs = line.split(fieldsep)
			yield list(map(float, strs))

	rows = list(inner())
	if not rows:
		diefunc('Input data was empty!')
		return []

	len0 = len(rows[0])
	assert len0 > 0, "Empty lines were skipped"
	if not all(len(x)==len0 for x in rows):
		warnings.warn('Mismatched row lengths! Incomplete columns will be ignored')

	cols = list(zip(*rows))
	return cols

if __name__ == '__main__':
	main()
