using JuMP, Gurobi
"""
linearize the dynamics of the cart pull mechanism through the small angle assumption -> first term
in taylor series expansion for sin(q) and cos(q) and quazi-static assumption

inputs:
    x0 = [q₁ q₂ q₁dot q₂dot]'      #initial system state
    xf = [q₁ q₂ q₁dot q₂dot]'      #final system state
    tf = [sec]                     #time to get to final state
    h =  [sec]                     #time between collocation points

outputs:
    path = [4xk]                   #state matrix that specifies path 
"""
function SolvLinearizedCartPole(x0, xf, tf, h , q2const = false, itype = "Trap")
    
    #define model
    m = Model(solver = GurobiSolver(OutputFlag=0))
    T = 0:h:tf                   #time vector
    K = length(T)                 #number of collocation points
    @variable(m,    x[1:4,1:K])   #system state at discrete times
    @variable(m, xdot[1:4,1:K])   #derivative of system state at discrete times
    @variable(m,    u[1:1,1:K])   #control input to system along q1

    #boundry constraints - initial and final state
    @constraint(m, x[:,1] .== x0)
    @constraint(m, x[:,K] .== xf)
    
    #path constraints
    #q₂ limit
    if q2const
        for k in 1:K
            @constraint(m, x[2,k] >= xf[2,1] - pi/4)
            @constraint(m, x[2,k] <= xf[2,1] + pi/4)
        end
    end
    
    
    #q₁ limits
#     for k in 1:K
#         @constraint(m, x[1,k] >= -10)
#         @constraint(m, x[1,k] <=  10)
#     end
        
    #dynamics (path) constraints, integral form
    #model constants
    l = 1         #length of arm
    m1 = 3        #mass of cart
    m2 = 1        #mass of pendulum point mass
    g = 9.81
    linAngle = 0 #the q2 angle we linearize the system about 
    
    
    #setup the dynamics constraints linearized about q1 = 0
    if linAngle == 0
        for k in 1:K
            #notes - linearized system with first term in taylor series, and neglected coreolis acceleration
            @constraint(m,xdot[1,k] == x[3,k])
            @constraint(m,xdot[2,k] == x[4,k])
            @constraint(m,xdot[3,k] == u[1,k] + m2*g*1*x[2,k] / m1) 
            @constraint(m,xdot[4,k] == -u[1,k]  -(m1 + m2)*g*x[2,k] / m1) 

        end
    end
    
    if linAngle == pi
        for k in 1:K
            #notes - linearized system with first term in taylor series, and neglected coreolis acceleration
            @constraint(m,xdot[1,k] == x[3,k])
            @constraint(m,xdot[2,k] == x[4,k])
            @constraint(m,xdot[3,k] == u[1,k]  + m2*g*1*x[2,k] / m1) 
            @constraint(m,xdot[4,k] == u[1,k]  + (m1 + m2)*g*x[2,k] / m1) 

        end
    end

    if itype == "forwardEuler" 
        #add forward euler dynamics constraint
        for k in 1:K-1
            @constraint(m, x[:,k+1] .== x[:,k] + h*xdot[:,k])
        end
        
    elseif itype == "Trap"
        for k in 1:K-1
             #add backwards euler dynamics constraint
            @constraint(m, x[:,k+1] .== x[:,k] + .5*h*(xdot[:,k] + xdot[:,k+1]))
        end     
    end
    
    
    # minimize 2-norm (THIS IS LEAST-SQUARES)
    @objective(m, Min, sum(u.^2) ) #+ sum(x[2,:]).^2))
    solve(m)
    #print(m)
    
    control = getvalue(u)
    xopt = getvalue(x)
    xdotopt = getvalue(xdot)
    
    return (control, xopt,xdotopt)
end
 

using JuMP, Ipopt

