#!/usr/bin/env python

"""Multi-thread stuff for python."""

# Copyright Greg Kochanski, 2007
# Released under the LGPL version 2 or later.

import g_mailbox as GM
# DEBUG = 2
DEBUG = 0
# import g_foo as GM
import sys
from gmisclib import die

EOF = GM.EOF
prinfo = GM.prinfo
Lock = GM.Lock
Condition = GM.Condition
start_new_thread = GM.start_new_thread


def ncpu():
	n = 0
	for x in open('/proc/cpuinfo', 'r'):
		if x.split(':')[0].strip() == 'processor':
			n += 1
	assert n > 0, "Silly!"
	return n


class CountedCondition(GM.Condition):
	def __init__(self, lock=None, verbose=0, name=None):
		GM.Condition.__init__(self, lock, verbose, name)
		self.i = 0

	def incr(self, delta):
		self.i += delta
		self.notify()



class default_pool_initializer(object):
	def __init__(self, info, common):
		pass
	
	def close(self):
		pass
	
	def other_get_err(self):
		pass
	
	def catch_exception(self, server, inp):
		die.catch('catch exception')
		pass

	def get_fcn(self, fcn):
		return fcn


class rmap_pool_initializer(default_pool_initializer):
	import g_remote_exec as RE

	def __init__(self, info, path):
		# print "info=", info
		execpath, startargs = info
		self.c = self.RE.connection(execpath, startargs)
		# print 'self.c=', self.c
		fsp = self.RE.remote_fcn(self.c, None, 'gyropy', ('g_remote_exec', 'set_remote_path'))
		# print 'fsp=', fsp
		fsp(path)
		# print 'rmap_pool_initializer.__init__ done'
	

	def close(self):
		self.c.close()
	

	def other_get_err(self):
		self.close()
	

	def get_fcn(self, mod_fcn_name):
		modname, fcnname = mod_fcn_name
		return self.RE.remote_fcn(self.c, None, modname, fcnname)


class pool_c(object):
	def __init__(self, poolinfo, initialization=default_pool_initializer, common=None):
		"""When you're doing a rmap_pool_initializer, then common is used
		for pypath.
		"""
		self.mbox = GM.mailbox()
		self.startlock = Lock()

		self.rc = CountedCondition()
		for info in poolinfo:
			self.startlock.acquire()
			GM.start_new_thread( self.__a_thread,
						(self.mbox, initialization, info, common, self.startlock, self.rc)
						)
		self.startlock.acquire()
		self.startlock.release()
		del self.startlock


	def put(self, server, inp):
		self.mbox.put( (server, inp) )


	def join(self):
		self.mbox.putclose()
		self.rc.acquire()
		while self.rc.i > 0:
			self.rc.wait()
		self.rc.release()


	@staticmethod
	def __a_thread(mbox, initialization, info, common, startlock, runcount):
		runcount.acquire()
		runcount.incr(1)
		runcount.release()

		init = initialization(info, common)
		startlock.release()
		while True:
			# print '__a_thread loop top'
			try:
				server, inp = mbox.get()
			except GM.EOF:
				# print '__a_thread EOF'
				init.close()
				break
			except:
				# print '__a_thread other except!'
				init.other_get_err()
				break

			# print '__a_thread', server, inp, init
			try:
				server(inp, init)
			except:
				# print '__a_thread except2!'
				init.catch_exception(server, inp)
			# print '__a_thread_loop_bot'
		init.close()
		runcount.acquire()
		runcount.incr(-1)
		runcount.release()



class ssh_rmap_pool_c(pool_c):
	POOLTEMPLATE  = ['ssh', None, 'python', '-c', '"import gyropy; gyropy.g_remote_exec.server()"']

	def __init__(self, poolinfo, initialization=rmap_pool_initializer, pypath=None,
				pooltemplate=None):
		if pooltemplate is None:
			pooltemplate = self.POOLTEMPLATE
		pooli = []
		for pi in poolinfo:
			q = list(pooltemplate)
			q[q.index(None)] = pi
			pooli.append( (None,q) )
		pool_c.__init__(self, pooli, initialization=initialization, common=pypath)



def execute(fcn, inputlist, kwargs={}, pool=None):
	map(fcn, inputlist, kwargs, pool)






