#!/usr/bin/python
"""Scheduler component: compute (near-)optimal tuner allocations
"""

import logging
import logging.config

import calliope.db
from calliope.recording import Recording
from calliope.tuner import Tuner
from calliope.channel import Channel

log = logging.getLogger("ocarina")
#logging.basicConfig(level=logging.DEBUG)
logging.config.fileConfig("/etc/calliope/logging.conf")

def partition(lst, fn):
	"""Return two lists, of elements of "set" that fn maps to
	True, and elements thatfn maps to False
	"""
	rv = [[], []]
	for i in lst:
		if fn(i):
			rv[0].append(i)
		else:
			rv[1].append(i)
	return rv

class Problem(object):
	def __init__(self, nodes=None, active=None):
		if nodes is None:
			nodes = set()
		if active is None:
			active = []

		self.nodes = nodes
		self.active = active
		self.bestvalue = 0

	def add_recording(self, rec):
		# Clear out the active recordings that finish before this new
		# recrding starts
		while self.active and rec.start > self.active[0].recs[0].end:
			self.active.pop(0)

		# Now add the new recording
		# Initially, every node is a single recording
		cur_node = Node(rec)
		self.nodes.add(cur_node)

		# When we add something to the active list, we draw an edge to
		# everything else on the active list
		for n in self.active:
			cur_node.add_edge(n)

		self.active.append(cur_node)
		self.active.sort(lambda a, b: cmp(a.recs[0].end, b.recs[0].end))

	def completed(self):
		self.active = None

	def split(self):
		subproblems = []
		nodes = self.nodes.copy()
		while nodes:
			Q = Problem()
			# Take an arbitrary node off the list, and follow its
			# arcs recursively
			n = nodes.pop()
			Q.nodes.add(n)
			self._follow_edges(Q, nodes, n)
			subproblems.append(Q)
		return subproblems

	def copy(self):
		"""Make a shallowish copy of this object. Generate new nodes,
		but keep the recordings uncopied.
		"""
		rv = Problem()
		rv.nodes = set()
		nmap = {}
		# Duplicate all the nodes, and keep track of which new node
		# resulted from which old one.
		for n in self.nodes:
			newn = Node(node=n)
			nmap[id(n)] = newn
			rv.nodes.add(newn)
		# Go back through them all and add the edges in, now that we
		# have them.
		for n in self.nodes:
			for e in n.edges:
				nmap[id(n)].add_edge(nmap[id(e)])
				
		return rv

	def _follow_edges(self, Q, nodes, n):
		"""From node n, find any nodes in the set nodes connected to
		n, and add them to the problem P. Remove the nodes from the
		set nodes, and continue recursively.
		"""
		for e in n.edges:
			if e in nodes:
				nodes.remove(e)
				Q.nodes.add(e)
				self._follow_edges(Q, nodes, e)

	def allocate(self):
		"""Generate a feasible allocation for this problem, by
		attempting to resolve collisions.
		"""
		self.bestvalue = len(self.nodes)

		nodes = list(self.nodes)
		nodes.sort(lambda a, b: cmp(a.d(), b.d()))

		rv = self.colour_graph(nodes[:])
		if rv == 0:
			return 0
		
		# First, find all of the possible edges we could remove
		sacrifices = []
		for n in self.nodes:
			for e in n.edges:
				if n.recs[0].mergable(e.recs[0]):
					if (e, n) not in sacrifices:
						sacrifices.append((n, e))
		# Partition the set into merges across channels (preferred)
		# and within channels
		sacrifices = partition(sacrifices, lambda x: x[0].recs[0].overlap(x[1].recs[0]))
		sacrifices = sacrifices[0] + sacrifices[1]

		S = []
		for s in sacrifices:
			S.append(s)
			rv = self.colour_graph(nodes[:], S)
			if rv == 0:
				return 0

		return self.bestvalue

	def colour_graph(self, nodes, merges=[], missing=0):
		"""Generate a feasible colouring for this graph, or None.
		"""
		# No uncoloured nodes, we have a solution
		if not nodes:
			if missing < self.bestvalue:
				self.bestvalue = missing
				for n in self.nodes:
					n.bestcolour = n.colour
			return missing

		# If we're already below the optimum, skip the rest of this
		# search
		if missing >= self.bestvalue:
			#print "Solution already worse than current best: skipping", missing, self.bestvalue
			return missing

		maxn = nodes.pop(0)

		#print "Colouring node", maxn
		#print "Merge candidate pairs", merges

		# Work out our feasibles
		feasibles = maxn.colours.copy()
		# Remove colours surrounding us, where we could clash
		for nb in maxn.edges:
			# No colour on the neighbour, so skip to the next
			if nb.colour is None:
				continue
			# If the neighbour is a possible merge, leave the colour
			# in and skip to the next
			if (maxn, nb) in merges or (nb, maxn) in merges:
				continue
			# Otherwise, we have no chance of merging, so we can't
			# use this colour
			feasibles.discard(nb.colour)

		#print "My feasibles are", feasibles

		# Go through all the feasible colours of this node (its own
		# possibles, less the neighbouring colours)
		for c in feasibles:
			#print "For node", maxn
			#print "Trying colour", c
			maxn.colour = c
			rv = self.colour_graph(nodes[:], merges, missing)
			if rv == 0:
				#print "WIN!"
				return rv

		# If we don't find anything, clear this node, and try the rest
		# of the graph
		#print "FAIL!"
		maxn.colour = None
		rv = self.colour_graph(nodes[:], merges, missing+1)
		return rv

	def take_best(self):
		# Get the best-allocated colours
		for n in self.nodes:
			n.colour = n.bestcolour

	def set_rgroups(self):
		# Compute the rgroups: take the nodes in arbitrary order.
		# Check the node's neighbours: if all are different colours to
		# this node, go to the next one.
		# If some or more are of the same colour, look at their rgroups:
		# All zero: pick the next rgroup number and add me to it (keep
		# a list of nodes for each number)
		# Otherwise, set me to max(e_i.rgroup), and move all the
		# rgroups of smaller size into that one

		log.debug("rgroups processing starts")
		rgroups = [ set(), ]
		for n in self.nodes:
			log.debug("For node %s", str(n))
			same = set() # Set of neighbours of the same colour
			this_groups = set() # Set of rgroups amongst those neighbours
			max_rg = 0 # Largest rgroup we've seen in the neighbours
			for e in n.edges:
				if n.colour is not e.colour:
					# Differing colours: ignore this neighbour
					continue
				# Same colours: accumulate data
				same.add(e)
				this_groups.add(e.rgroup)
				max_rg = max(max_rg, e.rgroup)

			log.debug("  Same-colour neighbours: %s", str(same))
			log.debug("  Rgroups: %s", str(this_groups))
			log.debug("  max_rg: %d", max_rg)
				
			if len(this_groups) == 0:
				# No neighbours of the same colour
				log.debug("No neighbours... next.")
				continue
			if len(this_groups) == 1 and iter(this_groups).next() == 0:
				# Only one rgroup seen nearby: use that
				log.debug("Only one rgroup: setting to %d", len(rgroups))
				n.rgroup = len(rgroups)
				rgroups.append(set((n,)))
			else:
				# Several rgroups seen nearby: merge them all into the
				# maximum one
				log.debug("Several rgroups")
				n.rgroup = max_rg
				for rg in this_groups:
					if rg == max_rg:
						continue
					for p in rgroups[rg]:
						p.rgroup = max_rg
					rgroups[max_rg] |= rgroups[rg]
					rgroups[rg] = None

				
	def write_back(self):
		for n in self.nodes:
			n.recs[0].tuner = n.colour
			n.recs[0].rgroup = n.rgroup
			if n.recs[0].state == "Waiting":
				n.recs[0].state = "Allocated"

	def __repr__(self):
		return "Problem(nodes=set(\n  %s\n), active=%s)" \
			   % (",\n  ".join([str(n) for n in self.nodes]), str(self.active))