"""
simplify the dynamics of the cart pull mechanism through the small angle assumption -> first term
in taylor series expansion for sin(q) and cos(q)

inputs:
    x0 = [q₁ q₂ q₁dot q₂dot]'      #initial system state
    xf = [q₁ q₂ q₁dot q₂dot]'      #final system state
    tf = [sec]                     #time to get to final state
    h =  [sec]                     #time between collocation points

outputs:
    path = [4xk]                   #state matrix that specifies path 
"""
function SolvNonStaticCartPole(x0, xf, tf, h ; q2const = false,  itype = "Trap")
    
    #define model
    #m = Model(solver = GurobiSolver(OutputFlag=0))
    m = Model(solver = IpoptSolver(print_level=0))
    T = 0:h:tf                    #time vector
    K = length(T)                 #number of collocation points
    @variable(m,    x[1:4,1:K])   #system state at discrete times
    @variable(m, xdot[1:4,1:K])   #derivative of system state at discrete times
    @variable(m,    u[1:1,1:K])   #control input to system along q1

    #boundry constraints - initial and final state
    @constraint(m, x[:,1] .== x0)
    @constraint(m, x[:,K] .== xf)
    
    
    #path constraints
    #q₂ limit
    if q2const
        for k in 1:K
            @constraint(m, x[2,k] >= xf[2,1] - pi/4)
            @constraint(m, x[2,k] <= xf[2,1] + pi/4)
        end
    end
    
    #q₁ limits
    for k in 1:K
        @constraint(m, x[1,k] >= -4)
        @constraint(m, x[1,k] <=  4)
    end
        
    
    #dynamics (path) constraints, integral form
    #model constants
    l = 1         #length of arm
    m1 = 3        #mass of cart
    m2 = 1        #mass of pendulum point mass
    g = 9.81
    linAngle = pi
    
    #setup the dynamics constraints linearized about q1 = 0
    if linAngle == 0
        for k in 1:K
            #notes - linearized system with first term in taylor series
            @constraint(m,xdot[1,k] == x[3,k])
            @constraint(m,xdot[2,k] == x[4,k])
            
            q1 = x[1,k] ; q2 = x[2,k] ; q1dot = x[3,k] ; q2dot = x[4,k]
            @NLconstraint(m,xdot[3,k] == (l*m2*q2*q2dot^2 + u[1,k] + m2*g*1*q2) / m1)
            @NLconstraint(m,xdot[4,k] == -(l*m2*1*q2*q2dot^2 + u[1,k]*1 + (m1 + m2)*g*q2) / l*m1) 
        end
    end
    
    if linAngle == pi
        for k in 1:K
            #notes - linearized system with first term in taylor series
            @constraint(m,xdot[1,k] == x[3,k])
            @constraint(m,xdot[2,k] == x[4,k])
            
            q1 = x[1,k] ; q2 = x[2,k] ; q1dot = x[3,k] ; q2dot = x[4,k]
            @NLconstraint(m,xdot[3,k] == (l*m2*q2*q2dot^2 + u[1,k] + m2*g*1*q2) / m1)
            @NLconstraint(m,xdot[4,k] == -(l*m2*1*q2*q2dot^2 - u[1,k]*1 -(m1 + m2)*g*q2) / l*m1) 

        end
    end
    
    if itype == "forwardEuler" 
        #add forward euler dynamics constraint
        for k in 1:K-1
            @constraint(m, x[:,k+1] .== x[:,k] + h*xdot[:,k])
        end
        
    elseif itype == "Trap"
        #add backwards euler dynamics constraint
        for k in 1:K-1
            @constraint(m, x[:,k+1] .== x[:,k] + .5*h*(xdot[:,k] + xdot[:,k+1]))
        end     
    end
    
    # minimize 2-norm (THIS IS LEAST-SQUARES)
    #@objective(m, Min, .5*h*sum(u[1:end-1].^2 + u[2:end].^2))
    @objective(m, Min, sum(u.^2)) 
    solve(m)
    
    control = getvalue(u)
    xopt = getvalue(x)
    xdotopt = getvalue(xdot)
    
    return (control, xopt, xdotopt)
end

using JuMP, PyPlot, Ipopt

#define initial and final states, x
x0 = [0 0 0 0]' 
xf = [3 pi 0 0]'   #wind up at the same angle, but 

"""
solve the cart pole problem with and full dynamics

inputs:
    x0 = [q₁ q₂ q₁dot q₂dot]'      #initial system state
    xf = [q₁ q₂ q₁dot q₂dot]'      #final system state
    tf = [sec]                     #time to get to final state
    h =  [sec]                     #time between collocation points

outputs:
    path = [4xk]                   #state matrix that specifies path 
"""
function SolvFullCartPole(x0, xf, tf, h; q2const = false, xStart="None" , itype="Trap")
    
    #define model
    m = Model(solver = IpoptSolver(print_level=0))
    T = 0:h:tf                   #time vector
    K = length(T)                 #number of collocation points
    @variable(m, x[1:4,1:K])   #system state at discrete times
    @variable(m, xdot[1:4,1:K])   #derivative of system state at discrete times
    @variable(m,    u[1:1,1:K])   #control input to system along q1
    if xStart != "None"  setvalue(x[:,:] , xStart) end

    #boundry constraints - initial and final state
    @constraint(m, x[:,1] .== x0)
    @constraint(m, x[:,K] .== xf)

    #path constraints
    if q2const
        for k in 1:K
            @constraint(m, x[2,k] >= xf[2,1] - pi/4)
            @constraint(m, x[2,k] <= xf[2,1] + pi/4)
        end
   end
    
      for k in 1:K
        @constraint(m, x[1,k] >= -10)
        @constraint(m, x[1,k] <=  10)
    end
    
    #model constants
    l = 1         #length of arm
    m1 = 3        #mass of cart
    m2 = 1        #mass of pendulum point mass
    g = 9.81
    
    
    #setup the full dynamics constraints
    for k in 1:K
        @constraint(m,xdot[1,k] == x[3,k])                      
        @constraint(m,xdot[2,k] == x[4,k])                         
        q1 = x[1,k] ; q2 = x[2,k] ; q1dot = x[3,k] ; q2dot = x[4,k]
        @NLconstraint(m,xdot[3,k] ==  (l*m2*sin(q2)*q2dot^2 + u[1,k] + m2*g*cos(q2)*sin(q2)) /
                                    (m1 + m2*(1-cos(q2)^2)))
        @NLconstraint(m,xdot[4,k] == -(l*m2*cos(q2)*sin(q2)*q2dot^2 + u[1,k]*cos(q2) + (m1 + m2)*g*sin(q2))/
                                    (l*m1 + l*m2*(1-cos(q2)^2)))
    end
    

    if itype == "forwardEuler" 
        #add forward euler dynamics constraint
        for k in 1:K-1
            @constraint(m, x[:,k+1] .== x[:,k] + h*xdot[:,k])
        end
        
    elseif itype == "Trap"
        for k in 1:K-1
             #add backwards euler dynamics constraint
            @constraint(m, x[:,k+1] .== x[:,k] + .5*h*(xdot[:,k] + xdot[:,k+1]))
        end     
    end

    # minimize 2-norm (THIS IS LEAST-SQUARES)
    #@objective(m, Min, .5*h*sum(u[1:end-1].^2 + u[2:end].^2))
    @objective(m, Min, sum(u.^2) ) # + 10000*sum(x[2,:] - pi).^2)
    solve(m)
    #print(m)
    
    control = getvalue(u)
    xopt = getvalue(x)
    xdotopt = getvalue(xdot)
    
    return (control, xopt, xdotopt)
