<commit_before>import dlvhex

import shallowhexparser as shp

import logging

def id(p):
	for x in dlvhex.getTrueInputAtoms():
		tup = x.tuple()
		if len(tup) != 2:
			raise Exception("this external atom processes only arity 1 predicate inputs")
		dlvhex.output( (tup[1],) )

def idc(c):
	dlvhex.output((c,))

def testZeroArity0():
	pass

def testZeroArity1():
	dlvhex.output(())

def testA(pred):
	if len(dlvhex.getTrueInputAtoms()) == 0:
		dlvhex.output(('foo',))
	else:
		dlvhex.output(('bar',))

def testB(pred1, pred2):
	if len(dlvhex.getTrueInputAtoms()) <= 1:
		dlvhex.output(('bar',))
	else:
		dlvhex.output(('foo',))

def testC(pred):
	for atom in dlvhex.getTrueInputAtoms():
		for x in atom.tuple()[1:]:
			# output arguments of input predicate
			dlvhex.output((x.value(),))

def testEven(pred1, pred2):
	true = [x for x in dlvhex.getTrueInputAtoms()]
	num = len(true)
	if num % 2 == 0:
		dlvhex.output(())

def testSubstr(string, start, length):
	stringv = string.value()
	startv = start.intValue()
	lengthv = start.intValue()
	unquoted = stringv.strip('"')
	if startv+lengthv < len(unquoted):
		dlvhex.output((unquoted[startv:startv+lengthv],))

def testStrlen(string):
	stringv = string.value()
	unquoted = stringv.strip('"')
	dlvhex.output( (len(unquoted),) )

def testSmallerThan(int1, int2):
	if int1.intValue() < int2.intValue():
		dlvhex.output( () )

def testConcat(strs):
	#logging.debug('testConcat got '+repr(strs)+' '+str(strs.__class__)+' '+str(strs[0].__class__))
	values = [s.value() for s in strs]
	#logging.debug('testConcat values '+repr(values))
	needquote = any(['"' in s for s in values])
	#logging.debug('testConcat needquote '+repr(needquote))
	unquoted = [s.strip('"') for s in values]
	result = ''.join(unquoted)
	if needquote:
		result = '"'+result+'"'
	#logging.debug('testConcat returns '+repr(result))
	dlvhex.output((result,))

def functionCompose(args):
	logging.debug('functionCompose got '+repr(args))
	if len(args) == 1:
		dlvhex.output((args[0],))
	else:
		apred = args[0].value()
		avalues = [a.value() for a in args[1:]]
		dlvhex.output(("{}({})".format(apred, ','.join(avalues)),))

def functionDecompose(term, narg):
	logging.debug('functionDecompose got {} and {}'.format(repr(term), narg))
	pinp = shp.parseTerm(term.value())
	logging.debug('parseTerm {}'.format(repr(pinp)))
	argidx = narg.intValue()
	if argidx < len(pinp):
		dlvhex.output( (shp.shallowprint(pinp[argidx]),) )

def functionDecomposeN(inp, N):
	logging.debug('functionDecomposeN got {} and {}'.format(repr(inp), N))
	pinp = shp.parseTerm(inp.value())
	logging.debug('parseTerm {}'.format(repr(pinp)))
	if len(pinp) == N+1:
		otuple = [ shp.shallowprint(x) for x in pinp ]
		dlvhex.output( tuple(otuple) )

def functionDecompose1(inp):
	functionDecomposeN(inp, 1)
def functionDecompose2(inp):
	functionDecomposeN(inp, 2)
def functionDecompose3(inp):
	functionDecomposeN(inp, 3)

def getArity(term):
	logging.debug('getArity got {}'.format(repr(term)))
	pinp = shp.parseTerm(term.value())
	logging.debug('parseTerm {}'.format(repr(pinp)))
	dlvhex.output( (len(pinp)-1,) )

def isEmpty(assignment):

	true = 0
	false = 0
	unknown = 0

	premisse = ()
	for x in dlvhex.getInputAtoms():
		if x.isTrue():
			true = true + 1
		elif x.isFalse():
			false = false + 1
		else:
			unknown = unknown + 1

	if true > 0:
		# external atom is true
		dlvhex.output(())
	elif (true + unknown) > 0:
		# external atom can be true
		dlvhex.outputUnknown(())
	else:
		# else case applies: (true + unknown) < min.intValue() or true > max.intValue()
		#
		# external atom is certainly not true
		v = 0

def numberOfBalls(assignment, min, max):

	true = 0
	false = 0
	unknown = 0

	premisse = ()
	for x in dlvhex.getInputAtoms():
		if x.isTrue():
			true = true + 1
		elif x.isFalse():
			false = false + 1
		else:
			unknown = unknown + 1
			v = 0

	if true >= min.intValue() and (true + unknown) <= max.intValue():
		# external atom is true
		dlvhex.output(())
	elif (true + unknown) >= min.intValue() and true <= max.intValue():
		# external atom can be true
		dlvhex.outputUnknown(())
	else:
		# else case applies: (true + unknown) < min.intValue() or true > max.intValue()
		#
		# external atom is certainly not true
		v = 0

def numberOfBallsSE(assignment, max):

	true = 0
	false = 0
	unknown = 0

	premisse = ()
	for x in dlvhex.getInputAtoms():
		if x.isTrue():
			true = true + 1
		elif x.isFalse():
			false = false + 1
		else:
			unknown = unknown + 1
			v = 0

	if (true + unknown) <= max.intValue():
		# external atom is true
		dlvhex.output(())
	elif true <= max.intValue():
		# external atom can be true
		dlvhex.outputUnknown(())
	else:
		# else case applies: if true > max.intValue()
		#
		# external
		v = 0

def numberOfBallsGE(assignment, min):

	true = 0
	false = 0
	unknown = 0

	premisse = ()
	for x in dlvhex.getInputAtoms():
		if x.isTrue():
			true = true + 1
		elif x.isFalse():
			false = false + 1
		else:
			unknown = unknown + 1
			v = 0

	if true >= min.intValue():
		# external atom is true
		dlvhex.output(())
	elif (true + unknown) >= min.intValue():
		# external atom can be true
		dlvhex.outputUnknown(())
	else:
		# else case applies: if (true + unknown) < min.intValue()
		#
		# external
		v = 0

# no native implementations for these, so let's use the non-native ones
testIsEmpty = isEmpty
testNumberOfBalls = numberOfBalls
testNumberOfBallsSE = numberOfBallsSE
testNumberOfBallsGE = numberOfBallsGE

def partialTest(assignment):

	true = 0
	false = 0
	unknown = 0

	premisse = ()
	for x in dlvhex.getInputAtoms():
		if x.isTrue():
			true = true + 1
#			premisse = premisse + (x, )
#			print "true input atom:", x.value()
		elif x.isFalse():
			false = false + 1
#			premisse = premisse + (x.negate(), )
#			print "false input atom:", x.value()
		else:
			unknown = unknown + 1
#			print "unknown input atom:", x.value()
			v = 0

	if true > 1:
#		dlvhex.learn(premisse + (dlvhex.storeOutputAtom((), False).negate(), ))
		dlvhex.output(())
	elif true + unknown > 1:
		dlvhex.outputUnknown(())

def testSetMinus(p, q):
	pset, qset = set(), set()
	for x in dlvhex.getTrueInputAtoms():
		tup = x.tuple()
		if tup[0].value() == p.value():
			pset.add(tup[1].value())
		elif tup[0].value() == q.value():
			qset.add(tup[1].value())
	rset = pset - qset
	for r in rset:
		dlvhex.output( (r,) )

def rdf(uri):
	logging.warning('TODO implement &rdf (and #namespace)')
	dlvhex.output(('s', 'p', 'o'))

def register():
	#XFAIL = expected failure
	dlvhex.addAtom("testA", (dlvhex.PREDICATE,), 1)
	dlvhex.addAtom("testB", (dlvhex.PREDICATE, dlvhex.PREDICATE), 1)
	dlvhex.addAtom("testC", (dlvhex.PREDICATE,), 1)
	dlvhex.addAtom("testZeroArity0", tuple(), 0)
	dlvhex.addAtom("testZeroArity1", tuple(), 0)
	#XFAIL unused dlvhex.addAtom("testConcatAll", (dlvhex.PREDICATE,), 1)
	#unused dlvhex.addAtom("testListDomain", (dlvhex.TUPLE,), 1)
	#unused dlvhex.addAtom("testListConcat", (dlvhex.TUPLE,), 1)
	#unused dlvhex.addAtom("testListLength", (dlvhex.CONSTANT,dlvhex.CONSTANT), 1)
	#unused dlvhex.addAtom("testListSplit", (dlvhex.CONSTANT,dlvhex.CONSTANT), 2)
	#unused dlvhex.addAtom("testListHalf", (dlvhex.CONSTANT,), 2)
	#unused dlvhex.addAtom("testListMerge", (dlvhex.CONSTANT,dlvhex.CONSTANT,dlvhex.CONSTANT), 2)
	dlvhex.addAtom("testSubstr", (dlvhex.CONSTANT,dlvhex.CONSTANT,dlvhex.CONSTANT), 1)
	dlvhex.addAtom("testStrlen", (dlvhex.CONSTANT,), 1)
	dlvhex.addAtom("testSmallerThan", (dlvhex.CONSTANT,dlvhex.CONSTANT), 0)
	dlvhex.addAtom("testEven", (dlvhex.PREDICATE,dlvhex.PREDICATE), 0)
	#unused dlvhex.addAtom("testOdd", (dlvhex.PREDICATE,dlvhex.PREDICATE), 0)
	#unused dlvhex.addAtom("testLessThan", (dlvhex.PREDICATE,dlvhex.PREDICATE), 0)
	#unused dlvhex.addAtom("testEqual", (dlvhex.PREDICATE,dlvhex.PREDICATE), 0)
	dlvhex.addAtom("id", (dlvhex.PREDICATE,), 1)
	#XFAIL partial dlvhex.addAtom("idp", (dlvhex.PREDICATE,), 1)
	dlvhex.addAtom("idc", (dlvhex.CONSTANT,), 1)
	#TODO testCautiousQuery
	dlvhex.addAtom("testSetMinus", (dlvhex.PREDICATE,dlvhex.PREDICATE), 1)

	prop = dlvhex.ExtSourceProperties()
	prop.setProvidesPartialAnswer(True)
	dlvhex.addAtom("isEmpty", (dlvhex.PREDICATE, ), 0, prop)
	dlvhex.addAtom("testIsEmpty", (dlvhex.PREDICATE, ), 0, prop)

	prop = dlvhex.ExtSourceProperties()
	prop.setProvidesPartialAnswer(True)
	dlvhex.addAtom("numberOfBalls", (dlvhex.PREDICATE, dlvhex.CONSTANT, dlvhex.CONSTANT), 0, prop)
	dlvhex.addAtom("testNumberOfBalls", (dlvhex.PREDICATE, dlvhex.CONSTANT, dlvhex.CONSTANT), 0, prop)

	prop = dlvhex.ExtSourceProperties()
	prop.setProvidesPartialAnswer(True)
	prop.addAntimonotonicInputPredicate(0)
	dlvhex.addAtom("numberOfBallsSE", (dlvhex.PREDICATE, dlvhex.CONSTANT), 0, prop)
	dlvhex.addAtom("testNumberOfBallsSE", (dlvhex.PREDICATE, dlvhex.CONSTANT), 0, prop)

	prop = dlvhex.ExtSourceProperties()
	prop.setProvidesPartialAnswer(True)
	prop.addMonotonicInputPredicate(0)
	dlvhex.addAtom("numberOfBallsGE", (dlvhex.PREDICATE, dlvhex.CONSTANT), 0, prop)
	dlvhex.addAtom("testNumberOfBallsGE", (dlvhex.PREDICATE, dlvhex.CONSTANT), 0, prop)

	prop = dlvhex.ExtSourceProperties()
	prop.setProvidesPartialAnswer(True)
	dlvhex.addAtom("partialTest", (dlvhex.PREDICATE, ), 0, prop)

	#XFAIL (TODO) sumD0
	#XFAIL getreq
	#unused mapping
	#unused getSizes
	#unused getSizesRestr
	#XFAIL getDiagnoses

	prop = dlvhex.ExtSourceProperties()
	prop.addFiniteOutputDomain(0)
	dlvhex.addAtom("testConcat", (dlvhex.TUPLE,), 1, prop)

	dlvhex.addAtom("functionCompose", (dlvhex.TUPLE,), 1)
	dlvhex.addAtom("functionDecompose", (dlvhex.CONSTANT,dlvhex.CONSTANT), 1)
	dlvhex.addAtom("functionDecompose1", (dlvhex.CONSTANT,), 2)
	dlvhex.addAtom("functionDecompose2", (dlvhex.CONSTANT,), 3)
	dlvhex.addAtom("functionDecompose3", (dlvhex.CONSTANT,), 4)
	dlvhex.addAtom("getArity", (dlvhex.CONSTANT,), 1)

	dlvhex.addAtom("rdf", (dlvhex.CONSTANT,), 3)

