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, Rval=0.001):
    U = [ sympy.symbols("U"+str(k)) for k in range(0,Nnodes) ]
    R = sympy.symbols('R')
    Y = [0]*Nnodes

    for type, i, j, value in Circuit:
        value = float(value)
        if type == 0:
            g = 1/value
        else:
            g = 1/R
        Y[i] += U[i]*g - U[j]*g
        Y[j] += U[j]*g - U[i]*g
        if type == 1:
            Y[i] += value/R
            Y[j] -= value/R
    Y.append(U[0])      # set U0 = 0
    Y.append(R-Rval)    # set R = Rval
    vars = U.copy()
    vars.append(R)
    print("Equations:")
    for y in Y:
        print("      0 =", y)
    sols = sympy.solve(Y, vars)[0]
    print("\nSolution:")
    for i,v in enumerate(vars):
        print("   {:>4} = {}".format(str(v), str(sols[i])))

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