end


using PyPlot

# Plotting the output of the Optimization in q coords and effort
function plotq2(x,h,tf)
    #figure(figsize=(12,4))
    title("q_2 trajectory")
    xlabel("time[sec]")
    ylabel("Angle [rad]")
    t = 0:h:tf
    plot( t, x[2,:], markersize = 2)
end

function plotq1(x,h,tf)
    title("q_1 trajectory")
    xlabel("time[sec]")
    ylabel("Distance [m]")
    t = 0:h:tf
    plot( t, x[1,:], markersize = 2)
end

function plotu(u,h,tf)
    title("Optimization Output Effort")
    xlabel("time[sec]")
    ylabel("Force [N]")
    plot(u[1,:])
end

# plotting the cart and pendulum position at each collocation point
function plotTipTraj(x)
    #Transform to x and y coordinates of the mass location
    l = 1
    K = size(x)[2]
    xCordM = zeros(1,K)
    yCordM = zeros(1,K)
    
    xCordM = x[1,:] + l*sin.(x[2,:])
    yCordM = -l*cos.(x[2,:])
    
    title("trajectory of tip in xy plane ")
    xlabel("Meters")
    ylabel("Meters")
    plot( xCordM, yCordM,".", markersize = 2)
end

# Plotting a stem plot of the cart and pendulum position at each colocation point
function plotXY(x)
    l = 1
    K = size(x)[2]
    
    # Transform to x and y coordinates of the tip location
    xCordM = zeros(1,K)
    yCordM = zeros(1,K)
    xCordM = x[1,:] + l*sin.(x[2,:])
    yCordM = -l*cos.(x[2,:])
    
    # Transform to x and y coordinates of the Cart
    xCordC = zeros(1,K)
    yCordC = zeros(1,K)
    xCordC = x[1,:]'
    
    #figure(figsize=(12,4))
    title("XY Transform of Colocation Results ")
    xlabel("Meters")
    ylabel("Meters")
    plot( xCordM, yCordM,".", markersize = 20);
    plot( xCordC, yCordC,"rs", markersize = 15);
    #legend(["Mass Location","Cart Location"])
    for i =1:length(xCordM)
     plot( [xCordM[i],xCordC[i]] , [yCordM[i], yCordC[i]],"k");
    end
end

function plotContour(xoff,yoff)
    xv = linspace(-6,4,1000)
    yv = linspace(-1.1,1.1,1000)
    function meshgrid{T}(vx::AbstractVector{T}, vy::AbstractVector{T})
        m, n = length(vy), length(vx)
        vx = reshape(vx, 1, n)
        vy = reshape(vy, m, 1)
        (repmat(vx, m, 1), repmat(vy, 1, n))
    end
    (X,Y) = meshgrid(xv,yv);

    z = [1/(sqrt(2*pi*sig^2))*exp(-((x-xoff)^2 + (y-yoff)^2)/(2*sig^2)) for y in yv, x in xv]
    contour(X, Y, z, zdir="z", 60, offset=0, linewidths=.5 );
end

# Plotting the Ellipse Constraining the motion of the pendulum
function plotEllipse(a,b,c,xoff,yoff) 
    angle = linspace(0,2*3.14595,360);
    r = zeros(1,length(angle));
    xE = zeros(1,length(angle));
    yE = zeros(1,length(angle));
    for i = 1:length(angle);
        r[i] = a*b/sqrt(b^2*cos(angle[i])^2 + a^2*sin(angle[i])^2);
        xE[i] = r[i]*cos(angle[i]) + xoff;
        yE[i] = r[i]*sin(angle[i]) + yoff;
    end
    plot(xE,yE,"g.--");
end



#define initial and final states
x0 = [0 0 0 0]' 
xf = [3 0 0 0]' 
tf = 3 ; h = .01

#evaluate each model
u1,x1,xdot1 =  SolvLinearizedCartPole(x0, xf, tf, h ; q2const = false , itype = "Trap")
u2,x2,xdot2 =  SolvNonStaticCartPole(x0, xf, tf, h ; q2const = false,  itype = "Trap")
u3,x3,xdot3 =  SolvFullCartPole(x0, xf, tf, h;  q2const = false, xStart="None" , itype= "Trap");


#plot the resulting tip trajectory
figure(figsize=(12,4))
plotTipTraj(x1)
plotTipTraj(x2)
plotTipTraj(x3)
legend(["small angle + quazi-static","small angle", "Full Dynamics" ])



#plot q1 vs. time
figure(figsize=(12,4)) 
plotq1(x1,h,tf)
plotq1(x2,h,tf)
plotq1(x3,h,tf)
legend(["small angle + quazi-static","small angle", "Full Dynamics" ])

figure(figsize=(12,4)) 
plotq2(x1,h,tf)
plotq2(x2,h,tf)
plotq2(x3,h,tf)
legend(["small angle + quazi-static","small angle", "Full Dynamics" ])

