import subprocess
import threading
from os import path,remove,devnull
from sys import stdin,stdout
import numpy as np

class VideoSink(object) :
	"""
	VideoSink: write numpy array to a movie file
	------------------------------------------------------------------------

	Requires mencoder:
	<http://www.mplayerhq.hu/design7/dload.html>

	Parameters:
	----------------
	filename	string	The path/name of the output file

	size		tuple	The row/column dimensions of the output movie

	rate		scalar	The framerate (fps)

	colorspace	string	The color space of the output frames, 8-bit RGB 
				by default ('$ mencoder -vf format=fmt=help' 
				for a list of valid color spaces)

	codec		string	The codec to use, libavcodecs by default 
				('$ mencoder -ovc -h' for a list of valid codecs)

	Methods:
	----------------
	VideoSink(array)	Write the input array to the specified .avi file
				- must be in the form [rows,columns,(channels)],
				and (rows,colums) must match the 'size' 
				specified when initialising the VideoSink.

	VideoSink.close()	Close the .avi file. The file *must* be closed 
				after writing, or the header information isn't 
				written correctly

	Example useage:
	----------------
	frames = np.random.random_integers(0,255,size=50,100,200,3).astype('uint8')
	vsnk = VideoSink('mymovie.avi',size=frames.shape[1:3],colorspace='rgb24')
	for frame in frames:
		vsnk(frame)
	vsnk.close()

	Alistair Muldal, Aug 2012

	Credit to VokkiCodder for this idea
	<http://vokicodder.blogspot.co.uk/2011/02/numpy-arrays-to-video.html>

	"""
	def __init__( self, filename='output.avi', size=(512,512), rate=25, colorspace='rgb24',codec='lavc'):

		# row/col --> x/y by swapping order
		self.size = size[::-1]

		cmdstring  = (	'mencoder',
				'/dev/stdin',
				'-demuxer', 'rawvideo',
				'-rawvideo', 'w=%i:h=%i'%self.size+':fps=%i:format=%s'%(rate,colorspace),
				'-o', filename,
				'-ovc', codec,
				'-nosound',
				'-really-quiet'
				)
		self.p = subprocess.Popen(cmdstring, stdin=subprocess.PIPE, shell=False)

	def __call__(self, frame) :
		assert frame.shape[0:2][::-1] == self.size
		# frame.tofile(self.p.stdin) # should be faster but it is indeed slower
		self.p.stdin.write(frame.tostring())
	def close(self) :
		self.p.stdin.close()
		self.p.terminate()