# vim:list:noet:
<commit_msg>added missing function symbols external atom<commit_after>import dlvhex

import shallowhexparser as shp

import logging

def id(p):
	for x in dlvhex.getTrueInputAtoms():
		tup = x.tuple()
		if len(tup) != 2:
			raise Exception("this external atom processes only arity 1 predicate inputs")
		dlvhex.output( (tup[1],) )

def idc(c):
	dlvhex.output((c,))

def testZeroArity0():
	pass

def testZeroArity1():
	dlvhex.output(())

def testA(pred):
	if len(dlvhex.getTrueInputAtoms()) == 0:
		dlvhex.output(('foo',))
	else:
		dlvhex.output(('bar',))

def testB(pred1, pred2):
	if len(dlvhex.getTrueInputAtoms()) <= 1:
		dlvhex.output(('bar',))
	else:
		dlvhex.output(('foo',))

def testC(pred):
	for atom in dlvhex.getTrueInputAtoms():
		for x in atom.tuple()[1:]:
			# output arguments of input predicate
			dlvhex.output((x.value(),))

def testEven(pred1, pred2):
	true = [x for x in dlvhex.getTrueInputAtoms()]
	num = len(true)
	if num % 2 == 0:
		dlvhex.output(())

def testSubstr(string, start, length):
	stringv = string.value()
	startv = start.intValue()
	lengthv = start.intValue()
	unquoted = stringv.strip('"')
	if startv+lengthv < len(unquoted):
		dlvhex.output((unquoted[startv:startv+lengthv],))

def testStrlen(string):
	stringv = string.value()
	unquoted = stringv.strip('"')
	dlvhex.output( (len(unquoted),) )

def testSmallerThan(int1, int2):
	if int1.intValue() < int2.intValue():
		dlvhex.output( () )

def testConcat(strs):
	#logging.debug('testConcat got '+repr(strs)+' '+str(strs.__class__)+' '+str(strs[0].__class__))
	values = [s.value() for s in strs]
	#logging.debug('testConcat values '+repr(values))
	needquote = any(['"' in s for s in values])
	#logging.debug('testConcat needquote '+repr(needquote))
	unquoted = [s.strip('"') for s in values]
	result = ''.join(unquoted)
	if needquote:
		result = '"'+result+'"'
	#logging.debug('testConcat returns '+repr(result))
	dlvhex.output((result,))

def isFunctionTerm(term):
	logging.debug('isFunctionTerm got '+repr(term))
	pinp = shp.parseTerm(term.value())
	logging.debug('parseTerm {}'.format(repr(pinp)))
	if len(pinp) > 1:
		# yes it is
		dlvhex.output( () )

def functionCompose(args):
	logging.debug('functionCompose got '+repr(args))
	if len(args) == 1:
		dlvhex.output((args[0],))
	else:
		apred = args[0].value()
		avalues = [a.value() for a in args[1:]]
		dlvhex.output(("{}({})".format(apred, ','.join(avalues)),))

def functionDecompose(term, narg):
	logging.debug('functionDecompose got {} and {}'.format(repr(term), narg))
	pinp = shp.parseTerm(term.value())
	logging.debug('parseTerm {}'.format(repr(pinp)))
	argidx = narg.intValue()
	if argidx < len(pinp):
		dlvhex.output( (shp.shallowprint(pinp[argidx]),) )

def functionDecomposeN(inp, N):
	logging.debug('functionDecomposeN got {} and {}'.format(repr(inp), N))
	pinp = shp.parseTerm(inp.value())
	logging.debug('parseTerm {}'.format(repr(pinp)))
	if len(pinp) == N+1:
		otuple = [ shp.shallowprint(x) for x in pinp ]
		dlvhex.output( tuple(otuple) )

def functionDecompose1(inp):
	functionDecomposeN(inp, 1)
def functionDecompose2(inp):
	functionDecomposeN(inp, 2)
def functionDecompose3(inp):
	functionDecomposeN(inp, 3)

def getArity(term):
	logging.debug('getArity got {}'.format(repr(term)))
	pinp = shp.parseTerm(term.value())
	logging.debug('parseTerm {}'.format(repr(pinp)))
	dlvhex.output( (len(pinp)-1,) )

def isEmpty(assignment):

	true = 0
	false = 0
	unknown = 0

	premisse = ()
	for x in dlvhex.getInputAtoms():
		if x.isTrue():
			true = true + 1
		elif x.isFalse():
			false = false + 1
		else:
			unknown = unknown + 1

	if true > 0:
		# external atom is true
		dlvhex.output(())
	elif (true + unknown) > 0:
		# external atom can be true
		dlvhex.outputUnknown(())
	else:
		# else case applies: (true + unknown) < min.intValue() or true > max.intValue()
		#
		# external atom is certainly not true
		v = 0

def numberOfBalls(assignment, min, max):

	true = 0
	false = 0
	unknown = 0

	premisse = ()
	for x in dlvhex.getInputAtoms():
		if x.isTrue():
			true = true + 1
		elif x.isFalse():
			false = false + 1
		else:
			unknown = unknown + 1
			v = 0

	if true >= min.intValue() and (true + unknown) <= max.intValue():
		# external atom is true
		dlvhex.output(())
	elif (true + unknown) >= min.intValue() and true <= max.intValue():
		# external atom can be true
		dlvhex.outputUnknown(())
	else:
		# else case applies: (true + unknown) < min.intValue() or true > max.intValue()
		#
		# external atom is certainly not true
		v = 0

def numberOfBallsSE(assignment, max):

	true = 0
	false = 0
	unknown = 0

	premisse = ()
	for x in dlvhex.getInputAtoms():
		if x.isTrue():
			true = true + 1
		elif x.isFalse():
			false = false + 1
		else:
			unknown = unknown + 1
			v = 0

	if (true + unknown) <= max.intValue():
		# external atom is true
		dlvhex.output(())
	elif true <= max.intValue():
		# external atom can be true
		dlvhex.outputUnknown(())
	else:
		# else case applies: if true > max.intValue()
		#
		# external
		v = 0

def numberOfBallsGE(assignment, min):

	true = 0
	false = 0
	unknown = 0

	premisse = ()
	for x in dlvhex.getInputAtoms():
		if x.isTrue():
			true = true + 1
		elif x.isFalse():
			false = false + 1
		else:
			unknown = unknown + 1
			v = 0

	if true >= min.intValue():
		# external atom is true
		dlvhex.output(())
	elif (true + unknown) >= min.intValue():
		# external atom can be true
		dlvhex.outputUnknown(())
	else:
		# else case applies: if (true + unknown) < min.intValue()
		#
		# external
		v = 0

# no native implementations for these, so let's use the non-native ones
testIsEmpty = isEmpty
testNumberOfBalls = numberOfBalls
testNumberOfBallsSE = numberOfBallsSE
testNumberOfBallsGE = numberOfBallsGE

def partialTest(assignment):

	true = 0
	false = 0
	unknown = 0

	premisse = ()
	for x in dlvhex.getInputAtoms():
		if x.isTrue():
			true = true + 1
#			premisse = premisse + (x, )
#			print "true input atom:", x.value()
		elif x.isFalse():
			false = false + 1
#			premisse = premisse + (x.negate(), )
#			print "false input atom:", x.value()
		else:
			unknown = unknown + 1
#			print "unknown input atom:", x.value()
			v = 0

	if true > 1:
#		dlvhex.learn(premisse + (dlvhex.storeOutputAtom((), False).negate(), ))
		dlvhex.output(())
	elif true + unknown > 1:
		dlvhex.outputUnknown(())

def testSetMinus(p, q):
	pset, qset = set(), set()
	for x in dlvhex.getTrueInputAtoms():
		tup = x.tuple()
		if tup[0].value() == p.value():
			pset.add(tup[1].value())
		elif tup[0].value() == q.value():
			qset.add(tup[1].value())
	rset = pset - qset
	for r in rset:
		dlvhex.output( (r,) )

def rdf(uri):
	logging.warning('TODO implement &rdf (and #namespace)')
	dlvhex.output(('s', 'p', 'o'))