figure(figsize=(12,4)) 
plotu(u1)
plotu(u2)
plotu(u3)
legend(["small angle + quazi-static","small angle", "Full Dynamics" ])

#define initial and final states
x0 = [0 pi 0 0]' 
xf = [3 pi 0 0]' 
tf = 3 ; h = .01

#evaluate each model
u1,x1,xdot1 =  SolvLinearizedCartPole(x0, xf, tf, h ; q2const = false , itype = "Trap")
u2,x2,xdot2 =  SolvNonStaticCartPole(x0, xf, tf, h ; q2const = false,  itype = "Trap")
u3,x3,xdot3 =  SolvFullCartPole(x0, xf, tf, h;  q2const = false, xStart="None" , itype= "Trap");


#plot q1 vs. time
figure(figsize=(12,4)) 
plotq1(x1,h,tf)
plotq1(x2,h,tf)
plotq1(x3,h,tf)
legend(["small angle + quazi-static","small angle", "Full Dynamics" ])
figure(figsize=(12,4)) 
plotq2(x1,h,tf)
plotq2(x2,h,tf)
plotq2(x3,h,tf)
legend(["small angle + quazi-static","small angle", "Full Dynamics" ])

figure(figsize=(12,4))
plotXY(x1)
legend(["Linearized and Quasistatic Model"])
figure(figsize=(12,4))
plotXY(x2)
legend(["Linearized Model"])
figure(figsize=(12,4))
plotXY(x3)
legend(["Full Dynamics"])

#define initial and final states
x0 = [0 pi 0 0]' 
xf = [3 pi 0 0]' 
tf = 3 ; h = .01

#Evaluate the Constrained Inverted Pendulem
u3,x3,xdot3 =  SolvFullCartPole(x0, xf, tf, h;  q2const = true, xStart="None" , itype= "Trap");

#plot q1 vs. time
figure(figsize=(12,4)) 
plotq1(x3,h,tf)
legend(["Full Dynamics" ])
figure(figsize=(12,4)) 
plotq2(x3,h,tf)
legend(["Full Dynamics" ])
figure(figsize=(12,4))
plotXY(x3)
legend(["Full Dynamics"])

#define initial and final states
x0 = [0 0 0 0]' 
xf = [3 pi 0 0]' 
tf = 3 ; h = .01

#evaluate each model
u1,x1,xdot1 =  SolvLinearizedCartPole(x0, xf, tf, h ; q2const = false , itype = "Trap")
u2,x2,xdot2 =  SolvNonStaticCartPole(x0, xf, tf, h ; q2const = false,  itype = "Trap")
u3,x3,xdot3 =  SolvFullCartPole(x0, xf, tf, h;  q2const = false, xStart="None" , itype= "Trap");

#plot the resulting tip trajectory
figure(figsize=(12,4))
plotTipTraj(x1)
plotTipTraj(x2)
plotTipTraj(x3)
legend(["small angle + quazi-static","small angle", "Full Dynamics" ]);


#plot q1 vs. time
figure(figsize=(12,4)) 
plotq1(x1,h,tf)
plotq1(x2,h,tf)
plotq1(x3,h,tf)
legend(["small angle + quazi-static","small angle", "Full Dynamics" ]);

figure(figsize=(12,4)) 
plotq2(x1,h,tf)
plotq2(x2,h,tf)
plotq2(x3,h,tf)
legend(["small angle + quazi-static","small angle", "Full Dynamics" ])

figure(figsize=(12,4)) 
plotu(u1)
plotu(u2)
plotu(u3)
legend(["small angle + quazi-static","small angle", "Full Dynamics" ])

using JuMP, PyPlot, Ipopt

#define initial and final states, x
x0 = [0 0 0 0]' 
xf = [3 pi 0 0]'   #wind up at the same angle, but 

