MA0301/exam_template/python/Graph.py

from sys import argv
from pathlib import Path
from math import sin, cos, pi

from common import printc, printerr, replaceContent

class Matrix:
  """ Adjacency matrix which supports 0 and 1 """

  def __init__(self, matrix):
    self.matrix = matrix

  def __iter__(self):
    return iter(self.matrix)
  
  def __len__(self):
    return len(self.matrix)

  def __eq__(self, other):
    return self.matrix == other

  def __str__(self):
    return "\n".join(' '.join('\033[32m1\033[0m' if b else '\033[31m0\033[0m' for b in line) for line in self)

  def __getitem__(self, key):
    return self.matrix[key]

  def __setitem__(self, key, item):
    self.matrix[key] = item
  
  @classmethod
  def fromGraph(cls, graph):
    return graph.toMatrix()
  
  @classmethod
  def fromString(cls, string):
    matrix = []
    for line in string.split('\n'):
      matrix.append([True if char == '1' else False for char in line])

    if len(matrix) != len(matrix[0]):
      raise ValueError('Matrix not covering all points')
  
    return cls(matrix)

  def toGraph(self):
    nodes = [chr(i + 65) for i in range(len(self))]

    edges = []
    for i, line in enumerate(self):
      edges += [(nodes[i],nodes[j]) for j, b in enumerate(line) if b]

    return Graph(nodes, edges)

  def toLaTeX(self):
    return '\\begin{pmatrix}\n' \
      + '\n'.join('  ' + ' & '.join('1' if b else '0' for b in line) + ' \\\\' for line in self.matrix) \
      + '\n\\end{pmatrix}'


class Graph:
  def __init__(self, nodes, edges): # Nodes = str, Edges = (str,str)
    self.nodes = nodes
    self.edges = edges
  
  def __str__(self):
    printc('Is undirected: ' + str(self.isUndirected()))
    return \
      f'Nodes: {" ".join(self.nodes)}\n' \
    + f'Edges: {" ".join(x+y for x,y in (self.undirectedEdgeSet() if self.isUndirected() else self.edges))}'
  
  @classmethod
  def fromMatrix(cls, matrix):
    return matrix.toGraph()

  @classmethod
  def fromString(cls, string):
    splitData = string.split('\n\n')

    if splitData[0] == 'complete':
      data1 = 'undirected'
      nodes = [chr(i + 65) for i in range(int(splitData[1]))]
      data2 = ' '.join(nodes)
      data3 = '\n'.join([a+b for a in nodes for b in nodes if a != b])
    elif splitData[0] == 'matrix':
      return Graph.fromMatrix(Matrix.fromString(splitData[1]))
    else:
      data1, data2, data3 = splitData

    graphType = data1
    nodes = data2.split(' ')
    edges = [(x,y) for x,y in data3.split('\n')]

    if graphType == 'undirected':
      edges = [(a,b) for a in nodes for b in nodes if (a,b) in edges or (b,a) in edges]

    return cls(nodes, edges)

  def toMatrix(self):
    rows = []
    for node in self.nodes:
      rows.append([(node,node2) in self.edges for node2 in self.nodes])
    return Matrix(rows)

  def isUndirected(self):
    matrix = self.toMatrix()

    flipv = lambda matrix: list(reversed(matrix))
    rot90cc = lambda matrix: list(list(x) for x in zip(*reversed(matrix)))

    return matrix == rot90cc(flipv(matrix)) \
      and all(matrix[i][i] == 0 for i in range(len(matrix)))
  
  def undirectedEdgeSet(self):
    edges = self.edges
    if self.isUndirected():
      edges = sorted(list(set((x,y) if x < y else (y,x) for x,y in edges)))
    return edges
  
  # def latexifyEulerPath(self):
  #   degrees = [sum(line) for line in self.toMatrix()]

  # def latexifyEulerCircuit():
  #   pass

  def getKruskalsSpanningTree(self, weigths): # weights = dict[str, int]
    edges = []
    connected_subgraphs = [set(v) for v in self.nodes]

    if not all(n in ''.join(x+y for x,y in self.edges) for n in self.nodes):
      printerr('Not all nodes are connected. There is no spanning tree to be made')
      return

    def find_subgraph_containing(v):
      return [x for x in connected_subgraphs if v in x][0]

    def merge_subgraphs_that_contains(u, v):
      subgraph_v = find_subgraph_containing(v)
      new_connected_subgraphs = [x for x in connected_subgraphs if v not in x]
      new_connected_subgraphs = [subgraph_v.union(x) if u in x else x for x in new_connected_subgraphs]
      return new_connected_subgraphs

    sorted_edges = [ e for e in sorted(self.edges, key=lambda e: weigths[e[0] + e[1]]) ]

    for u,v in sorted_edges:
      if find_subgraph_containing(u) != find_subgraph_containing(v):
        edges += [(u, v)]
        connected_subgraphs = merge_subgraphs_that_contains(u, v)
    
    edges += [(y,x) for x,y in edges]

    return Graph(self.nodes, edges)
  
  def toLaTeX(self):
    zippedNodes = zip(self.nodes, generateNodeCoords(len(self.nodes)))
    nodeString = '\n'.join(f'\\node ({name}) at ({x},{y}) {{${name}$}};' for name,(x,y) in zippedNodes)

    if self.isUndirected():
      edgeString = '\n'.join(f'\\draw ({x}) -- ({y});' for x,y in self.undirectedEdgeSet())
    else:
      edgeString = '\n'.join(
             f'\\draw [-{{Latex[scale=1]}}, bend left=8] ({x}) to ({y});' if y != x and (y,x) in self.edges 
        else f'\\draw [-{{Latex[scale=1]}}] ({x}) to [{generateLoopInOutAngle(x, self.nodes)},looseness=8] ({y});' if x == y
        else f'\\draw [-{{Latex[scale=1]}}] ({x}) to ({y});'
        for x,y in self.edges
      )

    return (nodeString, edgeString)

  def toLaTeXWithWeights(self, weights):
    zippedNodes = zip(self.nodes, generateNodeCoords(len(self.nodes)))
    nodeString = '\n'.join(f'\\node ({name}) at ({x},{y}) {{${name}$}};' for name,(x,y) in zippedNodes)

    if self.isUndirected():
      edgeString = '\n'.join(f'\\draw ({x}) -- ({y}) node [midway, above, sloped] (Tx{x+y}) {{{weights[x+y]}}};' for x,y in self.undirectedEdgeSet())

    return (nodeString, edgeString)