def register():
	#XFAIL = expected failure
	dlvhex.addAtom("testA", (dlvhex.PREDICATE,), 1)
	dlvhex.addAtom("testB", (dlvhex.PREDICATE, dlvhex.PREDICATE), 1)
	dlvhex.addAtom("testC", (dlvhex.PREDICATE,), 1)
	dlvhex.addAtom("testZeroArity0", tuple(), 0)
	dlvhex.addAtom("testZeroArity1", tuple(), 0)
	#XFAIL unused dlvhex.addAtom("testConcatAll", (dlvhex.PREDICATE,), 1)
	#unused dlvhex.addAtom("testListDomain", (dlvhex.TUPLE,), 1)
	#unused dlvhex.addAtom("testListConcat", (dlvhex.TUPLE,), 1)
	#unused dlvhex.addAtom("testListLength", (dlvhex.CONSTANT,dlvhex.CONSTANT), 1)
	#unused dlvhex.addAtom("testListSplit", (dlvhex.CONSTANT,dlvhex.CONSTANT), 2)
	#unused dlvhex.addAtom("testListHalf", (dlvhex.CONSTANT,), 2)
	#unused dlvhex.addAtom("testListMerge", (dlvhex.CONSTANT,dlvhex.CONSTANT,dlvhex.CONSTANT), 2)
	dlvhex.addAtom("testSubstr", (dlvhex.CONSTANT,dlvhex.CONSTANT,dlvhex.CONSTANT), 1)
	dlvhex.addAtom("testStrlen", (dlvhex.CONSTANT,), 1)
	dlvhex.addAtom("testSmallerThan", (dlvhex.CONSTANT,dlvhex.CONSTANT), 0)
	dlvhex.addAtom("testEven", (dlvhex.PREDICATE,dlvhex.PREDICATE), 0)
	#unused dlvhex.addAtom("testOdd", (dlvhex.PREDICATE,dlvhex.PREDICATE), 0)
	#unused dlvhex.addAtom("testLessThan", (dlvhex.PREDICATE,dlvhex.PREDICATE), 0)
	#unused dlvhex.addAtom("testEqual", (dlvhex.PREDICATE,dlvhex.PREDICATE), 0)
	dlvhex.addAtom("id", (dlvhex.PREDICATE,), 1)
	#XFAIL partial dlvhex.addAtom("idp", (dlvhex.PREDICATE,), 1)
	dlvhex.addAtom("idc", (dlvhex.CONSTANT,), 1)
	#TODO testCautiousQuery
	dlvhex.addAtom("testSetMinus", (dlvhex.PREDICATE,dlvhex.PREDICATE), 1)

	prop = dlvhex.ExtSourceProperties()
	prop.setProvidesPartialAnswer(True)
	dlvhex.addAtom("isEmpty", (dlvhex.PREDICATE, ), 0, prop)
	dlvhex.addAtom("testIsEmpty", (dlvhex.PREDICATE, ), 0, prop)

	prop = dlvhex.ExtSourceProperties()
	prop.setProvidesPartialAnswer(True)
	dlvhex.addAtom("numberOfBalls", (dlvhex.PREDICATE, dlvhex.CONSTANT, dlvhex.CONSTANT), 0, prop)
	dlvhex.addAtom("testNumberOfBalls", (dlvhex.PREDICATE, dlvhex.CONSTANT, dlvhex.CONSTANT), 0, prop)

	prop = dlvhex.ExtSourceProperties()
	prop.setProvidesPartialAnswer(True)
	prop.addAntimonotonicInputPredicate(0)
	dlvhex.addAtom("numberOfBallsSE", (dlvhex.PREDICATE, dlvhex.CONSTANT), 0, prop)
	dlvhex.addAtom("testNumberOfBallsSE", (dlvhex.PREDICATE, dlvhex.CONSTANT), 0, prop)

	prop = dlvhex.ExtSourceProperties()
	prop.setProvidesPartialAnswer(True)
	prop.addMonotonicInputPredicate(0)
	dlvhex.addAtom("numberOfBallsGE", (dlvhex.PREDICATE, dlvhex.CONSTANT), 0, prop)
	dlvhex.addAtom("testNumberOfBallsGE", (dlvhex.PREDICATE, dlvhex.CONSTANT), 0, prop)

	prop = dlvhex.ExtSourceProperties()
	prop.setProvidesPartialAnswer(True)
	dlvhex.addAtom("partialTest", (dlvhex.PREDICATE, ), 0, prop)

	#XFAIL (TODO) sumD0
	#XFAIL getreq
	#unused mapping
	#unused getSizes
	#unused getSizesRestr
	#XFAIL getDiagnoses

	prop = dlvhex.ExtSourceProperties()
	prop.addFiniteOutputDomain(0)
	dlvhex.addAtom("testConcat", (dlvhex.TUPLE,), 1, prop)

	dlvhex.addAtom("functionCompose", (dlvhex.TUPLE,), 1)
	dlvhex.addAtom("functionDecompose", (dlvhex.CONSTANT,dlvhex.CONSTANT), 1)
	dlvhex.addAtom("functionDecompose1", (dlvhex.CONSTANT,), 2)
	dlvhex.addAtom("functionDecompose2", (dlvhex.CONSTANT,), 3)
	dlvhex.addAtom("functionDecompose3", (dlvhex.CONSTANT,), 4)
	dlvhex.addAtom("getArity", (dlvhex.CONSTANT,), 1)
	dlvhex.addAtom("isFunctionTerm", (dlvhex.CONSTANT,), 0)

	dlvhex.addAtom("rdf", (dlvhex.CONSTANT,), 3)

# vim:list:noet:
<|endoftext|>


<commit_before># Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from tests.beeswax.impala_beeswax import ImpalaBeeswaxException
from tests.common.impala_test_suite import ImpalaTestSuite
from tests.common.test_dimensions import create_single_exec_option_dimension
from tests.common.test_dimensions import create_uncompressed_text_dimension

class TestScratchLimit(ImpalaTestSuite):
  """
  This class tests the functionality of setting the scratch limit as a query option
  """

  spilling_sort_query = """
      select o_orderdate, o_custkey, o_comment
      from tpch.orders
      order by o_orderdate
      """
  spilling_agg_query = """
      select count(*) from (
        select distinct o_orderdate, o_custkey, o_comment
        from tpch_parquet.orders) v;
      """
  spilling_join_query = """
      select count(*)
      from tpch_parquet.lineitem join tpch_parquet.orders on l_orderkey = o_orderkey
      """
  # The analytic function in this query spills, but the sort generally spills first.
  # Ideally we would have a spilling analytic without a sort to exercise it in isolation.
  spilling_analytic_query = """
      SELECT i_item_sk, i_current_price,
        SUM (i_current_price)
        OVER (ORDER BY i_item_sk rows between 500000 preceding and 500000 following) running_total
      FROM tpcds_parquet.item
      ORDER BY i_brand, i_item_sk;
      """

  spilling_queries = [spilling_sort_query, spilling_agg_query, spilling_join_query,
      spilling_analytic_query]

  # Block manager memory limit that is low enough to
  # force Impala to spill to disk when executing 'spilling_sort_query'
  buffer_pool_limit = "64m"

  @classmethod
  def get_workload(self):
    return 'functional-query'

  @classmethod
  def add_test_dimensions(cls):
    super(TestScratchLimit, cls).add_test_dimensions()
    # There is no reason to run these tests using all dimensions.
    cls.ImpalaTestMatrix.add_dimension(create_single_exec_option_dimension())
    cls.ImpalaTestMatrix.add_dimension(
        create_uncompressed_text_dimension(cls.get_workload()))

  def test_with_high_scratch_limit(self, vector):
    """
    Sort query runs to completion with a scratch limit well above
    its required scratch space which in this case is 128m.
    """
    exec_option = vector.get_value('exec_option')
    exec_option['buffer_pool_limit'] = self.buffer_pool_limit
    exec_option['scratch_limit'] = '500m'
    self.execute_query_expect_success(self.client, self.spilling_sort_query, exec_option)

  def test_with_low_scratch_limit(self, vector):
    """
    Sort query throws the appropriate exception with a scratch limit well below
    its required scratch space which in this case is 128m.
    """
    exec_option = vector.get_value('exec_option')
    exec_option['buffer_pool_limit'] = self.buffer_pool_limit
    exec_option['scratch_limit'] = '24m'
    expected_error = 'Scratch space limit of %s bytes exceeded'
    scratch_limit_in_bytes = 24 * 1024 * 1024
    try:
      self.execute_query(self.spilling_sort_query, exec_option)
      assert False, "Query was expected to fail"
    except ImpalaBeeswaxException as e:
      assert expected_error % scratch_limit_in_bytes in str(e)

  def test_with_zero_scratch_limit(self, vector):
    """
    Queries throws the appropriate exception with a scratch limit of
    zero which means no scratch space can be allocated.
    """
    exec_option = vector.get_value('exec_option')
    exec_option['buffer_pool_limit'] = self.buffer_pool_limit
    exec_option['scratch_limit'] = '0'
    for query in self.spilling_queries:
      self.execute_query_expect_failure(query, exec_option)

  def test_with_unlimited_scratch_limit(self, vector):
    """
    Sort query runs to completion with a scratch Limit of -1 means default/no limit.
    """
    exec_option = vector.get_value('exec_option')
    exec_option['buffer_pool_limit'] = self.buffer_pool_limit
    exec_option['scratch_limit'] = '-1'
    self.execute_query_expect_success(self.client, self.spilling_sort_query, exec_option)

  def test_without_specifying_scratch_limit(self, vector):
    """
    Sort query runs to completion with the default setting of no scratch limit.
    """
    exec_option = vector.get_value('exec_option')
    exec_option['buffer_pool_limit'] = self.buffer_pool_limit
    self.execute_query_expect_success(self.client, self.spilling_sort_query, exec_option)

  def test_with_zero_scratch_limit_no_memory_limit(self, vector):
    """
    Queries run to completion without spilling as there is no limit on block memory manger.
    Scratch limit of zero ensures spilling is disabled.
    """
    exec_option = vector.get_value('exec_option')
    exec_option['scratch_limit'] = '0'
    for query in self.spilling_queries:
      self.execute_query_expect_success(self.client, query, exec_option)
<commit_msg>IMPALA-5772: also fix TestScratchLimit<commit_after># Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from tests.beeswax.impala_beeswax import ImpalaBeeswaxException
from tests.common.impala_test_suite import ImpalaTestSuite
from tests.common.test_dimensions import create_single_exec_option_dimension
from tests.common.test_dimensions import create_uncompressed_text_dimension

