from optimake import *

prob = multi_stage_problem('traj_tracking', 20)

xweight, yweight, vweight = prob.parameters(['xweight', 'yweight', 'vweight'], stage_dependent=False)
x0, y0, phi0, v0 = prob.parameters(['x0', 'y0', 'phi0', 'v0'], stage_dependent=False)
vref, yref = prob.parameters(['vref', 'yref'], stage_dependent=False)
xref = prob.parameter('xref', stage_dependent=True)

x = prob.variable('x')
y = prob.variable('y')
phi = prob.variable('phi')
delta = prob.variable('delta', hard_lowerbound=-0.2, hard_upperbound=0.2)
v = prob.variable('v')
a = prob.variable('a', hard_lowerbound=-1.0, hard_upperbound=1.0)

obj = general_objective(xweight * (x - xref)**2 + yweight * (y - yref)**2 + vweight * (v - vref)**2)
prob.objective(obj)

""" ode """
length = 1.0
ode = differential_equation(
    state=[x, y, phi, v], 
    state_dot=[v * cos(phi), v * sin(phi), v * tan(delta) / length, a], 
    stepsize=0.2, 
    discretization_method='erk4')
prob.equality(ode)

seq = general_equality([x - x0, y - y0, phi - phi0, v - v0])
prob.start_equality(seq)


option = codegen_option()
option.platform = 'windows-x86_64-mingw'
codegen = code_generator()
codegen.codegen(prob, option)