def generateLoopInOutAngle(node, nodes):
  baseAngle = 360 / len(nodes)
  nodeNum = [i for i,n in enumerate(nodes) if n == node][0]
  angle = nodeNum * baseAngle + 90
  return f'out={angle + 15},in={angle - 15}'


def generateNodeCoords(n):
  vectorLength = n / 2
  degreeToTurn = (2 * pi) / n

  
  nodeCoords = [(0, vectorLength)]
  for node in range(n):
    prev_x = nodeCoords[-1][0]
    prev_y = nodeCoords[-1][1]

    nodeCoords.append((
      round(cos(degreeToTurn) * prev_x - sin(degreeToTurn) * prev_y, 5),
      round(sin(degreeToTurn) * prev_x + cos(degreeToTurn) * prev_y, 5)
    ))

  return nodeCoords

def extractWeights(string):
  weights = dict()
  lines = string.split('\n')
  for i in range(4, len(lines)):
    edge, weight = lines[i].split(' ')
    weights[edge] = int(weight)
    weights[edge[1] + edge[0]] = int(weight)
    lines[i] = edge
  return ('\n'.join(lines), weights)

def processFileContent(raw):
  lines = raw.split('\n')
  lines.pop(1)
  outputType = lines.pop(0)

  if lines[0] == 'kruskals':
    lines[0] = 'undirected'
    string, weights = extractWeights('\n'.join(lines))
    graph1 = Graph.fromString(string)
    graph2 = graph1.getKruskalsSpanningTree(weights)
    return replaceContent(
      (graph1.toLaTeXWithWeights(weights), graph2.toLaTeXWithWeights(weights)),
      'Kruskals',
      replaceF = lambda temp, cont: 
        temp.replace('%NODES1', cont[0][0])
            .replace('%EDGES1', cont[0][1])
            .replace('%NODES2', cont[1][0])
            .replace('%EDGES2', cont[1][1])
    )


  graph = Graph.fromString('\n'.join(lines))

  print(graph)
  print(graph.toMatrix())

  if outputType == 'toGraph':
    content = graph.toLaTeX()
    return replaceContent(
      content,
      'Graph',
      replaceF = lambda temp, cont: temp.replace('%NODES', cont[0]).replace('%EDGES', cont[1])
    )
  else:
    content = graph.toMatrix().toLaTeX()
    return replaceContent(content, 'Matrix')

if __name__ == '__main__':
  g = Graph.fromString('complete\n\n6')
  print(g)

  import random

  d = dict()
  for e in g.edges:
    d[str(e)] = random.randint(0, 10)

  print(g.getKruskalsSpanningTree(d))