class VideoSource(object):
	"""
	VideoSource: create numpy arrays from frames in a movie file
	------------------------------------------------------------------------

	Requires mencoder:
	<http://www.mplayerhq.hu/design7/dload.html>

	Parameters:
	----------------
	filename	string	The path/name of the output file.

	colorspace	string	The color space of the output frames, 8-bit RGB 
				by default*.

	cache		string	Temporary cache for the raw frames,
				'/tmp/tmpframes_<filename>' by default. This 
				file is deleted when .close() is called.

	blocking	bool	If True, VideoSource will block control until 
				all of the source frames have been decoded 
				(useful if VideoSource is being created within a
				function). By default it will return control 
				immediately, but frames will not be accessible 
				until decoding is complete.
	
	*Only the following colorspaces are currently supported:
	8-bit monochrome: 	'y800','y8'
	8-bit RGB:		'rgb24','bgr24'
	8-bit RGBA:		'rgba','argb','rgb32','bgra','abgr','bgr32'

	Methods:
	----------------
	VideoSource[1:10:2]	Grab frames from the source movie as numpy 
				arrays. Currently only a single set of indices 
				is supported. It is also possible to iterate
				over frames.

	VideoSource.close()	Closes and deletes the temporary frame cache

	VideoSink.decode()	Decode the frames again and create a new cache.
				Do this if you called .close() and want to 
				access the frames again.

	Example useage:
	----------------
	vsrc = VideoSource('mymovie.avi',blocking=True)
	oddframes = vsrc[1::2]			# get item(s)
	framelist = [frame for frame in vsrc]	# iterate over frames
	vsrc.close()				# close and delete cache

	Alistair Muldal, Aug 2012

	Inspired by VokkiCodder's VideoSink class
	<http://vokicodder.blogspot.co.uk/2011/02/numpy-arrays-to-video.html>

	"""
	def __init__(self,filename,colorspace='rgb24',cache=None,blocking=False):

		# check that the file exists
		try:
			open(filename)
		except IOError:
			raise IOError('Movie "%s" does not exist!' %filename)

		# default cache location
		if cache == None:
			cache = '/tmp/%s_tmpframes' %path.basename(filename)

		# get the format of the movie
		format = getformat(filename)
		nrows = int(format['ID_VIDEO_HEIGHT'])
		ncols = int(format['ID_VIDEO_WIDTH'])
		seconds = float(format['ID_LENGTH'])
		framerate = float(format['ID_VIDEO_FPS'])
		nframes = int(seconds*framerate)

		self.colorspace = colorspace

		# get the number of color channels (so we can know the size of a
		# single frame)
		if colorspace in ['y800','y8']:
			nchan = 1
			bitdepth = 'uint8'
			atomsize = 1
		elif colorspace in ['rgb16','rgb24','bgr16','bgr24']:
			nchan = 3
			bitdepth = 'uint8'
			atomsize = 1
		elif colorspace in ['rgba','argb','rgb32','bgra','abgr','bgr32']:
			nchan = 4
			atomsize = 1
			bitdepth = 'uint8'
		else:
			raise Exception('Sorry - "%s" is not a currently supported colorspace',colorspace)

		self.blocking = blocking
		self.info = format
		self.filename = filename
		self.framerate = framerate
		self.duration = seconds
		self.shape = (nframes,nrows,ncols,nchan)
		self.bitdepth = bitdepth

		self._framesize = nrows*ncols*nchan*atomsize
		self._offsets = self._getoffsets()
		self._cache_p = cache

		# make temporary file containing the raw frames
		self.frames = self.decode()

	def __getitem__(self,key):
		"""
		Grab frames from the source movie as numpy arrays
		"""

		nframes,nrows,ncols,nchan = self.shape
		bits = self.bitdepth
		rcount = self._framesize
		offsets = self._offsets
		raw = self._cache_h
		thread = self._thread

		if thread.isAlive():
			print 'Not done decoding frames yet'
			return

		if raw.closed:
			print 'Temporary frame cache is closed. Call .decode() to create a new frame cache.'
			return

		if isinstance(key,tuple):
			raise IndexError('Too many indices')
		elif isinstance(key,int):
			# -ve indices read from the end
			if key < 0:
				key = nframes+key
			indices = (key,)
		elif isinstance(key,slice):
			indices = np.arange(*key.indices(nframes),dtype='uint16')

		outsize = (len(indices),nrows,ncols,nchan)
		framesout = np.empty(outsize,dtype=bits)

		for ii in xrange(len(indices)):
			oo = offsets[indices[ii]]
			raw.seek(oo,0)
			frame = np.fromfile(raw,count=rcount,dtype=bits)
			framesout[ii] = frame.reshape(self.shape[1:])

		return framesout.squeeze()

	def __iter__(self):
		"""
		Iterate over the source movie, return frames as numpy arrays
		"""
		nframes = self.shape[0]
		bits = self.bitdepth
		rcount = self._framesize
		raw = self._cache_h
		thread = self._thread

		if thread.isAlive():
			print 'Not done decoding frames yet'
			return

		if raw.closed:
			print 'Temporary frame cache is closed. Call .decode() to create a new frame cache.'
			return

		raw.seek(0,0)
		for ii in xrange(nframes):
			frame = np.fromfile(raw,count=rcount,dtype=bits)
			yield frame.reshape(self.shape[1:])

	def close(self):
		"""
		Close the temporary frame cache and delete it from the disk
		"""
		print "Deleting temporary frame cache"
		self._cache_h.close()
		remove(self._cache_p)

	def decode(self):

		# wipe any existing data in the cache
		file(self._cache_p,'w')

		# read the .avi, dump the output in a temporary file
		cmdstring = (	'mencoder',
				self.filename,
				'-ac','none',
				'-ovc','raw',
				'-nosound',
				'-vf','format=%s' %self.colorspace,
				'-of','rawvideo',
				'-o',self._cache_p,
				'-really-quiet'
				)

		def _notifyOnExit(onExit, popenArgs,blocking):
			def runInThread(onExit, popenArgs):
				proc = subprocess.Popen(popenArgs[0],**popenArgs[1])
				proc.wait()
				onExit()
				return
			thread = threading.Thread(target=runInThread,args=(onExit,popenArgs))
			thread.start()
			if blocking:
				# wait until thread completes
				thread.join()
			return thread
		def onExit():
			print "Finished decoding frames in %s" %self.filename
			return

		with open(devnull,"w") as fnull:
			print "Decoding frames in %s to temporary cache. Be patient..." %self.filename
			self._thread = _notifyOnExit(onExit,(cmdstring,dict(shell=False)),self.blocking)
			# self._getframes = subprocess.Popen(cmdstring,stdout=fnull,stderr=fnull,shell=False)

		# open the temporary file for reading
		self._cache_h = file(self._cache_p,'r')

	def _getoffsets(self):
		offsets = []
		for index in xrange(self.shape[0]):
			offsets.append(self._framesize*index)
		return offsets


def getformat(filename):
	"""
	Grab the header information from a movie file using mplayer, return it
	as a dict
	"""

	# use mplayer to grab the header information
	cmdstring = (	'mplayer',
			'-vo','null',
			'-ao','null',
			'-frames','0',
			'-really-quiet',
			'-identify',
			filename
			)

	# suppress socket error messages
	with open(devnull, "w") as fnull:
		formatstr = subprocess.check_output(cmdstring,stderr=fnull,shell=False)
	lines = formatstr.splitlines()
	format = {}

	for line in lines:
		name,value = line.split('=')
		format.update({name:value})

	return format