class TestScratchLimit(ImpalaTestSuite):
  """
  This class tests the functionality of setting the scratch limit as a query option
  """

  spilling_sort_query = """
      select o_orderdate, o_custkey, o_comment
      from tpch.orders
      order by o_orderdate
      """
  spilling_agg_query = """
      select count(*) from (
        select distinct o_orderdate, o_custkey, o_comment
        from tpch_parquet.orders) v;
      """
  spilling_join_query = """
      select count(*)
      from tpch_parquet.lineitem join tpch_parquet.orders on l_orderkey = o_orderkey
      """
  # The analytic function in this query spills, but the sort generally spills first.
  # Ideally we would have a spilling analytic without a sort to exercise it in isolation.
  spilling_analytic_query = """
      SELECT i_item_sk, i_current_price,
        SUM (i_current_price)
        OVER (ORDER BY i_item_sk rows between 500000 preceding and 500000 following) running_total
      FROM tpcds_parquet.item
      ORDER BY i_brand, i_item_sk;
      """

  spilling_queries = [spilling_sort_query, spilling_agg_query, spilling_join_query,
      spilling_analytic_query]

  # Buffer pool limit that is low enough to force Impala to spill to disk when executing
  # spill_query.
  buffer_pool_limit = "32m"

  @classmethod
  def get_workload(self):
    return 'functional-query'

  @classmethod
  def add_test_dimensions(cls):
    super(TestScratchLimit, cls).add_test_dimensions()
    # There is no reason to run these tests using all dimensions.
    cls.ImpalaTestMatrix.add_dimension(create_single_exec_option_dimension())
    cls.ImpalaTestMatrix.add_dimension(
        create_uncompressed_text_dimension(cls.get_workload()))

  def test_with_high_scratch_limit(self, vector):
    """
    Sort query runs to completion with a scratch limit well above
    its required scratch space which in this case is 128m.
    """
    exec_option = vector.get_value('exec_option')
    exec_option['buffer_pool_limit'] = self.buffer_pool_limit
    exec_option['scratch_limit'] = '500m'
    self.execute_query_expect_success(self.client, self.spilling_sort_query, exec_option)

  def test_with_low_scratch_limit(self, vector):
    """
    Sort query throws the appropriate exception with a scratch limit well below
    its required scratch space which in this case is 128m.
    """
    exec_option = vector.get_value('exec_option')
    exec_option['buffer_pool_limit'] = self.buffer_pool_limit
    exec_option['scratch_limit'] = '24m'
    expected_error = 'Scratch space limit of %s bytes exceeded'
    scratch_limit_in_bytes = 24 * 1024 * 1024
    try:
      self.execute_query(self.spilling_sort_query, exec_option)
      assert False, "Query was expected to fail"
    except ImpalaBeeswaxException as e:
      assert expected_error % scratch_limit_in_bytes in str(e)

  def test_with_zero_scratch_limit(self, vector):
    """
    Queries throws the appropriate exception with a scratch limit of
    zero which means no scratch space can be allocated.
    """
    exec_option = vector.get_value('exec_option')
    exec_option['buffer_pool_limit'] = self.buffer_pool_limit
    exec_option['scratch_limit'] = '0'
    for query in self.spilling_queries:
      self.execute_query_expect_failure(query, exec_option)

  def test_with_unlimited_scratch_limit(self, vector):
    """
    Sort query runs to completion with a scratch Limit of -1 means default/no limit.
    """
    exec_option = vector.get_value('exec_option')
    exec_option['buffer_pool_limit'] = self.buffer_pool_limit
    exec_option['scratch_limit'] = '-1'
    self.execute_query_expect_success(self.client, self.spilling_sort_query, exec_option)

  def test_without_specifying_scratch_limit(self, vector):
    """
    Sort query runs to completion with the default setting of no scratch limit.
    """
    exec_option = vector.get_value('exec_option')
    exec_option['buffer_pool_limit'] = self.buffer_pool_limit
    self.execute_query_expect_success(self.client, self.spilling_sort_query, exec_option)

  def test_with_zero_scratch_limit_no_memory_limit(self, vector):
    """
    Queries run to completion without spilling as there is no limit on block memory manger.
    Scratch limit of zero ensures spilling is disabled.
    """
    exec_option = vector.get_value('exec_option')
    exec_option['scratch_limit'] = '0'
    for query in self.spilling_queries:
      self.execute_query_expect_success(self.client, query, exec_option)
<|endoftext|>


<commit_before>#!/usr/bin/python

from itaps import iMesh, iBase
import argparse
from yt.utilities.lib.geometry_utils import triangle_plane_intersect
import numpy as np
import matplotlib.cm as colormap
from matplotlib.collections import PatchCollection
from matplotlib.patches import Polygon
from matplotlib.path import Path
from matplotlib.patches import PathPatch
import matplotlib.pyplot as plt
from pylab import *

import time
mesh = iMesh.Mesh()

CCW = 1
CW = -1

########

def point_match(pnt1, pnt2):

    b = False
    x1, y1, z1 = round(pnt1[0],6),round(pnt1[1],6),round(pnt1[2],6)
    x2, y2, z2 = round(pnt2[0],6),round(pnt2[1],6),round(pnt2[2],6)
    if x1 == x2 and y1 == y2 and z1 == z2:
        b = True
    return b

def surface_intersections(tris, axis, coord):

    lines=[]
    #get all intersections for these surface triangles
    for tri in tris:
        intersect, line = intersection(axis, coord, tri)
        if intersect:
            lines.append(line[0])

    #time to start building intersections
    intersections=[]

    while len(lines) !=0:
        #arbitrarily start a new intersection
        current_intersection = lines[0]
        del lines[0]

        i=0
        while i < len(lines):
            line = lines[i]
            #insers points in the appropriate place if there is a match at the front or the back
            if point_match(line[0],current_intersection[0]):
                current_intersection = np.insert(current_intersection, 0, line[1], 0)
                del lines[i]
                i=0
            elif point_match(line[0],current_intersection[-1]):
                current_intersection = np.append(current_intersection, [line[1]], 0)
                del lines[i]
                i=0
            elif point_match(line[1], current_intersection[0]):
                current_intersection = np.insert(current_intersection, 0, line[0], 0)
                del lines[i]
                i=0
            elif point_match(line[1],current_intersection[-1]):
                current_intersection = np.append(current_intersection, [line[0]], axis=0)
                del lines[i]
                i=0
            else:
                i+=1
        intersections.append(current_intersection)

    return intersections

def intersection(axis, coord, triangle):

    #create an array for the coordinates
    triangle_verts=np.array([],ndmin=3)
    triangle_verts.shape =(0,3,3)

    verts = mesh.getEntAdj(triangle,iBase.Type.vertex)
    vert1 = mesh.getVtxCoords(verts[0])
    vert2 = mesh.getVtxCoords(verts[1])
    vert3 = mesh.getVtxCoords(verts[2])

    temp = np.vstack((vert1,vert2,vert3))
    #insert the coordinates into the array
    triangle_verts = np.append(triangle_verts,[temp],axis = 0)

    #check for an intersection
    line = triangle_plane_intersect(axis,coord,triangle_verts)

    #if a line is returned, indicate that we have an intersection
    intersect = True if line.size is not 0 else False

    return intersect, line

def insert_pnt( line, coll):

    if point_match(line[0],coll[0]):
        coll = np.insert(coll, 0, line[1], 0)
    elif point_match(line[0],coll[-1]):
        coll = np.append(coll, [line[1]], 0)
    elif point_match(line[1], coll[0]):
        coll = np.insert(coll, 0, line[0], 0)
    elif point_match(line[1],coll[-1]):
        coll = np.append(coll, [line[0]], axis=0)
    else:
        print "Warning: no points coincide in this tri, moving on"

    return coll

#########

def get_surfaces():

    #get all the entsets
    sets = mesh.getEntSets(1)

    surfs = []
    #filter out the surfaces using the CATEGORY tag
    for set in sets:
        tags = mesh.getAllTags(set)

        for tag in tags:
            if tag.name == "CATEGORY":
                tag_handle = mesh.getTagHandle(tag.name)
                category_type = filter(lambda null: null != 0 , tag_handle[set])
                category_type = ''.join(chr(item) for item in category_type)
                if category_type == "Surface": surfs.append(set)

    print "There are " + str(len(surfs)) + " surfaces in this model."

    return surfs

def get_all_volumes():

    #get all the entsets
    sets = mesh.getEntSets(0)

    vols = []
    #filter out all the volumes using the CATEGORY tag
    for set in sets:
        tags = mesh.getAllTags(set)

        for tag in tags:
            if tag.name == "CATEGORY":
                tag_handle = mesh.getTagHandle(tag.name)
                category_type = filter(lambda null: null != 0 , tag_handle[set])
                category_type = ''.join(chr(item) for item in category_type)
                if category_type == "Volume": vols.append(set)

    print "There are " + str(len(vols)) + " volumes in this model."
    return vols

def get_vol_intersections(volume, intersect_dict):

    #get the surfaces for this volume
    surfs = volume.getChildren(0)

    intersections = []
    # return all intersections from the dictionary with
    # the surfaces of the volume as the key
    for surf in surfs:
        intersections += intersect_dict[surf]

    return intersections

# this function assumes that each plane-volume intersection will result
# in some number of complete loops
def stitch(intersections):

    colls = []
    #first, check for complete loops
    i=0
    while i < len(intersections):
        intersection = intersections[i]
        # add if this is full loop
        if point_match(intersection[0],intersection[-1]) and len(intersection) != 2:
            colls.append(intersection)
            del intersections[i]
            i=0
        # if this is a segment with a near-zerio length, remove it
        elif point_match(intersection[0],intersection[-1]) and len(intersection) == 2:
            del intersections[i]
            i=0
        i+=1

    if 0 == len(intersections):
        return colls

    #add the last intersection (arbitrary starting point)
    colls.append(intersections.pop())

    #do until all intersections are matched
    while len(intersections) != 0:

        #if the current collection is a loop, move to
        #the next arb starting point
        if point_match(colls[-1][0],colls[-1][-1]):
            colls.append(intersections.pop())

        else:
            i=0
            while i < len(intersections):
                intersection=intersections[i]
                if point_match(colls[-1][0],intersection[0]):
                    #reverse and attach to the front of colls[-1]
                    colls[-1]=np.append(intersection[::-1],colls[-1],axis=0)
                    del intersections[i]
                    i=0
                elif point_match(colls[-1][0], intersection[-1]):
                    #attach to the front of colls[-1]
                    colls[-1]=np.append(intersection,colls[-1],axis=0)
                    #                   print type(intersection)
                    del intersections[i]
                    i=0
                elif point_match(colls[-1][-1], intersection[0]):
                    #attach to the back of colls[-1]
                    colls[-1]=np.append(colls[-1],intersection,axis=0)
                    del intersections[i]
                    i=0
                elif point_match(colls[-1][-1], intersection[-1]):
                    #reverse and attach to the back of colls[-1]
                    colls[-1]=np.append(colls[-1],intersection[::-1],axis=0)
                    del intersections[i]
                    i=0
                #if no match is found, move to next intersection
                i+=1

    return colls

def parsing():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '-f', action='store', dest='filename', help='The path to the .h5m file')

    args = parser.parse_args()

    if not args.filename:
        raise Exception('h5m file path not specified!!. [-f] is not set')

    return args

# This function sets up coding for the path such that interior loops
# will not be filled for a volume cross-section
def return_coding(ob):
    # The codes will be all "LINETO" commands, except for "MOVETO"s at the
    # beginning of each subpath
    n = len(ob)
    codes = np.ones(n, dtype=Path.code_type) * Path.LINETO
    codes[0] = Path.MOVETO
    return codes

def create_surface_intersections(surfs, axis, coord):

    intersection_dict={}
    for surf in surfs:
        # get the surface's triangles
        surf_tris = surf.getEntities(iBase.Type.all, iMesh.Topology.triangle)
        #print "Retrieved " + str(len(surf_tris)) + " triangles from a surface set."
        # generate the surface intersections
        surf_intersections = surface_intersections(surf_tris, axis, coord )
        #add the surface's entry to the dictionary
        intersection_dict[surf] = surf_intersections
    return intersection_dict

def slice_faceted_model(filename, coord, axis):

    mesh.load(filename)

    #get all surfaces in the file
    surfs = get_surfaces()

    #get the intersections on each surface and store in dict with surfs as the keys
    intersection_dict = create_surface_intersections(surfs, axis, coord)

    #get all the volumes
    vols = get_all_volumes()

    groups = [vols,vols]

    all_group_paths=[]
    for group in groups:
        group_paths = get_group_volume_paths(vols, axis, intersection_dict)
        all_group_paths.append(group_paths)

    return all_group_paths