"""
solve the cart pull problem using direct collocation and full dynamics

inputs:
    x0 = [q₁ q₂ q₁dot q₂dot]'      #initial system state
    xf = [q₁ q₂ q₁dot q₂dot]'      #final system state
    tf = [sec]                     #time to get to final state
    h =  [sec]                     #time between collocation points

outputs:
    path = [4xk]                   #state matrix that specifies path 
"""
function SolvFullCartPullObstical(x0, xf, tf, h, a, b, c, xoff, yoff, itype = "backwardEuler")
    
    #define model
    m = Model(solver = IpoptSolver(print_level=0))
    T = 0:h:tf                   #time vector
    K = length(T)                 #number of collocation points
    @variable(m,    x[1:4,1:K])   #system state at discrete times
    @variable(m, xdot[1:4,1:K])   #derivative of system state at discrete times
    @variable(m,    u[1:1,1:K])   #control input to system along q1

    #boundry constraints - initial and final state
    @constraint(m, x[:,1] .== x0)
    @constraint(m, x[:,K] .== xf)
    
    #constraints on the total cart range
    @constraint(m, x[1,:] .<= 4)
    @constraint(m,-6 .<= x[1,:])

    
    #model constants
    l = 1         #length of arm
    m1 = 3        #mass of cart
    m2 = 1        #mass of pendulum point mass
    g = 9.81
    
    
        # Elypse constraint on the mass location
    for i = 1:K
        @NLconstraint(m, ((x[1,i] + l*sin(x[2,i])) - xoff)^2/a^2 + ((-l*cos(x[2,i])) - yoff)^2/b^2 >= c )
    end
    
    #setup the full dynamics constraints
    for k in 1:K
        @constraint(m,xdot[1,k] == x[3,k])                      #this step could be wrong!!! I may need to integrate here
        @constraint(m,xdot[2,k] == x[4,k])                      #look here for an issue
         #calculate xdot    
        q1 = x[1,k] ; q2 = x[2,k] ; q1dot = x[3,k] ; q2dot = x[4,k]
        @NLconstraint(m,xdot[3,k] ==  (l*m2*sin(q2)*q2dot^2 + u[1,k] + m2*g*cos(q2)*sin(q2)) /
                                    (m1 + m2*(1-cos(q2)^2)))
        @NLconstraint(m,xdot[4,k] == -(l*m2*cos(q2)*sin(q2)*q2dot^2 + u[1,k]*cos(q2) + (m1 + m2)*g*sin(q2))/
                                    (l*m1 + l*m2*(1-cos(q2)^2)))

    end
    

    if itype == "forwardEuler" 
        #add forward euler dynamics constraint
        for k in 1:K-1
            @constraint(m, x[:,k+1] .== x[:,k] + h*xdot[:,k])
        end
        
    elseif itype == "backwardEuler"
        for k in 1:K-1
             #add backwards euler dynamics constraint
            @constraint(m, x[:,k+1] .== x[:,k] + .5*h*(xdot[:,k] + xdot[:,k+1]))
        end     
    end

    # minimize 2-norm (THIS IS LEAST-SQUARES)
    #@objective(m, Min, .5*h*sum(u[1:end-1].^2 + u[2:end].^2))
    @objective(m, Min, sum(u.^2) ) #+ sum(x[2,:]).^2))
    solve(m)
    #print(m)
    
    control = getvalue(u)
    xopt = getvalue(x)
    xdotopt = getvalue(xdot)
    
    return ( m, control, xopt, xdotopt)
end


# Elipse Constarint Paramaters
a = 1
b = 0.5
c = 1
xoff = 1.5
yoff = -1
tf = 3 ; h = .03

(m, c, x, xdot) = SolvFullCartPullObstical(x0,xf,tf,h, a, b, c, xoff, yoff)
figure(figsize=(12,4))
plotq1(x,h,tf)
plotq2(x,h,tf)
figure(figsize=(12,4))
plotu(u,h,tf)
figure(figsize=(12,4))
plotEllipse(a,b,c,xoff,yoff) 
plotXY(x)


xv = linspace(-4,6,1000)
yv = linspace(-5,5,1000)
xoff = 1.5
yoff = -2
sig = 1

# swap x and y to obtain standard coordinates
z = [1/(sqrt(2*pi*sig^2))*exp(-((x-xoff)^2 + (y-yoff)^2)/(2*sig^2)) for y in yv, x in xv]

function meshgrid{T}(vx::AbstractVector{T}, vy::AbstractVector{T})
    m, n = length(vy), length(vx)
    vx = reshape(vx, 1, n)
    vy = reshape(vy, m, 1)
    (repmat(vx, m, 1), repmat(vy, 1, n))
end
(X,Y) = meshgrid(xv,yv);

pygui(false)
figure(figsize=(12,12))
surf(X,Y, z+0.5, rstride=8, cstride=8,cmap="rainbow",edgecolor="black", linewidths=.05)
contour(X, Y, z, zdir="z", 60, offset=0, origin="lower", linewidths=.5 )
xlabel("X Coord"); ylabel("Y Coord"); zlabel("Penalty");
tight_layout()

using JuMP, PyPlot, Ipopt

#define initial and final states, x
x0 = [0 0 0 0]' 
xf = [3 pi 0 0]'   #wind up at the same angle, but 

"""
solve the cart pull problem using direct collocation and full dynamics

inputs:
    x0 = [q₁ q₂ q₁dot q₂dot]'      #initial system state
    xf = [q₁ q₂ q₁dot q₂dot]'      #final system state
    tf = [sec]                     #time to get to final state
    h =  [sec]                     #time between collocation points

outputs:
    path = [4xk]                   #state matrix that specifies path 
"""
function SolveRegularizedCartPull(x0, xf, tf, h, xoff, yoff, y, itype = "backwardEuler")
    
    #define model
    m = Model(solver = IpoptSolver(print_level=0))
    T = 0:h:tf                   #time vector
    K = length(T)                 #number of collocation points
    @variable(m,    x[1:4,1:K])   #system state at discrete times
    @variable(m, xdot[1:4,1:K])   #derivative of system state at discrete times
    @variable(m,    u[1:1,1:K])   #control input to system along q1

    #boundry constraints - initial and final state
    @constraint(m, x[:,1] .== x0)
    @constraint(m, x[:,K] .== xf)
    
    #constraints on the total cart range
    @constraint(m, x[1,:] .<= 4)
    @constraint(m,-6 .<= x[1,:])
    
    #model constants
    l = 1         #length of arm
    m1 = 3        #mass of cart
    m2 = 1        #mass of pendulum point mass
    g = 9.81
    
    #setup the full dynamics constraints
    for k in 1:K
        @constraint(m,xdot[1,k] == x[3,k])                      #this step could be wrong!!! I may need to integrate here
        @constraint(m,xdot[2,k] == x[4,k])                      #look here for an issue
         #calculate xdot    
        q1 = x[1,k] ; q2 = x[2,k] ; q1dot = x[3,k] ; q2dot = x[4,k]
        @NLconstraint(m,xdot[3,k] ==  (l*m2*sin(q2)*q2dot^2 + u[1,k] + m2*g*cos(q2)*sin(q2)) /
                                    (m1 + m2*(1-cos(q2)^2)))
        @NLconstraint(m,xdot[4,k] == -(l*m2*cos(q2)*sin(q2)*q2dot^2 + u[1,k]*cos(q2) + (m1 + m2)*g*sin(q2))/
                                    (l*m1 + l*m2*(1-cos(q2)^2)))

    end
    

    if itype == "forwardEuler" 
        #add forward euler dynamics constraint
        for k in 1:K-1
            @constraint(m, x[:,k+1] .== x[:,k] + h*xdot[:,k])
        end
        
    elseif itype == "backwardEuler"
        for k in 1:K-1
             #add backwards euler dynamics constraint
            @constraint(m, x[:,k+1] .== x[:,k] + .5*h*(xdot[:,k] + xdot[:,k+1]))
        end     
    end

    #xCordM = x[1,:] + l*sin.(x[2,:])
    #yCordM = -l*cos.(x[2,:])
    
    # minimize 2-norm (THIS IS LEAST-SQUARES)

    @NLexpression(m, reg[1,i=1:K],  1/(sqrt(2*pi*sig^2))*exp(-((((x[1,i] + l*sin(x[2,i]))-xoff)^2 + (-l*cos(x[2,i])-yoff)^2)/(2*sig^2))) ) 
    @NLobjective(m, Min, sum(u[j]^2 + y*reg[1,j] for j=1:K) ) 
    solve(m)
    #print(m)
    
    control = getvalue(u)
    xopt = getvalue(x)
    xdotopt = getvalue(xdot)
    
    return (m, control, xopt, xdotopt)