def map(fcn, inputlist, kwargs={}, pool=None):
	assert pool is not None
	exinfo = []
	o = []
	outlock = CountedCondition(name='map.server', verbose=DEBUG>1)

	def server(input_j, init):
		f = init.get_fcn(fcn)
		inp, j = input_j
		# print 'in-server', inp, j
		try:
			tmp = f(inp, **kwargs)
		except:
			ex = sys.exc_info()
			# print 'except!'
			outlock.acquire()
			exinfo.append(ex)
			outlock.incr(1)
			outlock.release()
		else:	
			# print 'in-server f-done', tmp
			outlock.acquire()
			o[j] = tmp
			outlock.incr(1)
			outlock.release()
		# print 'server done'

	# print 'map go'
	for (j, inp) in enumerate(inputlist):
		# print 'check exinfo'
		outlock.acquire()
		# print 'got outlock'
		if exinfo:
			outlock.release()
			break
		o.append(None)
		outlock.release()
		# print 'pool.put'
		pool.put( server, (inp, j) )

	# print 'wait loop', 'ex=', exinfo
	outlock.acquire()
	while not exinfo and outlock.i < len(o):
		# print 'outlock ready to wait'
		outlock.wait()
		# print 'outlock notified'
	# print 'after wait loop', 'ex=', exinfo
	ex = exinfo
	outlock.release()
	# print 'wait done', 'ex=', ex
	if ex:
		ex = ex[0]
		raise ex[0], ex[1], ex[2]
	return o


def rmap(modname, fcnname, inputlist, kwargs={}, pool=None):
	return map( (modname, fcnname), inputlist, kwargs, pool)


class remote_func(object):
	"""This sets up a thing that acts as a remote function call.
	The number of simultaneous calls is controlled by the
	size of the pool.

	So, you can do a map(func) and then within func,
	you can do a call to a remote machine.
	"""

	def __init__(self, modname, fcnname, pool=None):
		self.modname = modname
		self.funcname = fcnname
		self.pool = pool
	
	def __call__(self, input, **kwargs):
		return rmap((self.modname, self.funcname), [input], kwargs, self.pool)[0]


class _locked_within_max(object):
	import time as _t
	import math as _m
	# import os as _os

	def __init__(self):
		self.lock = Lock(name='locked_within_max', verbose=DEBUG>1)
		self.count = 0
		self.max = 0
		self.q = 0.0

	def inc(self, inc):
		self.lock.acquire()
		self.count += inc
		if self.count > self.max:
			self.max = self.count
		self.lock.release()

	def sleep(self):
		# sys.stdout.flush()
		self.inc(1)
		# self._os.system('sleep 2')
		self._t.sleep(self._m.exp(self._m.sin(self.q)-1.0))
		self.q += 1.0
		self.inc(-1)
		# sys.stdout.flush()


def test(nproc=0):
	print 'TEST', nproc
	mypool = pool_c([1]*nproc)
	sci = 0
	for scii in range(15):
		sci += 1 + sci//4 + scii//2
		sys.setcheckinterval(sci)
		i = range(10)
		tmp = map(lambda x:str(x), i, pool=mypool)
		# print 'tmp=', tmp
		if tmp != [str(x) for x in i]:
			print 'Whoops!', tmp, i
			sys.exit(1)
	mypool.join()


def test_to_string(x):
	return str(x)

def testrmap(nproc=1):
	print 'TESTRMAP', nproc
	g_r_p = pool_c([(None, ('python', '-c', 'import gyropy; gyropy.g_remote_exec.server()'))]*nproc,
			initialization = rmap_pool_initializer
			)
	for sci in range(0,100,11):
		sys.setcheckinterval(sci)
		i = range(10)
		tmp = rmap('gyropy', ('g_remote_exec', 'echo'), i, pool=g_r_p)
		if tmp != i:
			die.die('Whoops: rmap echo: %s %s' % (tmp, i))
		tmp = rmap('gyropy', ('g_threading', 'test_to_string'), i, pool=g_r_p)
		# print 'tmp=', tmp
		if tmp != [str(x) for x in i]:
			die.die('Whoops: %s %s' % (tmp, i))
	g_r_p.join()


def _ri(x):
	if x == 3:
		if __debug__:
			prinfo('Ready to raise Keyboard Interrunpt')
		raise KeyboardInterrupt('fake')
	return x

def _rj(x):
	if x == 3:
		if __debug__:
			prinfo('Ready to raise ValueError')
		raise ValueError, 'fake'
	return x

def testI():
	print 'TESTI'
	assert _localpool is not None
	try:
		i = range(10)
		assert _localpool is not None
		try:
			map(_ri, i, pool=_localpool)
			print 'Whoops -- lost a KeyboardInterrupt'
			sys.exit(1)
		except KeyboardInterrupt, x:
			prinfo('Caught KeyboardInterrupt')
			# tmp = sys.exc_info()
			# print 'testI caught KeyboardInterrupt', x, str(x)
			# print 'tmp=', tmp
			# print 'str(exc_info)=', str(tmp)
			sys.stdout.flush()
			if str(x) != 'fake':
				print 'Whoops! no fake'
				sys.exit(1)
			pass
		print 'TestI1'
		sys.stdout.flush()
		try:
			map(_rj, i, pool=_localpool)
		except ValueError, x:
			prinfo('Caught ValueError')
			sys.stdout.flush()
			if str(x) != 'fake':
				print 'Whoops no fake 2'
				sys.exit(1)
			pass
		else:
			print 'Whoops -- lost a ValueError'
			sys.exit(1)
		print 'TestI2'
		sys.stdout.flush()
	except:
		print 'UNEXPECTED EXCEPTION in testI'
		raise
	print 'TestI3'