def get_group_volume_paths(vols, axis, intersection_dict):

    all_coordinates=[]
    all_codes=[]
    orient_time=0
    for vol in vols:

        #get the intersections for this volume based on its child surfaces
        intersects = get_vol_intersections(vol, intersection_dict)

        #if no intersections are returned, move on to the next volume
        if len(intersects) == 0: continue
        print "Retrieved "+str(len(intersects))+" intersections for this volume."

        #order the intersections into loops
        loops = stitch(intersects)
        print "Found "+str(len(loops))+" poly collections for this volume."

        #remove the axis of intersection from the points
        loops = [np.delete(loop,axis,1) for loop in loops]

        if __name__ == "__main__":
            print "Re-orienting intersections..."
        start = time.clock()
        #orient the loops for proper plot fills
        loops = orient_loops(loops)
        orient_time += (time.clock()-start)

        #Reformat
        #rearrange coords into one long list and remove the coordinates for the slice
        coordinates = np.concatenate(loops[:],axis=0)

        #generate coding for the path that will allow for interior loops (see return_coding)
        codes=np.concatenate([return_coding(loop) for loop in loops])

        #add this volume's info to the global list
        all_coordinates.append(coordinates)
        all_codes.append(codes)
        coordinates=[]
        codes=[]

    if __name__ == "__main__":
        print "Took " + str(orient_time) + " seconds to reorient loops."
    return all_coordinates, all_codes

def orient_loops(loops):

    #first we will create a set of paths to represent the loops
    paths = []
    for loop in loops:
        paths.append(Path(loop))

    #now we will generate a containment matrix for the loops
    M = gen_containment(paths)
    #determine the current windings
    current_windings = get_windings(paths)
    #get the desired windings
    desired_windings = get_fill_windings(M)
    #alter the current windings to match the desired windings
    loops = set_windings(current_windings, desired_windings, loops)
    return loops

def set_windings(current_windings, desired_windings, loops):

    n = len(current_windings)
    assert(len(current_windings)==len(desired_windings))
    assert(len(current_windings)==len(loops))

    for i in range(n):

        if current_windings[i] != desired_windings[i]:
            loops[i]=loops[i][::-1]

    return loops

def find_winding(path):

    area = 0
    verts = path.vertices
    j = len(verts) -1
    for i in range(len(verts)-1):
        area += (verts[j,0]+verts[i,0]) * (verts[j,1]-verts[i,1])
        j=i

    return CW if area >= 0 else CCW

def get_windings(paths):

    windings=[]
    for path in paths:
        winding = find_winding(path)
        windings.append(winding)
    return windings

def gen_containment(paths):

    n = len(paths)
    mat = np.empty([n,n])
    for i in range(n):
        for j in range(n):
            mat[i,j] = 1 if paths[j].contains_path(paths[i]) else 0

    return mat

def get_fill_windings(M):

    a,b = M.shape
    assert(a==b)
    #use the sum of the rows of M to determine windings for path[i]
    windings=[]
    # 1 indicates CCW; -1 indicates CW
    for i in range(a):
        wind = CCW if sum(M[i])%2 == 0 or sum(M[i]) == 0 else CW
        windings.append(wind)

    return windings

def main():

    #parse arguments and load the file
    args = parsing()

    axis = 1
    coord = 5.0
    group_patches = slice_faceted_model(args.filename, coord, axis)
    patches=[]
    collections=[]

    if __name__ == "__main__":
        print "Plotting..."

    #create patches for the plot
    for patch in group_patches:
        coords = patch[0]
        codes = patch[1]
        #name = group[2]
        #select color here
        for i in range(len(codes)):
            path = Path(coords[i], codes[i])
            patches.append(PathPatch(path))

        colors = 100*np.random.rand(len(patches))
        p = PatchCollection(patches, alpha=0.4)
        collections.append(p)
        p.set_array(np.array(colors))

    #create a new figure
    fig, ax = plt.subplots()

    for collection in collections:
        ax.add_collection(collection)
    #show the plot!
    ax.autoscale_view()
    plt.show()

if __name__ == "__main__":
    start = time.clock()
    main()
    print "Took " + str((time.clock()-start)) + "seconds total."

<commit_msg>Updated to plotting by group/material if asked for<commit_after>#!/usr/bin/python

from itaps import iMesh, iBase
import argparse
from yt.utilities.lib.geometry_utils import triangle_plane_intersect
import numpy as np
import matplotlib.cm as colormap
from matplotlib.collections import PatchCollection
from matplotlib.patches import Polygon
from matplotlib.path import Path
from matplotlib.patches import PathPatch
import matplotlib.pyplot as plt
from pylab import *
from random import random

import time
mesh = iMesh.Mesh()

CCW = 1
CW = -1

########

def point_match(pnt1, pnt2):

    b = False
    x1, y1, z1 = round(pnt1[0],6),round(pnt1[1],6),round(pnt1[2],6)
    x2, y2, z2 = round(pnt2[0],6),round(pnt2[1],6),round(pnt2[2],6)
    if x1 == x2 and y1 == y2 and z1 == z2:
        b = True
    return b

def surface_intersections(tris, axis, coord):

    lines=[]
    #get all intersections for these surface triangles
    for tri in tris:
        intersect, line = intersection(axis, coord, tri)
        if intersect:
            lines.append(line[0])

    #time to start building intersections
    intersections=[]

    while len(lines) !=0:
        #arbitrarily start a new intersection
        current_intersection = lines[0]
        del lines[0]

        i=0
        while i < len(lines):
            line = lines[i]
            #insers points in the appropriate place if there is a match at the front or the back
            if point_match(line[0],current_intersection[0]):
                current_intersection = np.insert(current_intersection, 0, line[1], 0)
                del lines[i]
                i=0
            elif point_match(line[0],current_intersection[-1]):
                current_intersection = np.append(current_intersection, [line[1]], 0)
                del lines[i]
                i=0
            elif point_match(line[1], current_intersection[0]):
                current_intersection = np.insert(current_intersection, 0, line[0], 0)
                del lines[i]
                i=0
            elif point_match(line[1],current_intersection[-1]):
                current_intersection = np.append(current_intersection, [line[0]], axis=0)
                del lines[i]
                i=0
            else:
                i+=1
        intersections.append(current_intersection)

    return intersections

def intersection(axis, coord, triangle):

    #create an array for the coordinates
    triangle_verts=np.array([],ndmin=3)
    triangle_verts.shape =(0,3,3)

    verts = mesh.getEntAdj(triangle,iBase.Type.vertex)
    vert1 = mesh.getVtxCoords(verts[0])
    vert2 = mesh.getVtxCoords(verts[1])
    vert3 = mesh.getVtxCoords(verts[2])

    temp = np.vstack((vert1,vert2,vert3))
    #insert the coordinates into the array
    triangle_verts = np.append(triangle_verts,[temp],axis = 0)

    #check for an intersection
    line = triangle_plane_intersect(axis,coord,triangle_verts)

    #if a line is returned, indicate that we have an intersection
    intersect = True if line.size is not 0 else False

    return intersect, line

def insert_pnt( line, coll):

    if point_match(line[0],coll[0]):
        coll = np.insert(coll, 0, line[1], 0)
    elif point_match(line[0],coll[-1]):
        coll = np.append(coll, [line[1]], 0)
    elif point_match(line[1], coll[0]):
        coll = np.insert(coll, 0, line[0], 0)
    elif point_match(line[1],coll[-1]):
        coll = np.append(coll, [line[0]], axis=0)
    else:
        print "Warning: no points coincide in this tri, moving on"

    return coll

#########

def get_surfaces():

    #get all the entsets
    sets = mesh.getEntSets(1)

    surfs = []
    #filter out the surfaces using the CATEGORY tag
    for set in sets:
        tags = mesh.getAllTags(set)

        for tag in tags:
            if tag.name == "CATEGORY":
                tag_handle = mesh.getTagHandle(tag.name)
                category_type = filter(lambda null: null != 0 , tag_handle[set])
                category_type = ''.join(chr(item) for item in category_type)
                if category_type == "Surface": surfs.append(set)

    print "There are " + str(len(surfs)) + " surfaces in this model."

    return surfs

def get_vols_by_group():

    mesh_sets=mesh.getEntSets()

    all_vols = get_all_volumes()

    group_vols=[]
    group_names=[]
    for s in mesh_sets:
        tags=mesh.getAllTags(s)
        for t in tags:

            if t.name == 'NAME':
                group_name = tag_to_script(t[s])
                group_names.append(group_name)
                #assume only volumes in groups for now
                vols = []
                for vol in all_vols:
                    if s.contains(vol): vols.append(vol)

                group_vols.append(vols)
                print vols

    return group_vols, group_names

"""
function to transform the tags into strings
tag : string of the tag to add to tag_list
tag_list : vector of tags in the problem
returns tag_list
"""
def tag_to_script(tag):
    a = []
    # since we have a byte type tag loop over the 32 elements
    for part in tag:
        # if the byte char code is non 0
        if (part != 0):
            # convert to ascii
            a.append(str(unichr(part)))
            # join to end string
            test = ''.join(a)
            # the the string we are testing for is not in the list of found
            # tag values, add to the list of tag_values
    # if not already in list append to list
    #    if not any(test in s for s in tag_list):
    # the original code was incorrectly missing groups when one of the same
    # name with/without rho was added
    return test

def get_all_volumes():

    #get all the entsets
    sets = mesh.getEntSets(0)

    vols = []
    #filter out all the volumes using the CATEGORY tag
    for set in sets:
        tags = mesh.getAllTags(set)

        for tag in tags:
            if tag.name == "CATEGORY":
                tag_handle = mesh.getTagHandle(tag.name)
                category_type = filter(lambda null: null != 0 , tag_handle[set])
                category_type = ''.join(chr(item) for item in category_type)
                if category_type == "Volume": vols.append(set)

    print "There are " + str(len(vols)) + " volumes in this model."
    return vols

def get_vol_intersections(volume, intersect_dict):

    #get the surfaces for this volume
    surfs = volume.getChildren(0)

    intersections = []
    # return all intersections from the dictionary with
    # the surfaces of the volume as the key
    for surf in surfs:
        intersections += intersect_dict[surf]

    return intersections

# this function assumes that each plane-volume intersection will result
# in some number of complete loops
def stitch(intersections):

    colls = []
    #first, check for complete loops
    i=0
    while i < len(intersections):
        intersection = intersections[i]
        # add if this is full loop
        if point_match(intersection[0],intersection[-1]) and len(intersection) != 2:
            colls.append(intersection)
            del intersections[i]
            i=0
        # if this is a segment with a near-zerio length, remove it
        elif point_match(intersection[0],intersection[-1]) and len(intersection) == 2:
            del intersections[i]
            i=0
        i+=1

    if 0 == len(intersections):
        return colls

    #add the last intersection (arbitrary starting point)
    colls.append(intersections.pop())

    #do until all intersections are matched
    while len(intersections) != 0:

        #if the current collection is a loop, move to
        #the next arb starting point
        if point_match(colls[-1][0],colls[-1][-1]):
            colls.append(intersections.pop())

        else:
            i=0
            while i < len(intersections):
                intersection=intersections[i]
                if point_match(colls[-1][0],intersection[0]):
                    #reverse and attach to the front of colls[-1]
                    colls[-1]=np.append(intersection[::-1],colls[-1],axis=0)
                    del intersections[i]
                    i=0
                elif point_match(colls[-1][0], intersection[-1]):
                    #attach to the front of colls[-1]
                    colls[-1]=np.append(intersection,colls[-1],axis=0)
                    #                   print type(intersection)
                    del intersections[i]
                    i=0
                elif point_match(colls[-1][-1], intersection[0]):
                    #attach to the back of colls[-1]
                    colls[-1]=np.append(colls[-1],intersection,axis=0)
                    del intersections[i]
                    i=0
                elif point_match(colls[-1][-1], intersection[-1]):
                    #reverse and attach to the back of colls[-1]
                    colls[-1]=np.append(colls[-1],intersection[::-1],axis=0)
                    del intersections[i]
                    i=0
                #if no match is found, move to next intersection
                i+=1

    return colls

