def setupOcp(dae,conf,publisher,nk=50,nicp=1,deg=4):
    ocp = Coll(dae, nk=nk,nicp=nicp,deg=deg)
    
    # constrain invariants
    def invariantErrs():
        dcm = C.horzcat( [ C.veccat([dae.x('e11'), dae.x('e21'), dae.x('e31')])
                         , C.veccat([dae.x('e12'), dae.x('e22'), dae.x('e32')])
                         , C.veccat([dae.x('e13'), dae.x('e23'), dae.x('e33')])
                         ] ).trans()
        err = C.mul(dcm.trans(), dcm)
        dcmErr = C.veccat([ err[0,0]-1, err[1,1]-1, err[2,2]-1, err[0,1], err[0,2], err[1,2] ])
        f = C.SXFunction( [dae.xVec(),dae.uVec(),dae.pVec()]
                        , [dae.output('c'),dae.output('cdot'),dcmErr]
                        )
        f.setOption('name','invariant errors')
        f.init()
        return f

    [c0,cdot0,dcmError0] = invariantErrs().call([ocp.xVec(0),ocp.uVec(0),ocp.pVec()])
    ocp.constrain(c0,'==',0)
    ocp.constrain(cdot0,'==',0)
    ocp.constrain(dcmError0,'==',0)

    # constrain line angle
    for k in range(0,nk+1):
        ocp.constrain(kiteutils.getCosLineAngle(ocp,k),'>=',C.cos(55*pi/180))

    # constrain airspeed
    def constrainAirspeedAlphaBeta():
        f = C.SXFunction( [dae.xVec(),dae.uVec(),dae.pVec()]
                        , [dae.output('airspeed'),dae.output('alpha(deg)'),dae.output('beta(deg)')]
                        )
        f.setOption('name','airspeed/alpha/beta')
        f.init()

        for k in range(0,nk):
            [airspeed,alphaDeg,betaDeg] = f.call([ocp.xVec(k),ocp.uVec(k),ocp.pVec()])
            ocp.constrainBnds(airspeed,(10,50))
            ocp.constrainBnds(alphaDeg,(-5,10))
            ocp.constrainBnds(betaDeg,(-10,10))
    constrainAirspeedAlphaBeta()

    # make it periodic
    for name in [ #"x"   # state 0
                  "y"   # state 1
                , "z"   # state 2
#                , "e11" # state 3
#                , "e12" # state 4
#                , "e13" # state 5
#                , "e21" # state 6
#                , "e22" # state 7
#                , "e23" # state 8
#                , "e31" # state 9
#                , "e32" # state 10
#                , "e33" # state 11
#                , "dx"  # state 12
                , "dy"  # state 13
                , "dz"  # state 14
                , "w1"  # state 15
                , "w2"  # state 16
                , "w3"  # state 17
                , "r" # state 20
                , "dr" # state 21
                ]:
        ocp.constrain(ocp.lookup(name,timestep=0),'==',ocp.lookup(name,timestep=-1))

    # periodic attitude
    kiteutils.periodicDcm(ocp)

    # bounds
    ocp.bound('aileron',(-0.04,0.04))
    ocp.bound('elevator',(-0.1,0.1))

    ocp.bound('x',(-2,200))
    ocp.bound('y',(-200,200))
    ocp.bound('z',(0.5,200))
    ocp.bound('r',(1,30))
    ocp.bound('dr',(-10,10))
    ocp.bound('ddr',(-2.5,2.5))

    for e in ['e11','e21','e31','e12','e22','e32','e13','e23','e33']:
        ocp.bound(e,(-1.1,1.1))

    for d in ['dx','dy','dz']:
        ocp.bound(d,(-70,70))

    for w in ['w1','w2','w3']:
        ocp.bound(w,(-4*pi,4*pi))

    ocp.bound('endTime',(0.5,5))
    ocp.bound('w0',(10,10))
    ocp.bound('energy',(-1e6,1e6))

    # boundary conditions
    ocp.bound('energy',(0,0),timestep=0,quiet=True)
    ocp.bound('y',(0,0),timestep=0,quiet=True)
    
    # objective function
    obj = 0
    for k in range(nk):
        u = ocp.uVec(k)
        ddr = ocp.lookup('ddr',timestep=k)
        aileron = ocp.lookup('aileron',timestep=k)
        elevator = ocp.lookup('elevator',timestep=k)
        
        aileronSigma = 0.1
        elevatorSigma = 0.1
        ddrSigma = 5.0
        
        ailObj = aileron*aileron / (aileronSigma*aileronSigma)
        eleObj = elevator*elevator / (elevatorSigma*elevatorSigma)
        winchObj = ddr*ddr / (ddrSigma*ddrSigma)
        
        obj += ailObj + eleObj + winchObj
    ocp.setObjective( 1e1*C.sumAll(obj)/float(nk) + ocp.lookup('energy',timestep=-1)/ocp.lookup('endTime') )

    # callback function
    class MyCallback:
        def __init__(self):
            self.iter = 0 
        def __call__(self,f,*args):
            self.iter = self.iter + 1
            xOpt = numpy.array(f.input(C.NLP_X_OPT))

            opt = ocp.devectorize(xOpt)
            xup = opt['vardict']
            
            kiteProtos = []
            for k in range(0,nk):
                j = nicp*(deg+1)*k
                kiteProtos.append( kiteproto.toKiteProto(C.DMatrix(opt['x'][:,j]),
                                                         C.DMatrix(opt['u'][:,j]),
                                                         C.DMatrix(opt['p']),
                                                         conf['kite']['zt'],
                                                         conf['carousel']['rArm'],
                                                         zeroDelta=True) )