class Node(object):
	def __init__(self, recording=None, node=None):
		self.edges = set()
		if node is None:
			self.recs = [ recording, ]
			self._set_colours()
			self.colour = None
			self.bestcolour = None
			self.rgroup = 0
		else:
			self.recs = node.recs[:]
			self.colours = node.colours.copy()
			self.colour = node.colour
			self.bestcolour = node.bestcolour
			self.rgroup = node.rgroup

	def add_edge(self, node):
		self.edges.add(node)
		node.edges.add(self)

	def remove_edge(self, node):
		self.edges.remove(node)
		node.edges.remove(self)

	def d(self):
		return len(self.edges)

	def _set_colours(self):
		"""Generate the set of allowable colours for this node.
		"""
		self.colours = None
		for rec in self.recs:
			channel = rec.channel
			tuners = set()
			for ct in channel.tuners:
				tuners.add(ct.tuner)
			if self.colours is None:
				self.colours = tuners
			else:
				self.colours &= tuners # Intersection

	def __repr__(self):
		if self.colour is None:
			colour = " "
		else:
			colour = str(self.colour.tid)
		if self.bestcolour is None:
			bestcolour = " "
		else:
			bestcolour = str(self.bestcolour.tid)

		return "Node (%s.%s) r%d @ %x -> [%s] [%s = %d-%d]" \
			   % (colour, bestcolour, self.rgroup, id(self),
				  ", ".join(["%x" % id(e) for e in self.edges]),
				  self.recs[0].short_title(),
				  self.recs[0].start,
				  self.recs[0].end)



# Grab all recordings from the database in start order
recordings = list(Recording.select())
recordings.sort(lambda a, b: cmp(a.start, b.start))

# Go through the events in order and devise an interval graph. We add
# recordings to the active list in start order, and keep the active
# list in finish order. We then take start/end events in sequence,
# moving recordings onto and off the active list as appropriate.
P = Problem()

for rec in recordings:
	P.add_recording(rec)
P.completed()

# We can now run an allocator algorithm. First, break into distinct
# sub-problems.
problems = P.split()

# Then solve each sub-problem
for Q in problems:
	log.debug("Solving problem:")
	log.debug(Q)
	missing = Q.allocate()
	if missing > 0:
		Q.take_best()
	Q.set_rgroups()
	Q.write_back()
	log.debug(Q)