from fglib import graphs, nodes, inference, rv

# Create factor graph
fg = graphs.FactorGraph()

# Create variable nodes
x1 = nodes.VNode("x1", rv.Gaussian)  
x2 = nodes.VNode("x2", rv.Gaussian)
x3 = nodes.VNode("x3", rv.Gaussian)
x4 = nodes.VNode("x4", rv.Gaussian)

# Create factor nodes (with joint distributions)
fa = nodes.FNode("fa", rv.Gaussian([[1], [1]], [[1, 0], [0, 1]], x1, x2))
fb = nodes.FNode("fb", rv.Gaussian([[1], [1]], [[1, 0], [0, 1]], x2, x3))
fc = nodes.FNode("fc", rv.Gaussian([[1], [1]], [[1, 0], [0, 1]], x2, x4))

# Add nodes to factor graph
fg.set_nodes([x1, x2, x3, x4])
fg.set_nodes([fa, fb, fc])

# Add edges to factor graph
fg.set_edge(x1, fa)
fg.set_edge(fa, x2)
fg.set_edge(x2, fb)
fg.set_edge(fb, x3)
fg.set_edge(x2, fc)
fg.set_edge(fc, x4)

# Perform sum-product algorithm on factor graph
# and request belief of variable node x4
belief = inference.sum_product(fg, x4)

# Print belief of variables
print("Belief of variable node x4:")
print(belief)