end

y = 1000000
xoff = 1.5
yoff = -1
tf = 3 ; h = .03

(m, u, x, xdot) = SolveRegularizedCartPull( x0, xf, tf, h, xoff, yoff, y)
figure(figsize=(12,4))
plotq1(x,h,tf)
plotq2(x,h,tf)
legend(["q1 Cart Position","q2 Cart Pendulum Angle"])
figure(figsize=(12,4))
plotu(u,h,tf)
legend(["Control Effort"])
figure(figsize=(12,4))
plotXY(x)
plotContour(xoff,yoff)

using JuMP, PyPlot, Ipopt
"""
solve the cart pole double pendulum problem using direct collocation and full dynamics

inputs:
    x0 = [q₁ q₂ q₁dot q₂dot]'      #initial system state
    xf = [q₁ q₂ q₁dot q₂dot]'      #final system state
    tf = [sec]                     #time to get to final state
    h =  [sec]                     #time between collocation points

outputs:
    path = [4xk]                   #state matrix that specifies path 
"""
function SolvFullCartPoleDoublePendulum(x0, xf, tf, h ; xStart = "None" , itype = "backwardEuler")
    
    #define model
    m = Model(solver = IpoptSolver(print_level=0))
    T = 0:h:tf                   #time vector
    K = length(T)                 #number of collocation points
    @variable(m, x[1:6,1:K])   #system state at discrete times
    @variable(m, xdot[1:6,1:K])   #derivative of system state at discrete times
    @variable(m,    u[1:1,1:K])   #control input to system along q1
    if xStart != "None"  setvalue(x[:,:] , xStart) end


    #boundry constraints - initial and final state
    @constraint(m, x[:,1] .== x0)
    @constraint(m, x[:,K] .== xf)

    
    #model constants
    l1 = 1        #length of arm 1
    l2 = 1        #length of arm 2
    m1 = 3        #mass of cart
    m2 = 1        #mass of pendulum 1 point mass
    m3 = 1        #mass of pendulum 2 point mass 
    g = 9.81      #gravity
    
    #path Constraints
    for k in 1:K
        @constraint(m, x[1,k] <= 10)
        @constraint(m, x[1,k] >= -10)
    end
    
    
    #setup the full dynamics constraints
    for k in 1:K
        @constraint(m,xdot[1:3,k] .== x[4:6,k])
        
        #matrix variables (unsupported)        
#         M = [m1+m2+m3              l1*(m1 + m2)*cos(q2)   m2*l2*cos(q3);
#              l1*(m1+m2)*cos(q2)    l1^2*(m1+m2)           l1*l2*m2*cos(q2-q3);
#              l2*m2*cos(q3)         l1*l2*m2*cos(q2-q3)    l2^2*m2             ]
        
