from numpy import *
from scipy import sparse
from scipy import linalg


def path(n): # edge list of path on n vertices
  e = []
  for i in range(n-1):
    e.append([i, i+1])

  return e


def randGraph(n, m): # generate random graph on n verts with close to m edges
  e = []
  for i in range(m):
    edge = [random.randint(n), random.randint(n)]
    if (edge[0] != edge[1]):
      e.append(edge)

  return e


def edges2adjmat(e): # convert edge list to adjacency matrix
  ea = array(e)
  numverts = ea.max() + 1
  a = sparse.lil_matrix((numverts,numverts))

  for edge in e:
    a[edge[0].__int__(),edge[1].__int__()] = 1
    a[edge[1].__int__(),edge[0].__int__()] = 1

  return a

def adjmat2lap(a): # convert adjacency matrix to a laplacian
  numverts = a.shape[0]
  degrees = a*ones(numverts)
  degmat = sparse.lil_matrix((numverts,numverts))
  degmat.setdiag(degrees)
  lap = degmat - a

  return lap



def solvePotentials(e, verts, pots): 
  # given edge list e, compute the graph, then
  # set the potentials of verts to pots,
  # and return the resulting potential
  # note that verts are numbered from 0

  a = edges2adjmat(e)
  l = adjmat2lap(a)
  l = l.tocsc()

  n = l.shape[0]

  b = zeros([n,1]) + 0.0


  for i in range(verts.__len__()):
    v = verts[i]
    for j in range(n):
      l[v,j] = 0

    l[v,v] = 1
    b[v] = pots[i]

  ans = linalg.iterative.gmres(l,b)
  x = ans[0]

  return x