def teste():
	print 'TESTE'
	import random
	import math
	import time
	sci = 0
	for scii in range(15):
		sci += 1 + sci//4 + scii/2
		print 'TESTe start'
		q = sci
		def fcn(x):
			# print "Fcn(", x, ")"
			if random.random() < 0.5:
				time.sleep(0.0001*math.exp(random.random()*3))
			x.append(q)
			if random.random() < 0.5:
				time.sleep(0.0001*math.exp(random.random()*3))
			# print 'Fcn done'
		sys.setcheckinterval(sci)
		for lng in [1, 2, 3, 4, 11, 44]:
			tmp = [ [] for i in range(lng) ]
		execute(fcn, tmp, pool=_localpool)
		if len(tmp) != lng:
			print 'len(tmp)!=lng'
			sys.exit(1)
		for t in tmp:
			if t != [q]:
				print 't!=[q]'
				sys.exit(1)
		print 'TESTe end'


def testre_remote(q, x):
	print 'testre_remote', q, x
	q.append(x)



def testre():
	print 'TESTre start'
	tmp = [[],[],[],[],[],[],[],[]]
	lt = len(tmp)
	q = 3
	print 'pre-Rexecute'
	g_r_p = global_remote_pool(4)
	Rexecute('g_threading', 'testre_remote', tmp, kwargs={'x':q}, pool=g_r_p)
	if lt != len(tmp):
		print 'lt != len(tmp)'
		sys.exit(1)
	for t in tmp:
		if t != [q]:
			print "Whoops!  t=%s" % repr(t)
			sys.exit(1)
	print 'TESTre middle'
	x = []
	tmp  = [x,x,x,x]
	Rexecute('g_threading', 'testre_remote', tmp, kwargs={'x':q}, pool=g_r_p)
	for t in tmp:
		if t != [q,q,q,q]:
			print "Whoops!  t=%s" % repr(t)
			sys.exit(1)
	print 'TESTre end'
	g_r_p.join()
	print 'g_r_p joined'


def testre2_remote(q, o):
	import math
	o.append( (q, math.sqrt(q)) )

def testre2():
	tmp = [1, 4, 9, 0, -1, 1]
	g_r_p = global_remote_pool(4)
	o = []
	try:
		Rexecute('g_threading', 'testre2_remote', tmp, kwargs={'o':o}, pool=g_r_p)
	except ValueError, x:
		if 'math domain' not in str(x):
			print 'Whoops: wrong kind of exception'
			sys.exit(1)
		pass
	else:
		print "Missed exception in teste2"
		sys.exit(1)
	g_r_p.join()


def teste2():
	tmp = [1, 4, 9, 0, -1, 1]
	o = []
	try:
		execute(testre2_remote, tmp, kwargs={'o':o}, pool=_localpool)
	except ValueError, x:
		if 'math domain' not in str(x):
			print 'whoops2, wrong exception'
			sys.exit(1)
		pass
	else:
		print "Missed exception in teste2"
		sys.exit(1)

def testw(size, ps):
	print 'TESTW', size, ps
	lw = _locked_within_max()
	tmp = [lw] * size
	pool = pool_c([1]*ps)
	execute(lambda x: x.sleep(), tmp, pool=pool)
	if not ( 1 <= lw.max <= len(tmp) ):
		print 'whoops: wrong lw.max'
		sys.exit(1)
	prinfo('pool.join()')
	pool.join()
	prinfo('pool.joined')



if __name__ == '__main__':
	_localpool = pool_c([1])
	testI()
	teste()
	# teste2()
	# testre()
	# testre2()
	# sys.exit(0)
	# g_r_p.join()
	# sys.exit(0)
	test(1)
	# sys.exit(0)
	test(2)
	test(9)
	# test(0)
	# test(-1)
	# test(-7)
	testI()
	testw(1, 1)
	testw(1, 5)
	testw(3, 1)
	testw(3, 3)
	testw(3, 4)
	testw(13, 1)
	testw(13, 3)
	testw(13, 12)
	testw(13, 14)
	testw(13, 44)
	# testw(33, 1)
	# testw(33, 13)
	# testw(33, 33)
	# testw(33, 53)
	# testre()
	# testre2()
	testrmap(1)
	testrmap(2)
	testrmap(3)
	testrmap(11)
	print 'OK'