def parsing():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '-f', action='store', dest='filename', help='The path to the .h5m file')

    args = parser.parse_args()

    if not args.filename:
        raise Exception('h5m file path not specified!!. [-f] is not set')

    return args

# This function sets up coding for the path such that interior loops
# will not be filled for a volume cross-section
def return_coding(ob):
    # The codes will be all "LINETO" commands, except for "MOVETO"s at the
    # beginning of each subpath
    n = len(ob)
    codes = np.ones(n, dtype=Path.code_type) * Path.LINETO
    codes[0] = Path.MOVETO
    return codes

def create_surface_intersections(surfs, axis, coord):

    intersection_dict={}
    for surf in surfs:
        # get the surface's triangles
        surf_tris = surf.getEntities(iBase.Type.all, iMesh.Topology.triangle)
        #print "Retrieved " + str(len(surf_tris)) + " triangles from a surface set."
        # generate the surface intersections
        surf_intersections = surface_intersections(surf_tris, axis, coord )
        #add the surface's entry to the dictionary
        intersection_dict[surf] = surf_intersections
    return intersection_dict

def slice_faceted_model(filename, coord, axis, by_group=False):

    mesh.load(filename)

    #get all surfaces in the file
    surfs = get_surfaces()

    #get the intersections on each surface and store in dict with surfs as the keys
    intersection_dict = create_surface_intersections(surfs, axis, coord)

    if by_group:
        groups, group_names =get_vols_by_group()
    else:
        #get all the volumes
        vols = get_all_volumes()
        groups = [vols]

    all_group_paths=[]
    for group_vols in groups:
        group_paths = get_group_volume_paths(group_vols, axis, intersection_dict)
        all_group_paths.append(group_paths)

    return all_group_paths

def get_group_volume_paths(vols, axis, intersection_dict):

    all_coordinates=[]
    all_codes=[]
    orient_time=0
    for vol in vols:

        #get the intersections for this volume based on its child surfaces
        intersects = get_vol_intersections(vol, intersection_dict)

        #if no intersections are returned, move on to the next volume
        if len(intersects) == 0: continue
        print "Retrieved "+str(len(intersects))+" intersections for this volume."

        #order the intersections into loops
        loops = stitch(intersects)
        print "Found "+str(len(loops))+" poly collections for this volume."

        #remove the axis of intersection from the points
        loops = [np.delete(loop,axis,1) for loop in loops]

        if __name__ == "__main__":
            print "Re-orienting intersections..."
        start = time.clock()
        #orient the loops for proper plot fills
        loops = orient_loops(loops)
        orient_time += (time.clock()-start)

        #Reformat
        #rearrange coords into one long list and remove the coordinates for the slice
        coordinates = np.concatenate(loops[:],axis=0)

        #generate coding for the path that will allow for interior loops (see return_coding)
        codes=np.concatenate([return_coding(loop) for loop in loops])

        #add this volume's info to the global list
        all_coordinates.append(coordinates)
        all_codes.append(codes)
        coordinates=[]
        codes=[]

    if __name__ == "__main__":
        print "Took " + str(orient_time) + " seconds to reorient loops."
    return all_coordinates, all_codes

def orient_loops(loops):

    #first we will create a set of paths to represent the loops
    paths = []
    for loop in loops:
        paths.append(Path(loop))

    #now we will generate a containment matrix for the loops
    M = gen_containment(paths)
    #determine the current windings
    current_windings = get_windings(paths)
    #get the desired windings
    desired_windings = get_fill_windings(M)
    #alter the current windings to match the desired windings
    loops = set_windings(current_windings, desired_windings, loops)
    return loops

def set_windings(current_windings, desired_windings, loops):

    n = len(current_windings)
    assert(len(current_windings)==len(desired_windings))
    assert(len(current_windings)==len(loops))

    for i in range(n):

        if current_windings[i] != desired_windings[i]:
            loops[i]=loops[i][::-1]

    return loops

def find_winding(path):

    area = 0
    verts = path.vertices
    j = len(verts) -1
    for i in range(len(verts)-1):
        area += (verts[j,0]+verts[i,0]) * (verts[j,1]-verts[i,1])
        j=i

    return CW if area >= 0 else CCW

def get_windings(paths):

    windings=[]
    for path in paths:
        winding = find_winding(path)
        windings.append(winding)
    return windings

def gen_containment(paths):

    n = len(paths)
    mat = np.empty([n,n])
    for i in range(n):
        for j in range(n):
            mat[i,j] = 1 if paths[j].contains_path(paths[i]) else 0

    return mat

def get_fill_windings(M):

    a,b = M.shape
    assert(a==b)
    #use the sum of the rows of M to determine windings for path[i]
    windings=[]
    # 1 indicates CCW; -1 indicates CW
    for i in range(a):
        wind = CCW if sum(M[i])%2 == 0 or sum(M[i]) == 0 else CW
        windings.append(wind)

    return windings

def main():

    #parse arguments and load the file
    args = parsing()

    axis = 1
    coord = 0.0
    group_patches = slice_faceted_model(args.filename, coord, axis, by_group=True)
    patches=[]
    collections=[]

    if __name__ == "__main__":
        print "Plotting..."

    #create patches for the plot
    for patch in group_patches:
        coords = patch[0]
        codes = patch[1]
        #name = group[2]
        #select color here
        color = np.random.rand(3,1)
        for i in range(len(codes)):
            path = Path(coords[i], codes[i])
            patches.append(PathPatch(path, color=color, ec='black', lw=1, ls='solid'))

    p = PatchCollection(patches, match_original=True, alpha=0.5)

    #create a new figure
    fig, ax = plt.subplots()

    #for collection in collections:
    ax.add_collection(p)

    #show the plot!
    ax.autoscale_view()
    plt.show()

if __name__ == "__main__":
    start = time.clock()
    main()
    print "Took " + str((time.clock()-start)) + " seconds total."

<|endoftext|>


<commit_before>import requests

def safe_urlopen(
    url,
    method=None,
    params=None,
    data=None,
    json=None,
    headers=None,
    allow_redirects=False,
    timeout=30,
    verify_ssl=True,
):
    """
    A slightly safer version of ``urlib2.urlopen`` which prevents redirection
    and ensures the URL isn't attempting to hit a blacklisted IP range.
    """

    session = requests.Session()

    kwargs = {}

    if json:
        kwargs['json'] = json
        if not headers:
            headers = {}
        headers.setdefault('Content-Type', 'application/json')

    if data:
        kwargs['data'] = data

    if params:
        kwargs['params'] = params

    if headers:
        kwargs['headers'] = headers

    if method is None:
        method = 'POST' if (data or json) else 'GET'

    response = session.request(
        method=method,
        url=url,
        allow_redirects=allow_redirects,
        timeout=timeout,
        verify=verify_ssl,
        **kwargs
    )

    return response
<commit_msg>Add download request<commit_after>import requests

def safe_urlopen(
    url,
    method=None,
    params=None,
    data=None,
    json=None,
    headers=None,
    allow_redirects=False,
    timeout=30,
    verify_ssl=True,
):
    """
    A slightly safer version of ``urlib2.urlopen`` which prevents redirection
    and ensures the URL isn't attempting to hit a blacklisted IP range.
    """

    session = requests.Session()

    kwargs = {}

    if json:
        kwargs['json'] = json
        if not headers:
            headers = {}
        headers.setdefault('Content-Type', 'application/json')

    if data:
        kwargs['data'] = data

    if params:
        kwargs['params'] = params

    if headers:
        kwargs['headers'] = headers

    if method is None:
        method = 'POST' if (data or json) else 'GET'

    response = session.request(
        method=method,
        url=url,
        allow_redirects=allow_redirects,
        timeout=timeout,
        verify=verify_ssl,
        **kwargs
    )

    return response

def download(url, access_token, filename, logger, headers=None, timeout=60):
    """Download the file from the given url at the current path"""
    logger.info("Downloading file from %s" % url)
    # Auth headers if access_token is present
    request_headers = {}
    if access_token:
        request_headers["Authorization"] = "Token " + access_token
    # Add any additional headers
    if headers:
        request_headers.update(headers)

    try:
        response = requests.get(url,
                                headers=request_headers,
                                timeout=timeout,
                                stream=True)

        if response.status_code != 200:
            logger("Failed to download file from %s: %s" % (url, response.status_code))
            return None

        with open(filename, 'wb') as f:
            for chunk in response.iter_content(chunk_size=1024):
                if chunk:
                    f.write(chunk)
        return filename
    except requests.exceptions.RequestException as e:
        logger.warning("Exception: %s" % e)
        return None
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# Import Python libs
from __future__ import absolute_import

# Import salt module
from salt.modules import mac_pkgutil

# Import Salt Testing libs
from salttesting import TestCase, skipIf
from salttesting.mock import NO_MOCK, NO_MOCK_REASON, MagicMock, patch

darwin_pkgutil.__salt__ = {}

@skipIf(NO_MOCK, NO_MOCK_REASON)
class DarwingPkgutilTestCase(TestCase):
    def test_list_installed_command(self):
        # Given
        r_output = "com.apple.pkg.iTunes"

        # When
        mock_cmd = MagicMock(return_value=r_output)
        with patch.dict(darwin_pkgutil.__salt__, {'cmd.run_stdout': mock_cmd}):
            output = darwin_pkgutil.list_()

        # Then
        mock_cmd.assert_called_with(
            ['/usr/sbin/pkgutil', '--pkgs'],
            python_shell=False
        )

    def test_list_installed_output(self):
        # Given
        r_output = "com.apple.pkg.iTunes"

        # When
        mock_cmd = MagicMock(return_value=r_output)
        with patch.dict(darwin_pkgutil.__salt__, {'cmd.run_stdout': mock_cmd}):
            output = darwin_pkgutil.list_()

        # Then
        self.assertEqual(output, r_output)

    def test_is_installed(self):
        # Given
        r_output = "com.apple.pkg.iTunes"

        # When
        mock_cmd = MagicMock(return_value=r_output)
        with patch.dict(darwin_pkgutil.__salt__, {'cmd.run_stdout': mock_cmd}):
            ret = darwin_pkgutil.is_installed("com.apple.pkg.iTunes")

        # Then
        self.assertTrue(ret)

        # When
        with patch.dict(darwin_pkgutil.__salt__, {'cmd.run_stdout': mock_cmd}):
            ret = darwin_pkgutil.is_installed("com.apple.pkg.Safari")

        # Then
        self.assertFalse(ret)

    def test_install(self):
        # Given
        source = "/foo/bar/fubar.pkg"
        package_id = "com.foo.fubar.pkg"

        # When
        mock_cmd = MagicMock(return_value=True)
        with patch("salt.modules.darwin_pkgutil.is_installed",
                   return_value=False) as is_installed:
            with patch("salt.modules.darwin_pkgutil._install_from_path",
                   return_value=True) as _install_from_path:
                ret = darwin_pkgutil.install(source, package_id)

        # Then
        _install_from_path.assert_called_with(source)

    def test_install_already_there(self):
        # Given
        source = "/foo/bar/fubar.pkg"
        package_id = "com.foo.fubar.pkg"

        # When
        mock_cmd = MagicMock(return_value=True)
        with patch("salt.modules.darwin_pkgutil.is_installed",
                   return_value=True) as is_installed:
            with patch("salt.modules.darwin_pkgutil._install_from_path",
                   return_value=True) as _install_from_path:
                ret = darwin_pkgutil.install(source, package_id)

        # Then
        self.assertEqual(_install_from_path.called, 0)
