# gta extension.py
import gta

class Unifier:
  def __init__(self, sym1, sym2):
    if not isinstance(sym1, gta.Symbol) or not isinstance(sym2, gta.Symbol):
      raise RuntimeError()
    symbol = None
    if sym1.symType == gta.SymbolType.SkolemConst:
      if sym2.symType == gta.SymbolType.FreeVar:
        symbol = sym1
      else: pass
    elif sym1.symType == gta.SymbolType.FreeVar:
      if (sym2.symType == gta.SymbolType.SkolemConst or
          sym2.symType == gta.SymbolType.FreeVar):
        symbol = sym2
    self.sym = symbol
  def substitute(self, node):
    if self.sym is None: return None
    if node.varSym.symType == gta.SymbolType.FreeVar:
      return gta.Node(node.funSym, self.sym)
    else: return node
  def __str__(self):
    return str(self.sym)


# construct new literal (a.k.a arc) from 2 literals
def combine2Edges(cls, edge1, edge2):
  if not isinstance(edge1, gta.Literal) or not isinstance(edge2, gta.Literal):
    raise RuntimeError()
  if edge1.des != edge2.src:
    raise RuntimeError()
  combined = cls(edge1.src, edge2.des, edge1.weight + edge2.weight)
  combined.history = (edge1, edge2)
  return combined
setattr(gta.Literal, 'combine2Edges', classmethod(combine2Edges))


def combine2EdgesWithUnifier(cls, edge1, edge2, uni):
  if not isinstance(edge1, gta.Literal) or not isinstance(edge2, gta.Literal):
    raise RuntimeError()
  combined = None
  src = uni.substitute(edge1.src)
  des = uni.substitute(edge2.des)
  if src is not None and des is not None:
    combined = cls(src, des, edge1.weight + edge2.weight)
    combined.history = (edge1, edge2)
  return combined
setattr(gta.Literal, 'combine2EdgesWithUnifier',
        classmethod(combine2EdgesWithUnifier))


# check if a literal is a self-arc
def isSelfArc(literal):
  return literal.src == literal.des
setattr(gta.Literal, 'isSelfArc', isSelfArc)


# trace arcs which were combined into this arc/path
def tracePath(literal):
  edges = []
  if literal.history is not None:
    edges = edges + literal.history[0].tracePath()
    edges = edges + literal.history[1].tracePath()
  else:
    edges.append(literal)
  return edges
setattr(gta.Literal, 'tracePath', tracePath)


def validateAndGetCycle(literal):
  cycle = literal.tracePath()
  nodes = []
  valid = True
  for arc in cycle:
    incluster = False
    for node in nodes:
      if node.funSym == arc.src.funSym:
        incluster = True
        break;
    if incluster:
      valid = False
      break
    nodes.append(arc.src)
  if valid: return cycle
  else: return None
setattr(gta.Literal, 'validateAndGetCycle', validateAndGetCycle)

# remove a node, preserving cycles
# return cycles found so far
def cycleInvariantRemove(graph, node):
  if not isinstance(node, gta.Node):
    raise RuntimeException()
  incommingArcs = []
  outgoingArcs = []
  cyclesAsSelfArcs = []
  for literal in graph.literals:
    if literal.isSelfArc():
      cyclesAsSelfArcs.append(literal)
      continue
    if literal.des == node:
      incommingArcs.append(literal)
    if literal.src == node:
      outgoingArcs.append(literal)
  for left in incommingArcs:
    for right in outgoingArcs:
      newArc = gta.Literal.combine2Edges(left, right)
      graph.addLiteral(newArc)
  for left in incommingArcs:
    graph.removeLiteralKeepingNode(left)
  for right in outgoingArcs:
    graph.removeLiteralKeepingNode(right)
  for arc in cyclesAsSelfArcs:
    graph.removeLiteralKeepingNode(arc)
  graph.removeIsolatedNode(node)
  return cyclesAsSelfArcs