#         C = [l1*(m1+m2)*q2dot^2*sin(q2) + m2*l2*q3dot^2*sin(q3);
#              -l1*l2*m2*q3dot^2*sin(q2-q3) + g*(m1+m2)*l1*sin(q2);
#              l1*l2*m2*q2dot^2*sin(q2-q3) + g*l2*m2*sin(q3)         ] 
        
                
        #solve using scalar methods (can't use matrix formulations)
        
        #extract q's from matrix variables
        q1 = x[1,k] ; q2 = x[2,k] ; q3 = x[3,k] ; q1dot = x[4,k] ; q2dot = x[5,k] ; q3dot = x[6,k]
        q1ddot = xdot[4,k] ; q2ddot = xdot[5,k] ; q3ddot = xdot[6,k]
        
        #calculate qddot's
        @NLconstraint(m,(m1+m2+m3)*q1ddot + l1*(m1 + m2)*cos(q2)*q2ddot + m2*l2*cos(q3)*q3ddot
                        -(l1*(m1+m2)*q2dot*q2dot*sin(q2) + m2*l2*q3dot*q3dot*sin(q3) + u[1,k]) == 0)
        
        @NLconstraint(m, l1*(m1+m2)*cos(q2)*q1ddot + l1*l2*(m1+m2)*q2ddot + l1*l2*m2*cos(q2-q3)*q3ddot
                        -(-l1*l2*m2*q3dot*q3dot*sin(q2-q3) + g*(m1+m2)*l1*sin(q2)) == 0)
        
         @NLconstraint(m, l2*m2*cos(q3)*q1ddot + l1*l2*m2*cos(q2-q3)*q2ddot + l2*l2*m2*q3ddot
                         -(l1*l2*m2*q2dot*q2dot*sin(q2-q3) + g*l2*m2*sin(q3)) == 0 )     
    end
        

    if itype == "forwardEuler" 
        #add forward euler dynamics constraint
        for k in 1:K-1
            @constraint(m, x[:,k+1] .== x[:,k] + h*xdot[:,k])
        end
        
    elseif itype == "backwardEuler"
        for k in 1:K-1
             #add backwards euler dynamics constraint
            @constraint(m, x[:,k+1] .== x[:,k] + .5*h*(xdot[:,k] + xdot[:,k+1]))
        end     
    end

    # minimize 2-norm (THIS IS LEAST-SQUARES)
    #@objective(m, Min, .5*h*sum(u[1:end-1].^2 + u[2:end].^2))
    @objective(m, Min, sum(u.^2))
    solve(m)
    #print(m)
    
    control = getvalue(u)
    xopt = getvalue(x)
    xdotopt = getvalue(xdot)
    #return m
    return (m,control, xopt,xdotopt)
end

x0 = [0 pi pi 0 0 0]' 
xf = [3 0 0  0 0 0]' 

m,c,x,xdot = SolvFullCartPoleDoublePendulum(x0, xf, 3, .1)

using PyPlot
figure(figsize=(12,4))
#plot(c[1,:])
plot( x[1,:], "b.-", markersize=4)
plot(x[2,:],"r")
plot(x[3,:],"r")
legend(["Angle 1","Angle 2","Cart Position"])

xx = resample(x,.1,1/45)
process2CSVPen2(xx)

"""takes a 3x3 rotation matrix and converts it to a 4x1 array of euler parameters"""
function A2P(A::Array) 
  e0 = sqrt((trace(A) + 1)/4)
  if e0 != 0
    e1 = (A[3,2] - A[2,3])/(4*e0)
    e2 = (A[1,3] - A[3,1])/(4*e0)
    e3 = (A[2,1] - A[1,2])/(4*e0)
  end

  if e0 == 0  #implies Χ = π
    #figure out which e terms are non-zero
    e1flag = false; e2flag = false; e3flag = false;
    if A[1,1] + 1 != 0  e1flag = true end
    if A[2,2] + 1 != 0  e2flag = true end
    if A[3,3] + 1 != 0  e3flag = true end

    if e1flag
      e1 = sqrt((A[1,1] + 1)/2)
      e2 = (A[2,1] + A[1,2])/(4*e1)
      e3 = (A[3,1] + A[1,3])/(4*e1)
    elseif e2flag
      e2 = sqrt((A[2,2] + 1)/2)
      e1 = (A[2,1] + A[1,2])/(4*e2)
      e3 = (A[3,2] + A[2,3])/(4*e2)
    elseif e3flag
      e3 = sqrt((A[3,3] + 1)/2)
      e1 = (A[3,1] + A[1,3])/(4*e3)
      e2 = (A[3,2] + A[2,3])/(4*e3)
    else
      warn("something is wrong with A value  $A ")
    end
  end

  p = [e0 e1 e2 e3]'
end

#principle rotations
Rx(Θ) = [1 0 0 ; 0 cos(Θ) -sin(Θ) ; 0 sin(Θ) cos(Θ)]
Ry(Θ) = [ cos(Θ) 0 sin(Θ) ; 0 1 0 ; -sin(Θ) 0 cos(Θ)]
Rz(Θ) = [ cos(Θ) -sin(Θ) 0 ; sin(Θ) cos(Θ) 0  ; 0 0 1]



"""
treat the cart pull problem as an initial value problem and numerically integrate forward in time 
with the full dynamics to get the best estimate of the "real" trajectory

the approach will be a simple forward euler integrator

inputs:
    x0 =   [q₁ q₂ q₁dot q₂dot]'      #initial system state
    u[k] = [u₁ u₂ ... uₙ]            #discrete sampled control value
    tf =   [sec]                     #time between collocation points
    h =    [sec]                     #time between collocation points
    method = str                     #level of assumptions being made 

outputs:
    x[k] = [4 x n]                   
    

"""
function pullCartForwardDynamics(x0,u,tf,h,method = 1)
    #setup
    tspan = 0:h:tf
    x =    zeros(4,length(tspan)+1)
    xdot = zeros(4,length(tspan)+1)
    x[:,1] = x0
    u = resample(u,tf/(length(u)-1),h) #make u the right size to be sampled
    
    
    
    #forward euler integration
    for (k,t) in enumerate(tspan)
        xdot[:,k] = calcxdot(x[:,k],u[k],h,method) 
        x[:,k+1]  = x[:,k] + h*xdot[:,k]
    end
    
    return x
end