<commit_msg>More lint<commit_after># -*- coding: utf-8 -*-

# Import Python libs
from __future__ import absolute_import

# Import salt module
from salt.modules import mac_pkgutil

# Import Salt Testing libs
from salttesting import TestCase, skipIf
from salttesting.mock import NO_MOCK, NO_MOCK_REASON, MagicMock, patch

mac_pkgutil.__salt__ = {}

@skipIf(NO_MOCK, NO_MOCK_REASON)
class DarwingPkgutilTestCase(TestCase):
    def test_list_installed_command(self):
        # Given
        r_output = "com.apple.pkg.iTunes"

        # When
        mock_cmd = MagicMock(return_value=r_output)
        with patch.dict(darwin_pkgutil.__salt__, {'cmd.run_stdout': mock_cmd}):
            output = darwin_pkgutil.list_()

        # Then
        mock_cmd.assert_called_with(
            ['/usr/sbin/pkgutil', '--pkgs'],
            python_shell=False
        )

    def test_list_installed_output(self):
        # Given
        r_output = "com.apple.pkg.iTunes"

        # When
        mock_cmd = MagicMock(return_value=r_output)
        with patch.dict(darwin_pkgutil.__salt__, {'cmd.run_stdout': mock_cmd}):
            output = darwin_pkgutil.list_()

        # Then
        self.assertEqual(output, r_output)

    def test_is_installed(self):
        # Given
        r_output = "com.apple.pkg.iTunes"

        # When
        mock_cmd = MagicMock(return_value=r_output)
        with patch.dict(darwin_pkgutil.__salt__, {'cmd.run_stdout': mock_cmd}):
            ret = darwin_pkgutil.is_installed("com.apple.pkg.iTunes")

        # Then
        self.assertTrue(ret)

        # When
        with patch.dict(darwin_pkgutil.__salt__, {'cmd.run_stdout': mock_cmd}):
            ret = darwin_pkgutil.is_installed("com.apple.pkg.Safari")

        # Then
        self.assertFalse(ret)

    def test_install(self):
        # Given
        source = "/foo/bar/fubar.pkg"
        package_id = "com.foo.fubar.pkg"

        # When
        mock_cmd = MagicMock(return_value=True)
        with patch("salt.modules.darwin_pkgutil.is_installed",
                   return_value=False) as is_installed:
            with patch("salt.modules.darwin_pkgutil._install_from_path",
                   return_value=True) as _install_from_path:
                ret = darwin_pkgutil.install(source, package_id)

        # Then
        _install_from_path.assert_called_with(source)

    def test_install_already_there(self):
        # Given
        source = "/foo/bar/fubar.pkg"
        package_id = "com.foo.fubar.pkg"

        # When
        mock_cmd = MagicMock(return_value=True)
        with patch("salt.modules.darwin_pkgutil.is_installed",
                   return_value=True) as is_installed:
            with patch("salt.modules.darwin_pkgutil._install_from_path",
                   return_value=True) as _install_from_path:
                ret = darwin_pkgutil.install(source, package_id)

        # Then
        self.assertEqual(_install_from_path.called, 0)
<|endoftext|>


<commit_before>from __future__ import unicode_literals

from smartmin.models import SmartModel
from django.db import models
from django.db.models import Avg, Sum, Q

WEEK_AVG = 8
AVG_GAMES = 72

class Player(SmartModel):
    name = models.CharField(max_length=128,
                            help_text="The player's name")

    def get_games(self):
        scores = list(PlayerScore.objects.filter(player=self).order_by('-match__date'))

        # create a dict of match / game to score
        match_games = dict()
        for score in scores:
            match_games['%d_%d' % (score.match_id, score.game)] = [score]

        q = Q(match_id__lt=0)
        for score in scores:
            q |= (Q(match=score.match, game=score.game) & ~Q(player=self))

        for opp_score in PlayerScore.objects.filter(q):
            match_games['%d_%d' % (opp_score.match_id, opp_score.game)].append(opp_score)

        games = []
        for score in scores:
            opp_score = match_games['%d_%d' % (score.match_id, score.game)][1]
            games.append(dict(date=score.match.date, score=score.score, opponent=opp_score.player.name, opp_score=opp_score.score))

        return games

    def get_week_averages(self):
        scores = list(PlayerScore.objects.filter(player=self).order_by('match__date'))

        weeks = list()
        current_week = dict()

        for score in scores:
            if not current_week or score.match.date != current_week['date']:
                if current_week:
                    current_week['avg'] = float(current_week['total']) / current_week['games']
                    weeks.append(current_week)

                current_week = dict(date=score.match.date, total=score.score, games=1)

            else:
                current_week['total'] += score.score
                current_week['games'] += 1

        if current_week:
            current_week['avg'] = float(current_week['total']) / current_week['games']
            weeks.append(current_week)

        # now build the week avg for each week
        for i in range(len(weeks)):
            game_count = sum(w['games'] for w in weeks[max(0, i-WEEK_AVG):i+1])
            game_total = sum(w['total'] for w in weeks[max(0, i-WEEK_AVG):i+1])
            weeks[i]['running_avg'] = game_total / float(game_count)

        return weeks

    def avg(self, before_date):
        last = list(PlayerScore.objects.filter(match__date__lt=before_date, player=self).order_by('-match__date')[:AVG_GAMES])

        # less than 6 games? default to average
        if len(last) < 6:
            return 7
        else:
            # build up our scores and sort them
            scores = sorted([ps.score for ps in last])

            # figure out our eighth size
            eighth = len(scores) / 8

            # strip off the top and bottom quartile
            scores = scores[eighth:-eighth]

            return float(sum(scores)) / len(scores)

    def set_season(self, season):
        self.season = season

        # get last games
        last = list(PlayerScore.objects.filter(player=self).order_by('-match__date')[:AVG_GAMES])

        self.wins = sum([1 if ps.score == 10 else 0 for ps in last])
        self.losses = sum([1 if ps.score < 10 else 0 for ps in last])

        self.games = len(last)
        self.points = sum([ps.score for ps in last])

        # build up our scores and sort them
        scores = sorted([ps.score for ps in last])

        # figure out our quartile size
        eighth = len(scores) / 8 if len(scores) > 0 else 0

        # strip off the top and bottom quartile
        scores = scores[eighth:-eighth]
        self.avg = float(sum(scores)) / len(scores) if len(scores) > 0 else 0

        q = Q(match_id__lt=0)
        for score in last:
            q |= (Q(match=score.match, game=score.game) & ~Q(player=self))

        last_opp = list(PlayerScore.objects.filter(q))

        self.opponent_points = sum([ps.score for ps in last_opp])
        self.mpg = float(self.points - self.opponent_points) / self.games if self.games > 0 else 0

    def games(self):
        return PlayerScore.objects.filter(match__season=self.season, player=self).count()

    def __unicode__(self):
        return self.name

    class Meta:
        ordering = ('name',)

class Season(SmartModel):
    name = models.CharField(max_length=128,
                            help_text="Season Name")
    start_date = models.DateField(help_text="When the season started")
    end_date = models.DateField(help_text="When the season ended")

    def __unicode__(self):
        return self.name

class Team(SmartModel):
    season = models.ForeignKey(Season)
    name = models.CharField(max_length=128,
                            help_text="The team name")

    @classmethod
    def season_summary(cls, season):
        teams = dict()
        for team in Team.objects.filter(season=season).distinct():
            team.points = 0
            team.games = 0
            team.wins = 0
            team.losses = 0
            teams[team] = team

        for match in Match.objects.filter(season=season):
            games = match.scores.count() / 2
            team1 = teams[match.team1]
            team2 = teams[match.team2]

            team1.points += match.points1
            team1_wins = match.scores.filter(team=match.team1, score=10).count()
            team1_losses = games - team1_wins
            team1.wins += team1_wins
            team1.losses += team1_losses
            team1.games += games

            team2.points += match.points2
            team2_wins = match.scores.filter(team=match.team2, score=10).count()
            team2_losses = games - team2_wins
            team2.wins += team2_wins
            team2.losses += team2_losses
            team2.games += games

        # sort our values by points
        return sorted(teams.values(), key=lambda t: t.points, reverse=True)

    def __unicode__(self):
        return self.name

class Match(SmartModel):
    season = models.ForeignKey(Season, related_name='matches')

    team1 = models.ForeignKey(Team, related_name='team1_matches')
    team2 = models.ForeignKey(Team, related_name='team2_matches')

    handicap1 = models.IntegerField(default=0)
    handicap2 = models.IntegerField(default=0)

    points1 = models.DecimalField(max_digits=3, decimal_places=1, default=0)
    points2 = models.DecimalField(max_digits=3, decimal_places=1, default=0)

    date = models.DateField()

    def sum_for_players(self, players):
        player_sum = 0
        for player in players:
            player_avg = player.avg(self.date)
            print "%s: %f" % (player.name, player_avg)
            player_sum += player_avg

        return player_sum

    def calculate_handicaps(self):
        # calculate average score for team 1 players for this season before this date
        print "Team1----------"
        team1_sum = self.sum_for_players(Player.objects.filter(scores__match=self, scores__team=self.team1).distinct())
        print team1_sum
        print "Team2----------"
        team2_sum = self.sum_for_players(Player.objects.filter(scores__match=self, scores__team=self.team2).distinct())
        print team2_sum

        if team1_sum > team2_sum:
            self.handicap1 = 0
            self.handicap2 = round(team1_sum - team2_sum - .01)
        else:
            self.handicap1 = round(team2_sum - team1_sum - .01)
            self.handicap2 = 0

    def summary(self):
        rounds = []
        match = dict(rounds=rounds, score1=0, score2=0, points1=0, points2=0)

        # are we 16 games or 25?
        last_game = PlayerScore.objects.filter(match=self).order_by('-game').first()
        games = 25
        if last_game:
            games = last_game.game

        if games <= 9:
            num_rounds = 3
        elif games <= 16:
            num_rounds = 4
        elif games <= 24:
            num_rounds = 5
        else:
            raise Exception("unknown number of games")

        for game_num in range(1, games+1):
            if (game_num - 1) % num_rounds == 0:
                games = []
                round = dict(games=games, round=len(rounds)+1, score1=0, score2=0)
                rounds.append(round)

            p1_score = PlayerScore.objects.filter(match=self, team=self.team1, game=game_num).first()
            p2_score = PlayerScore.objects.filter(match=self, team=self.team2, game=game_num).first()

            if p1_score and p2_score:
                round['score1'] += p1_score.score
                round['score2'] += p2_score.score

                games.append(dict(game=game_num, round=(game_num-1) / num_rounds,
                                  player1=p1_score.player, score1=p1_score.score,
                                  player2=p2_score.player, score2=p2_score.score))

        # add in handicaps
        for round in rounds:
            round['hscore1'] = round['score1'] + self.handicap1
            round['hscore2'] = round['score2'] + self.handicap2
            match['score1'] += round['hscore1']
            match['score2'] += round['hscore2']

        # calculate wins
        for round in rounds:
            if round['hscore1'] > round['hscore2']:
                match['points1'] += 1
                round['win1'] = True

            elif round['hscore2'] > round['hscore1']:
                match['points2'] += 1
                round['win2'] = True

            else:
                match['points1'] += .5
                match['points2'] += .5
                round['win1'] = True
                round['win2'] = True

        if match['score1'] > match['score2']:
            match['points1'] += 1
            match['score1_win'] = True

        elif match['score2'] > match['score1']:
            match['points2'] += 1
            match['score2_win'] = True

        if match['points1'] > match['points2']:
            match['win1'] = True
        elif match['points2'] > match['points1']:
            match['win2'] = True

        match['handicap1'] = self.handicap1
        match['handicap2'] = self.handicap2

        return match

    def calculate_stats(self):
        # calculate handicaps
        self.calculate_handicaps()

        # calculate wins
        summary = self.summary()
        self.points1 = summary['points1']
        self.points2 = summary['points2']
        self.save()

    def __unicode__(self):
        return "%s vs %s on %s" % (str(self.team1), str(self.team2), str(self.date))

    class Meta:
        verbose_name_plural = "Matches"

