MA0301/exam_template/python/Graph.py

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

# from run import printred

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):
    self.nodes = nodes
    self.edges = edges
  
  def __str__(self):
    print(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 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 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 processFileContent(raw, template):
  lines = raw.split('\n')
  lines.pop(1)
  outputType = lines.pop(0)

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

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

  if outputType == 'toGraph':
    content = graph.toLaTeX()
    return template.replace('%NODES', content[0]).replace('%EDGES', content[1])
  else:
    content = graph.toMatrix().toLaTeX()
    return template.replace('%CONTENT', content)



if __name__ == '__main__':
  matrix = Matrix([[False, False, True], [True, False, True], [True, False, False]])
  print(matrix)
  print(matrix.toGraph())
  print(matrix.toGraph().isUndirected())