#            kiteProtos = [kiteproto.toKiteProto(C.DMatrix(opt['x'][:,k]),C.DMatrix(opt['u'][:,k]),C.DMatrix(opt['p']), conf['kite']['zt'], conf['carousel']['rArm'], zeroDelta=True) for k in range(opt['x'].shape[1])]
            
            mc = kite_pb2.MultiCarousel()
            mc.css.extend(list(kiteProtos))
            
            mc.messages.append("endTime: "+str(xup['endTime']))
            mc.messages.append("w0: "+str(xup['w0']))
            mc.messages.append("iter: "+str(self.iter))

            # bounds feedback
#            lbx = ocp.solver.input(C.NLP_LBX)
#            ubx = ocp.solver.input(C.NLP_UBX)
#            violations = boundsFeedback(xOpt,lbx,ubx,ocp.bndtags,tolerance=1e-9)
#            for name in violations:
#                violmsg = "violation!: "+name+": "+str(violations[name])
#                mc.messages.append(violmsg)
            
            publisher.send_multipart(["multi-carousel", mc.SerializeToString()])


    # solver
    solverOptions = [ ("expand_f",True)
                    , ("expand_g",True)
                    , ("generate_hessian",True)
#                     ,("qp_solver",C.NLPQPSolver)
#                     ,("qp_solver_options",{'nlp_solver': C.IpoptSolver, "nlp_solver_options":{"linear_solver":"ma57"}})
                    , ("linear_solver","ma57")
                    , ("max_iter",1000)
                    , ("tol",1e-9)
#                    , ("Timeout", 1e6)
#                    , ("UserHM", True)
#                    , ("ScaleConIter",True)
#                    , ("ScaledFD",True)
#                    , ("ScaledKKT",True)
#                    , ("ScaledObj",True)
#                    , ("ScaledQP",True)
                    ]
    
    # initial guess
#    ocp.guessX(x0)
#    for k in range(0,nk+1):
#        val = 2.0*pi*k/nk
#        ocp.guess('delta',val,timestep=k,quiet=True)
#
#    ocp.guess('aileron',0)
#    ocp.guess('elevator',0)
#    ocp.guess('tc',0)
    ocp.guess('endTime',5.4)
