# -*- coding: utf-8 -*-
"""
Runs gathering dependency data (see winarch_processor) on multiple processors using MPI.

Author: Bartosz Broda, bartosz.broda@gmail.com
"""

from mpi4py import MPI

import os, time

from memprocessor import MemProcessor

from architect2 import  filters
from architect2.triples import  triplesgen_facades

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

MSG_END = 0 # end of processing
MSG_FILE = 1 # sending file
MSG_TRIPLES = 2  # sending triples

class Message:
	"""
	Very simple class used in message passing
	"""

	def __init__(self, msgtype, data):
		self.type = msgtype
		self.data = data
		
	def __repr__(self):
		return  "Message," + str( (self.type, self.data) )

def master(indir):
	"""
	Master-processing.
	The pseudo code for the master is:
	- gather all files to process
	- distribute calculations for all the files to nodes


	TODO extract file prefix
	"""
	allfiles = []
	for x, ddd in sorted([(r,f) for (r,d,f) in os.walk(indir) if len(f)] ):
		allfiles.extend( [os.path.join(x, w) for w in ddd if w.startswith('xx')] )

	allfiles = sorted(allfiles)


	free_set = set(range(1,size))
	busy_set = set([])
	
	file_no = 0
	while file_no < len(allfiles):
		# send files for processing
		while len(free_set):
			i = free_set.pop()
			busy_set.add(i)
			m = Message(MSG_FILE, allfiles[file_no])
			print "Master: Sending mesage %s to %d | (%d / %d) " % (allfiles[file_no], i, file_no, len(allfiles))
			comm.send((m), dest=i, tag=i)
			file_no += 1
		
		# wait for any free node...
		while len(free_set) == 0:
			status = MPI.Status()
			while not comm.Iprobe(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status):
				time.sleep(0.01)
			src = status.Get_source()
			free_set.add(src)
			busy_set.remove(src)
			comm.recv(source = src, tag = MPI.ANY_TAG) # not receiving data
		
	while len(busy_set):
		status = MPI.Status()
		print "Master: Probing"
		while not comm.Iprobe(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status):
			time.sleep(0.01)
		src = status.Get_source()
		free_set.add(src)
		busy_set.remove(src)
		comm.recv(source = src, tag = MPI.ANY_TAG)
		print "Master: Queues: free %s; busy %s" % (str(free_set), str(busy_set))
	
	for i in free_set:
		m = Message(MSG_END, "EOF")
		print "Master: Sending mesage EOF to %d " % (i)
		comm.send(m, dest=i, tag=i)
	
def slave(token_filter, left_size = 3, right_size = 3, dump_interval = 2000, outdir = 'data'):
	"""
	Slave process for processing documents.

	It waits for MSG_FILE and processes given document. After receiving MSG_END it writes
	down gathered data to outdir. Data is also written in regular intervals.
	"""
	processor = MemProcessor()
	i = 0
	doc = 0
	while True:
		msg = comm.recv(source = MPI.ANY_SOURCE, tag=MPI.ANY_TAG)
		
		if msg.type == MSG_END:
			processor.dump( os.path.join(outdir, 'data%d_%d.txt' % (rank,doc) ) )
			print "Slave: Rank %d QUITTING" % rank
			break
		elif msg.type == MSG_FILE:
			gen = triplesgen_facades.words_lower_triples_from_filename(msg.data, left_size, right_size, token_filter)
			processor.process( gen )
			i+=1
			if i % dump_interval == 0:
				doc += 1
				print "####### DUMP", rank
				processor.dump( os.path.join(outdir, 'data%d_%d.txt' % (rank,doc) ) )
			
			comm.send("Finished %s" % msg.data , dest=0, tag=11)
		else:
			print "Slave: Unknown message %d, Data: %s" % (msg.type, str(msg.data) )

def main():
	import argparse
	parser = argparse.ArgumentParser(description='Converts directory with dependency triples into a SuperMatrix',
									epilog = "Note, this program should be invoked via mpirun, e.g., mpirun -n 5 %(prog)s [options]")
	parser.add_argument('-i', '--indir', dest = 'indir', required=True, 
			help = 'Input directory with dependency tripeles.')
	parser.add_argument('-o', '--outdir', dest = 'outdir', default = './', 
			help = 'Output directory for writing SuperMatrix.')
	parser.add_argument('-u', '--udict', dest = 'udict', required=True, 
			help = 'File with unigram dictionary.')
	
	args = parser.parse_args()
	
	indir = args.indir
	outdir = args.outdir
	dictfile = args.udict
	
	token_filter = filters.dict_eng_filter(dictfile)

	if rank == 0:
		if not os.path.exists(outdir):
			os.makedirs(outdir)
		master(indir)
	else:
		slave(token_filter = token_filter, outdir = outdir)

if __name__ == '__main__':
	main()

