# Solution for Orthogonal Collocation on Finite Elements
import numpy as np
from gekko import GEKKO

tf = 10
u  = 0.5

x10 = -0.5
x20 =  0.0

m = GEKKO()
x11,x12,x13 = m.Array(m.Var,3)
x21,x22,x23 = m.Array(m.Var,3)
dx11,dx12,dx13 = m.Array(m.Var,3)
dx21,dx22,dx23 = m.Array(m.Var,3)

m.Equations([tf * (0.436*dx11 -0.281*dx12 +0.121*dx13) == x11 - x10, \
             tf * (0.614*dx11 +0.064*dx12 +0.046*dx13) == x12 - x10, \
             tf * (0.603*dx11 +0.230*dx12 +0.167*dx13) == x13 - x10, \
             tf * (0.436*dx21 -0.281*dx22 +0.121*dx23) == x21 - x20, \
             tf * (0.614*dx21 +0.064*dx22 +0.046*dx23) == x22 - x20, \
             tf * (0.603*dx21 +0.230*dx22 +0.167*dx23) == x23 - x20, \
             dx11 == u, \
             dx12 == u, \
             dx13 == u, \
             dx21 == x11**2-u, \
             dx22 == x12**2-u, \
             dx23 == x13**2-u])

m.solve()

print('States with Orthogonal Collocation')
print('x1 = ', x10, x11.value[0], x12.value[0], x13.value[0])
print('x2 = ', x20, x21.value[0], x22.value[0], x23.value[0])
print(' ')
print('Derivatives with Orthogonal Collocation')
print('dx1/dt = ', dx11.value[0], dx12.value[0], dx13.value[0])
print('dx2/dt = ', dx21.value[0], dx22.value[0], dx23.value[0])