# Matrix solution for Orthogonal Collocation on Finite Elements
import numpy as np

# Solve A x = b
# Matrix A
A = np.array([[ 1,-1, 0, 0, 5, 0, 0, 0],\
              [-1, 1, 0, 0, 0, 0, 3, 0], \
              [ 0, 0, 1,-1, 0, 5, 0, 0], \
              [ 0, 0,-1, 1, 0, 0, 0, 3], \
              [-1, 0, 0, 0, 0.75,-0.25, 0, 0], \
              [ 0, 0,-1, 0, 1, 0, 0, 0], \
              [ 0,-1, 0, 0, 0, 0, 0.75, -0.25], \
              [ 0, 0, 0,-1, 0, 0, 1, 0]])
# Column vector b
b = np.array([2,0,2,0,0,0,0,0])
# Solve A x = b as x = A^-1 * b
sol = np.linalg.solve(A,b)

print('States with Orthogonal Collocation')
print(['x11 = ' + str(sol[0])])
print(['x21 = ' + str(sol[1])])
print(['x12 = ' + str(sol[2])])
print(['x22 = ' + str(sol[3])])
print(' ')
print('Derivatives with Orthogonal Collocation')
print(['d(x11)/dt = ' + str(sol[4])])
print(['d(x21)/dt = ' + str(sol[5])])
print(['d(x12)/dt = ' + str(sol[6])])
print(['d(x22)/dt = ' + str(sol[7])])