import sympy
from sympy.utilities.codegen import codegen

phi = sympy.Symbol("phi")       #angle of the pendulum
dphi = sympy.Symbol("dphi")     #anglevelocity of the pendulum
ddphi = sympy.Symbol("ddphi")   #angle acceleration of the pendulum 

x = sympy.Symbol("x")           #position of the wagon
dx = sympy.Symbol("dx")         #velocity of the wagon
ddx = sympy.Symbol("ddx")       #acceleration of the wagon

mB = sympy.Symbol("mB")         #mass of the pendulum
mS = sympy.Symbol("mS")         #mass of the wagon

fx = sympy.Symbol("fx")         #friction of the pendulum
fphi = sympy.Symbol("fphi")     #friction of the wagon

l = sympy.Symbol("l")           #length of the pendulum
g = sympy.Symbol("g")           #gravity constant

#position of the pendulum

bx = x + sympy.sin(phi) * l     
by = sympy.cos(phi) * l

def timederivative(o):
	return o.diff(phi) * dphi + o.diff(dphi) * ddphi + o.diff(x) * dx + o.diff(dx) * ddx 


kineticenergy = (timederivative(bx) ** 2 + timederivative(by) ** 2) * mB / 2 + dx ** 2 * mS / 2
potentialenergy = by * g * mB

lagrangian = kineticenergy - potentialenergy

F = sympy.Symbol("F")            #Force on the wagon


sol = sympy.solve([
	timederivative(lagrangian.diff(dphi)) - lagrangian.diff(phi) + dphi * fphi,
	timederivative(lagrangian.diff(dx)) - lagrangian.diff(x) - mS * F + dx * fx,
], [ddphi, ddx])


eqddphi = sol[ddphi]
eqddphi.simplify()

eqddx = sol[ddx]
eqddx.simplify()


cg = codegen(("dinvertedpendulum", [dphi, eqddphi, dx, eqddx]), language = "Octave",
	argument_sequence = [phi, dphi, x, dx, F, g, l, mB, mS, fphi, fx]
)

# Two function to make the generated code copy an pasteable into an xcos block

def cleanup(s, outnames) :
	res = []
	for l in s.split("\n"):
		if not l.strip() or not l.strip()[0] == "%":
			res.append(l.strip())

	rejoined = "\n".join(res)

	for i, on in zip(range(1, len(outnames) + 1), outnames):
		rejoined = rejoined.replace("out%i" % i, "out_%s" % on)

	return rejoined

def blockify(s, outnames, argument_sequence):
	prefix = ""
	for ia, a in enumerate(argument_sequence):
		prefix += "%s = u%i\n" % (a, ia + 1)

	s = s.split("\n")
	s = "\n".join(s[1:-2])

	s += "\n"

	for i, on in zip(range(1, len(outnames) + 1), outnames):
		s += "y%u = out_%s\n" % (i, on)
	return prefix  + s


from sympy import S


# Linearizing the system

matr = sympy.Matrix([
	[S(0),              S(1),               S(0),            S(0)],              #dphi
	[eqddphi.diff(phi), eqddphi.diff(dphi), eqddphi.diff(x), eqddphi.diff(dx)],  #ddphi
	[S(0),              S(0),               S(0),            S(1)],              #dx
	[eqddx.diff(phi),   eqddx.diff(dphi),   eqddx.diff(x),   eqddx.diff(dx)],    #ddx
])

Fmat = sympy.Matrix([S(0), eqddphi.diff(F), S(0), eqddx.diff(F)])

invm = matr.subs({x : 0, phi : 0, dx : 0, dphi : 0})
invF = Fmat.subs({x : 0, phi : 0, dx : 0, dphi : 0})

parameter = {fx : 0.001, fphi : 0.001, l : 0.2, mS :  0.1, mB : 0.1, g : 9.81}


import numpy

import scipy
import scipy.linalg
import scipy.signal

A = numpy.array(invm.subs(parameter), dtype = numpy.float)
B = numpy.array(invF.subs(parameter), dtype = numpy.float)


#placing the poles
km = scipy.signal.place_poles(A, B, [-2, -2.1, -2.2, -2.3])


outnames = ["dphi", "ddphi", "dx", "ddx"]

print("Equations of motion for the pedulum on a wagon:")
print()
print(blockify(cleanup(cg[0][1], outnames), outnames, [phi, dphi, x, dx, F, g, l, mB, mS, fphi, fx]))

print()

print("Cotroll parameter:", km.gain_matrix[0])