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 solve(Circuit, Nnodes):
    U = [ sympy.symbols("U"+str(k)) for k in range(0,Nnodes) ]
    Y = [0]*Nnodes
    YR = [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
        else:
            YR[i] += U[i] - U[j] + value
            YR[j] += U[j] - U[i] - value
    eqns = []
    for i in range(Nnodes):
        if YR[i] != 0:
            eqns.append(YR[i])
        else:
            eqns.append(Y[i])
    eqns.append(U[0])          # set U0 = 0
    print("Equations:")
    for y in eqns:
        print("      0 =", y)
    sol = sympy.solve(eqns, U)
    print("\nSolution:")
    for i,v in enumerate(U):
        print("   {:>4} = {}".format(str(v), str(sol[v])))

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