#
#    ocp.guess('ddr',0)
    ocp.guess('w0',10)

    print "setting up collocation..."
    ocp.setupCollocation(ocp.lookup('endTime'))
    print "setting up solver..."
    ocp.setupSolver( solverOpts=solverOptions,
                     callback=MyCallback() )
    return ocp
def main():
    nk = 15

    print "creating model"
    dae = pendulum_model.pendulum_model()
    dae.addP('endTime')

    print "setting up OCP"
    ocp = Coll(dae, nk=nk,nicp=1,deg=4)
    
    # constrain invariants
    def invariantErrs():
        f = C.SXFunction( [dae.xVec(),dae.uVec(),dae.pVec()]
                        , [dae.output('invariants')]
                        )
        f.setOption('name','invariant errors')
        f.init()
        return f
    
    [c0] = invariantErrs().call([ocp.xVec(0),ocp.uVec(0),ocp.pVec()])
    ocp.constrain(c0,'==',0)

    # bounds
    r = 0.3
    ocp.bound('x',(-0.5,0.5))
    ocp.bound('z',(-0.5,0.5))
    ocp.bound('dx',(-5,5))
    ocp.bound('dz',(-5,5))
    ocp.bound('torque',(-50,50))
    ocp.bound('m',(0.3,0.5))
    ocp.bound('endTime',(0.01,1.5))

    # boundary conditions
    ocp.bound('x',(r,r),timestep=0)
    ocp.bound('z',(0,0),timestep=0)
    ocp.bound('x',(0,0),timestep=-1)
    ocp.bound('z',(-r*1.5,-r/2),timestep=-1)
    ocp.bound('dx',(0,0),timestep=0)
    ocp.bound('dz',(0,0),timestep=0)
    ocp.bound('dx',(0,0),timestep=-1)
    ocp.bound('dz',(0,0),timestep=-1)

    # make the solver
    ocp.setObjective(ocp.lookup('endTime'))
    
    context   = zmq.Context(1)
    publisher = context.socket(zmq.PUB)
    publisher.bind("tcp://*:5563")

    # callback function
    class MyCallback:
      def __init__(self):
        self.iter = 0 
      def __call__(self,f,*args):
          xOpt = numpy.array(f.input(C.NLP_X_OPT))

          self.iter = self.iter + 1
          xup = ocp.devectorize(xOpt)['vardict']
          
          po = kite_pb2.PendulumOpt()
          po.x.extend(list(xup['x']))
          po.z.extend(list(xup['z']))
          po.endTime = xup['endTime']
          po.iters = self.iter
          publisher.send_multipart(["pendulum-opt", po.SerializeToString()])
        
    # solver
    solverOptions = [ ("linear_solver","ma27")
#                    , ("derivative_test","first-order")
                    , ("expand_f",True)
                    , ("expand_g",True)
                    , ("generate_hessian",True)
#                    , ("max_iter",1000)
                    , ("tol",1e-4)
                    ]

    constraintFunOptions = [('numeric_jacobian',False)]

    # initial conditions
    ocp.guessX([r,0,0,0])
    ocp.guess('torque',0)
    ocp.guess('m',0)
    ocp.guess('endTime',0.3)

    ocp.setupCollocation( ocp.lookup('endTime') )
    ocp.setupSolver( solverOpts=solverOptions,
                     constraintFunOpts=constraintFunOptions,
                     callback=MyCallback() )
    
    opt = ocp.solve()

    # Plot the results
    plt.figure(1)
    plt.clf()
    legend = []
    for name in ocp.dae.xNames():
        legend.append(name)
        plt.plot(opt['tgrid'],opt['vardict'][name])#,'--')
    for name in ocp.dae.uNames():
        legend.append(name)
        plt.plot(opt['tgrid'],opt['vardict'][name]/20)#,'--')
    plt.title("pendulum swingup optimization")
    plt.xlabel('time')
    plt.legend(legend)
    plt.grid()
    plt.show()