import sympy as sympy

# Experiment-5 circuit
Example = [
    [0,0,5,1000],
    [0,1,2,1000],
    [0,1,4,1000],
    [0,1,5,1000],
    [0,2,3,1000],
    [0,2,4,1000],
    [0,4,5,1000],
    [1,0,1,5],
    [1,0,3,3.3]
]

# simple 3.3V into a 500R resistor
Example2 = [
    [0, 1, 0, 500],
    [1, 0, 1, 3.3]
]

# voltage divider - 5V into a 100/200 divider
Example3 = [
    [1, 0, 1, 5],
    [0, 1, 2, 100],
    [0, 2, 0, 200]
]

def merge(nodes, i, j):
    ith = [ k for k,v in enumerate(nodes) if i in v ][0]
    jth = [ k for k,v in enumerate(nodes) if j in v ][0]
    if ith != jth:
        nodes[ith] = nodes[ith] + nodes[jth]
        del nodes[jth]

def solve(Circuit, Nnodes):
    U = [ sympy.symbols("U"+str(k)) for k in range(0,Nnodes) ]
    Y = [0]*Nnodes
    supernodes = [ [i] for i in range(0,Nnodes) ]

    for type, i, j, value in Circuit:
        value = float(value)
        if type == 0:
            g = 1/value
            Y[i] += U[i]*g - U[j]*g
            Y[j] += U[j]*g - U[i]*g
        elif type == 1:
            Y.append(U[j]-U[i]-value)
            merge(supernodes, i, j)
    print("Supernodes:", [ s for s in supernodes if len(s) > 1 ])
    for s in supernodes:
        if len(s) > 1:
            i = s[0]
            for j in s[1:]:
                Y[i] += Y[j]
                Y[j] = 0       # remove equations for other nodes
    Y.append(U[0])      # set U0 = 0
    print("Equations:")
    for y in Y:
        print("      0 =", y)
    sols = sympy.solve(Y, U)
    if sols:
        sol = sols
        print("\nSolution:")
        for i,v in enumerate(U):
            print("   {:>4} = {}".format(str(v), sol[v]))
    else:
        print("\nNo Solution")

solve(Example, 6)
# solve(Example2, 2)
# solve(Example3, 3)