class PlayerScore(SmartModel):
    match = models.ForeignKey(Match, related_name='scores')
    game = models.IntegerField()

    team = models.ForeignKey(Team, related_name='scores')
    player = models.ForeignKey(Player, related_name='scores')
    score = models.IntegerField()

    def __unicode__(self):
        return "%s: %d" % (self.player, self.score)
<commit_msg>deal with 0 game players<commit_after>from __future__ import unicode_literals

from smartmin.models import SmartModel
from django.db import models
from django.db.models import Avg, Sum, Q

WEEK_AVG = 8
AVG_GAMES = 72

class Player(SmartModel):
    name = models.CharField(max_length=128,
                            help_text="The player's name")

    def get_games(self):
        scores = list(PlayerScore.objects.filter(player=self).order_by('-match__date'))

        # create a dict of match / game to score
        match_games = dict()
        for score in scores:
            match_games['%d_%d' % (score.match_id, score.game)] = [score]

        q = Q(match_id__lt=0)
        for score in scores:
            q |= (Q(match=score.match, game=score.game) & ~Q(player=self))

        for opp_score in PlayerScore.objects.filter(q):
            match_games['%d_%d' % (opp_score.match_id, opp_score.game)].append(opp_score)

        games = []
        for score in scores:
            opp_score = match_games['%d_%d' % (score.match_id, score.game)][1]
            games.append(dict(date=score.match.date, score=score.score, opponent=opp_score.player.name, opp_score=opp_score.score))

        return games

    def get_week_averages(self):
        scores = list(PlayerScore.objects.filter(player=self).order_by('match__date'))

        weeks = list()
        current_week = dict()

        for score in scores:
            if not current_week or score.match.date != current_week['date']:
                if current_week:
                    current_week['avg'] = float(current_week['total']) / current_week['games']
                    weeks.append(current_week)

                current_week = dict(date=score.match.date, total=score.score, games=1)

            else:
                current_week['total'] += score.score
                current_week['games'] += 1

        if current_week:
            current_week['avg'] = float(current_week['total']) / current_week['games']
            weeks.append(current_week)

        # now build the week avg for each week
        for i in range(len(weeks)):
            game_count = sum(w['games'] for w in weeks[max(0, i-WEEK_AVG):i+1])
            game_total = sum(w['total'] for w in weeks[max(0, i-WEEK_AVG):i+1])
            weeks[i]['running_avg'] = game_total / float(game_count)

        return weeks

    def avg(self, before_date):
        last = list(PlayerScore.objects.filter(match__date__lt=before_date, player=self).order_by('-match__date')[:AVG_GAMES])

        # less than 6 games? default to average
        if len(last) < 6:
            return 7
        else:
            # build up our scores and sort them
            scores = sorted([ps.score for ps in last])

            # figure out our eighth size
            eighth = len(scores) / 8

            # strip off the top and bottom quartile
            scores = scores[eighth:-eighth]

            return float(sum(scores)) / len(scores) if len(scores) > 0 else 0

    def set_season(self, season):
        self.season = season

        # get last games
        last = list(PlayerScore.objects.filter(player=self).order_by('-match__date')[:AVG_GAMES])

        self.wins = sum([1 if ps.score == 10 else 0 for ps in last])
        self.losses = sum([1 if ps.score < 10 else 0 for ps in last])

        self.games = len(last)
        self.points = sum([ps.score for ps in last])

        # build up our scores and sort them
        scores = sorted([ps.score for ps in last])

        # figure out our quartile size
        eighth = len(scores) / 8 if len(scores) > 0 else 0

        # strip off the top and bottom quartile
        scores = scores[eighth:-eighth]
        self.avg = float(sum(scores)) / len(scores) if len(scores) > 0 else 0

        q = Q(match_id__lt=0)
        for score in last:
            q |= (Q(match=score.match, game=score.game) & ~Q(player=self))

        last_opp = list(PlayerScore.objects.filter(q))

        self.opponent_points = sum([ps.score for ps in last_opp])
        self.mpg = float(self.points - self.opponent_points) / self.games if self.games > 0 else 0

    def games(self):
        return PlayerScore.objects.filter(match__season=self.season, player=self).count()

    def __unicode__(self):
        return self.name

    class Meta:
        ordering = ('name',)

class Season(SmartModel):
    name = models.CharField(max_length=128,
                            help_text="Season Name")
    start_date = models.DateField(help_text="When the season started")
    end_date = models.DateField(help_text="When the season ended")

    def __unicode__(self):
        return self.name

class Team(SmartModel):
    season = models.ForeignKey(Season)
    name = models.CharField(max_length=128,
                            help_text="The team name")

    @classmethod
    def season_summary(cls, season):
        teams = dict()
        for team in Team.objects.filter(season=season).distinct():
            team.points = 0
            team.games = 0
            team.wins = 0
            team.losses = 0
            teams[team] = team

        for match in Match.objects.filter(season=season):
            games = match.scores.count() / 2
            team1 = teams[match.team1]
            team2 = teams[match.team2]

            team1.points += match.points1
            team1_wins = match.scores.filter(team=match.team1, score=10).count()
            team1_losses = games - team1_wins
            team1.wins += team1_wins
            team1.losses += team1_losses
            team1.games += games

            team2.points += match.points2
            team2_wins = match.scores.filter(team=match.team2, score=10).count()
            team2_losses = games - team2_wins
            team2.wins += team2_wins
            team2.losses += team2_losses
            team2.games += games

        # sort our values by points
        return sorted(teams.values(), key=lambda t: t.points, reverse=True)

    def __unicode__(self):
        return self.name

class Match(SmartModel):
    season = models.ForeignKey(Season, related_name='matches')

    team1 = models.ForeignKey(Team, related_name='team1_matches')
    team2 = models.ForeignKey(Team, related_name='team2_matches')

    handicap1 = models.IntegerField(default=0)
    handicap2 = models.IntegerField(default=0)

    points1 = models.DecimalField(max_digits=3, decimal_places=1, default=0)
    points2 = models.DecimalField(max_digits=3, decimal_places=1, default=0)

    date = models.DateField()

    def sum_for_players(self, players):
        player_sum = 0
        for player in players:
            player_avg = player.avg(self.date)
            print "%s: %f" % (player.name, player_avg)
            player_sum += player_avg

        return player_sum

    def calculate_handicaps(self):
        # calculate average score for team 1 players for this season before this date
        print "Team1----------"
        team1_sum = self.sum_for_players(Player.objects.filter(scores__match=self, scores__team=self.team1).distinct())
        print team1_sum
        print "Team2----------"
        team2_sum = self.sum_for_players(Player.objects.filter(scores__match=self, scores__team=self.team2).distinct())
        print team2_sum

        if team1_sum > team2_sum:
            self.handicap1 = 0
            self.handicap2 = round(team1_sum - team2_sum - .01)
        else:
            self.handicap1 = round(team2_sum - team1_sum - .01)
            self.handicap2 = 0

    def summary(self):
        rounds = []
        match = dict(rounds=rounds, score1=0, score2=0, points1=0, points2=0)

        # are we 16 games or 25?
        last_game = PlayerScore.objects.filter(match=self).order_by('-game').first()
        games = 25
        if last_game:
            games = last_game.game

        if games <= 9:
            num_rounds = 3
        elif games <= 16:
            num_rounds = 4
        elif games <= 24:
            num_rounds = 5
        else:
            raise Exception("unknown number of games")

        for game_num in range(1, games+1):
            if (game_num - 1) % num_rounds == 0:
                games = []
                round = dict(games=games, round=len(rounds)+1, score1=0, score2=0)
                rounds.append(round)

            p1_score = PlayerScore.objects.filter(match=self, team=self.team1, game=game_num).first()
            p2_score = PlayerScore.objects.filter(match=self, team=self.team2, game=game_num).first()

            if p1_score and p2_score:
                round['score1'] += p1_score.score
                round['score2'] += p2_score.score

                games.append(dict(game=game_num, round=(game_num-1) / num_rounds,
                                  player1=p1_score.player, score1=p1_score.score,
                                  player2=p2_score.player, score2=p2_score.score))

        # add in handicaps
        for round in rounds:
            round['hscore1'] = round['score1'] + self.handicap1
            round['hscore2'] = round['score2'] + self.handicap2
            match['score1'] += round['hscore1']
            match['score2'] += round['hscore2']

        # calculate wins
        for round in rounds:
            if round['hscore1'] > round['hscore2']:
                match['points1'] += 1
                round['win1'] = True

            elif round['hscore2'] > round['hscore1']:
                match['points2'] += 1
                round['win2'] = True

            else:
                match['points1'] += .5
                match['points2'] += .5
                round['win1'] = True
                round['win2'] = True

        if match['score1'] > match['score2']:
            match['points1'] += 1
            match['score1_win'] = True

        elif match['score2'] > match['score1']:
            match['points2'] += 1
            match['score2_win'] = True

        if match['points1'] > match['points2']:
            match['win1'] = True
        elif match['points2'] > match['points1']:
            match['win2'] = True

        match['handicap1'] = self.handicap1
        match['handicap2'] = self.handicap2

        return match

    def calculate_stats(self):
        # calculate handicaps
        self.calculate_handicaps()

        # calculate wins
        summary = self.summary()
        self.points1 = summary['points1']
        self.points2 = summary['points2']
        self.save()

    def __unicode__(self):
        return "%s vs %s on %s" % (str(self.team1), str(self.team2), str(self.date))

    class Meta:
        verbose_name_plural = "Matches"

class PlayerScore(SmartModel):
    match = models.ForeignKey(Match, related_name='scores')
    game = models.IntegerField()

    team = models.ForeignKey(Team, related_name='scores')
    player = models.ForeignKey(Player, related_name='scores')
    score = models.IntegerField()

    def __unicode__(self):
        return "%s: %d" % (self.player, self.score)
<|endoftext|>