#!/usr/bin/env python3

# input data
############

# name of ground potential
gnd = 'GND'

# resistors (name, node1, node2, value in Ohms
rs = [
  ( 'R1', 'A', 'B'  , 150 ),
  ( 'R2', 'A', 'C'  , 100 ),
  ( 'R3', 'B', 'C'  , 220 ),
  ( 'R4', 'B', 'GND', 330 ),
  ( 'R5', 'C', 'GND', 470 ),
]

# currant sources (node1, node2, value in Amperes
cs = [
  ( 'GND', 'A', 1.2e-3 ),
  ( 'A'  , 'C', 2.5e-3 ),
]


# solver
########

from collections import defaultdict
from types import SimpleNamespace
import numpy as np

def error(msg):
  exit(f'Error: {msg}')

def node_factory():
  # create empty node
  return SimpleNamespace(idx=None, rs=[], cs=[])

def create_structure(rs, cs):
  # create empty structure as a dictionary of nodes
  nodes = defaultdict(node_factory)
  # add resistors to structure
  for name, nd1, nd2, r in rs:
    if r == 0:
      error(f'Resistor {name} has value of zero')
    nodes[nd1].rs.append((nd2, r))
    nodes[nd2].rs.append((nd1, r))
  # add current sources
  for nd1, nd2, c in cs:
    nodes[nd1].cs.append(-c)
    nodes[nd2].cs.append( c)
  # remove GND node from structure if existing or print error message
  if gnd in nodes:
    del nodes[gnd]
  else:
    error(f'Circuit has no GND reference (node name {repr(gnd)})')
  # enumerate remaining nodes starting with zero
  for i, key in enumerate(nodes):
    nodes[key].idx = i
  return nodes

def solve(nodes):
  # size of equation system
  n = len(nodes)
  # nxn matrix filled with zeros
  m = np.zeros((n, n))
  # n dimensional vector filled with zeros
  b = np.zeros(n)
  # create equation system
  for val in nodes.values():
    i = val.idx
    # fill matrix with sums of conductances
    for nd, r in val.rs:
      g = 1 / r
      m[i,i] += g
      if nd != gnd:
        j = nodes[nd].idx
        m[i,j] -= g
    # fill vector with sums of currents
    for c in val.cs:
      b[i] += c

  # compute voltage for each node as solutions of the equation system
  try:
    vs = np.linalg.solve(m, b)
  except np.linalg.LinAlgError as e:
    error(e)
  return vs

def show_results(nodes, vs):
  # show voltage for each node
  for nd, val in nodes.items():
    print('V_%-6s = %13.6e V' % (nd, vs[val.idx]))
  print()
  # compute and show current for each resistor
  for name, nd1, nd2, r in rs:
    u1 = 0 if nd1 == gnd else vs[nodes[nd1].idx]
    u2 = 0 if nd2 == gnd else vs[nodes[nd2].idx]
    c = (u1 - u2) / r
    print('I_%-6s = %13.6e A' % (name, c))

nodes = create_structure(rs, cs)
vs = solve(nodes)
show_results(nodes, vs)