setattr(gta.Graph, 'cycleInvariantRemove', cycleInvariantRemove)


# remove a node, preserving cycles, try unifiers to generate new arcs
# return cycles found so far
def cycleInvariantRemoveWithUnification(graph, node, cluster):
  if not isinstance(node, gta.Node):
    raise RuntimeError()
  incommingArcs = []
  outgoingArcs = []
  cyclesAsSelfArcs = []

  for literal in graph.literals:
    if literal.des == node or literal.src == node:
      if literal.isSelfArc():
        cyclesAsSelfArcs.append(literal)
      elif literal.des == node:
        incommingArcs.append(literal)
      else:
        outgoingArcs.append(literal)

  for vertex in cluster.nodes:
    gta.Debug.trace('in-cluster vertex:', vertex)
    if vertex == node:
      for left in incommingArcs:
        for right in outgoingArcs:
          newArc = gta.Literal.combine2Edges(left, right)
          gta.Debug.trace('new arc added for', node, newArc, 'by:', left, right)
          if newArc is not None: graph.addLiteral(newArc)
    else:
      gta.Debug.trace('process node-vertex', node, vertex)
      uni = Unifier(node.varSym, vertex.varSym)
      extraIncommingArcs = []
      extraOutgoingArcs = []
      if uni.sym is not None:
        for literal in graph.literals:
          if literal.des == vertex or literal.src == vertex:
            if literal.isSelfArc():
              cyclesAsSelfArcs.append(literal)
            elif literal.des == vertex:
              extraIncommingArcs.append(literal)
            else:
              extraOutgoingArcs.append(literal)
        for left in extraIncommingArcs:
          for right in outgoingArcs:
            newArc = gta.Literal.combine2EdgesWithUnifier(left, right, uni)
            gta.Debug.trace('new arc added', node, vertex, uni, ':', newArc)
            if newArc is not None: graph.addLiteral(newArc)
        for left in incommingArcs:
          for right in extraOutgoingArcs:
            newArc = gta.Literal.combine2EdgesWithUnifier(left, right, uni)
            gta.Debug.trace('new arc added', node, '-', vertex, uni, ':', newArc)
            if newArc is not None: graph.addLiteral(newArc)
  return cyclesAsSelfArcs
setattr(gta.Graph, 'cycleInvariantRemoveWithUnification',
        cycleInvariantRemoveWithUnification)


def getIncidentArcs(graph, node):
  incident = []
  for literal in graph.literals:
    if literal.src == node or literal.des == node:
      incident.append(literal)
  return incident
setattr(gta.Graph, 'getIncidentArcs', getIncidentArcs)


# keep removing nodes and discover cycles along the process
def shrink(graph):
  cycles = []
  for funSym in sorted(graph.clusters.keys(), key = lambda x: x.name):
    cluster = graph.clusters[funSym]
    numSrinkedNodes = len(cluster.nodes) - 0
    while (numSrinkedNodes > 0 and len(cluster.nodes) > 0):
      prev = len(cluster.nodes)
      node = cluster.nodes[-1]
      incidentArcs = graph.getIncidentArcs(node)
      gta.Debug.trace('ready to process', node)
      cyclesAsSelfArcs = (
        graph.cycleInvariantRemoveWithUnification(node, cluster))
      for arc in incidentArcs:
        graph.removeLiteralKeepingNode(arc)
      for arc in cyclesAsSelfArcs:
        circuit = arc.validateAndGetCycle()
        if circuit is not None:
          cycle = gta.Cycle(*circuit)
          cycles.append(cycle)
          gta.Debug.trace('cycle found:', cycle)
        try:
          graph.removeLiteralKeepingNode(arc)
        except ValueError as e: pass # already removed, ok
      gta.Debug.trace('to remove ', node)
      graph.removeIsolatedNode(node)
      numSrinkedNodes = prev - len(cluster.nodes)
  return cycles
setattr(gta.Graph, 'shrink', shrink)