"""
calcuate xdot from the specified system dynamics using forward euler
"""
function calcxdot(x,u,h,method)
    #model constants
    l = 1         #length of arm
    m1 = 1        #mass of cart
    m2 = 1        #mass of pendulum 
    g = 9.81
    
    
    xdot = zeros(4,1)
    
    #calculate xdot    
    q1 = x[1,1] ; q2 = x[2,1] ; q1dot = x[3,1] ; q2dot = x[4,1]
    xdot[3,1] =  (l*m2*sin(q2)*q2dot^2 + u + m2*g*cos(q2)*sin(q2)) / (m1 + m2*(1-cos(q2)^2))
    xdot[4,1] = -(l*m2*cos(q2)*sin(q2)*q2dot^2 + u*cos(q2) + (m1 + m2)*g*sin(q2))/ (l*m1 + l*m2*(1-cos(q2)^2))
    
    #forward euler step
    xdot[1:2,1] = x[3:4,1] + h*xdot[3:4,1]
    
    return xdot
    
end


using DataFrames
"""
convert state vector x for the cart pull problem from a minimal set of generalized coordinates to a
maximal set of cartesian coordinates using euler parameters to describe rotation. additionally, 
rotations will be made so that things look right in unity

inputs:
    x - state vector resulting from from our control optimization
    fln - string of file location
outputs:
    csv - write a csv to the proper location
    
assumptions:
-This code assumes the original simulation is resampled at 30Hz so the visualization timing is correct
- left handed coordinates as defined in unity:
    * y is up
    * x is left
    * z is out of the page

movement of the cart occures along the x, and rotation of the pendulum occures about z axis
"""
function process2CSV(x,path = "./pullCartVis3D/Assets/Data/data.csv")
    
    #extract state, we only need q1 and q2
    x = x[1:2,:]
    
    #build out fill state matrix  xFull = [r₁;r₂;p₁;p₂]
    nb = 2;
    xFull = zeros(7*nb, size(x)[2])  
    
    #place location parameters appropriately
    xFull[1,:] = x[1,:]
    xFull[4,:] = x[1,:]
    xFull[6,:] = .15*ones(size(x)[2],1)                 #unity compensation
    
    #place orientation parameters appropriately
    for i in range(1,size(x)[2])  
        xFull[7:10,i]   = A2P(Rx(pi/2)*eye(3))          #unity compensation
        xFull[11:14,i] = A2P(Rx(pi/2)*Ry(x[2,i]))       #unity compensation
    end
        
    #print to CSV
    xFull = convert(DataFrame, xFull)
    writetable(path,xFull)
end


function process2CSVPen2(x,path = "./pullCartVis3D/Assets/Data/data.csv")
    
    #extract state, we only need q1 - q3
    x = x[1:3,:]
    
    #build out full state matrix  xFull = [r₁;r₂;r₃;p₁;p₂;p₃]
    nb = 3;
    xFull = zeros(7*nb, size(x)[2])  
    
    #place r locations
    xFull[1,:] =  -x[1,:]
    xFull[4,:] =  -x[1,:]
    xFull[7,:] =  -x[1,:] - .85*sin.(x[2,:])                #forward kinematics from unity origin
    xFull[8,:] =  .85*cos.(x[2,:])                         #forward kinematics from unity origin
    
    xFull[6,:] = .15*ones(size(x)[2],1)                   #unity compensation
    xFull[9,:] = .15*ones(size(x)[2],1)                   #unity compensation
    
    #place orientation parameters appropriately
    for i in range(1,size(x)[2])  
        xFull[10:13,i]   = A2P(Rx(pi/2)*eye(3))           #unity compensation
        xFull[14:17,i]   = A2P(Rx(pi/2)*Ry(x[2,i]))       #unity compensation
        xFull[18:21,i]   = A2P(Rx(pi/2)*Ry(pi)*Ry(x[3,i]))       #unity compensation
    end
        
    #print to CSV
    xFull = convert(DataFrame, xFull)
    writetable(path,xFull)
end


"""
linearly interpolate between points on a grid u[k]
"""
function interpolate(u::Array,h::Float64,t::Float64)
    if ndims(u) == 1 u = reshape(u,1,length(u)) end   #robust to 1d arrays
    k = Int(floor(t/h) + 1)
    if k == size(u)[2] return u[k] end              #you sampled at the exact last point
    deltat = t - (k-1)*h;                            #correcting for 1 index
    u_t = u[k] + deltat*(u[k+1]-u[k])/h
    return u_t
end
    
"""
resample an x array at the requested sample rate
"""
function resample(x,h0,h1)
    if ndims(x) == 1 x = reshape(x,1,length(x)) end   #robust to 1d arrays
    tf = h0*(size(x)[2] - 1)
    colsNew = Int(floor(tf/h1) + 1)
    xnew = zeros(size(x)[1],colsNew)
    for row in 1:size(x)[1]
        for col in 1:colsNew
                xnew[row,col] = interpolate(x[row:row,:],h0,(col-1)*h1)
        end
    end
    return xnew
end

"""
seed the starting state of the double pendulum with linear change between starting and ending state
"""
function seedDoublePendulum(x0, xf,tf,h)
    T = 0:h:tf                   
    K = length(T)
    x = zeros(6,K)
    x[1,:] = linspace(x0[1,1],xf[1,1],K)
    x[2,:] = linspace(x0[2,1],xf[2,1],K)
    x[3,:] = linspace(x0[3,1],xf[3,1],K)
end