# -------------------------------------------------------------------------- #
# OpenSim Moco: examplePredictAndTrack.py                                    #
# -------------------------------------------------------------------------- #
# Copyright (c) 2018 Stanford University and the Authors                     #
#                                                                            #
# Author(s): Christopher Dembia                                              #
#                                                                            #
# Licensed under the Apache License, Version 2.0 (the "License"); you may    #
# not use this file except in compliance with the License. You may obtain a  #
# copy of the License at http://www.apache.org/licenses/LICENSE-2.0          #
#                                                                            #
# Unless required by applicable law or agreed to in writing, software        #
# distributed under the License is distributed on an "AS IS" BASIS,          #
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   #
# See the License for the specific language governing permissions and        #
# limitations under the License.                                             #
# -------------------------------------------------------------------------- #

import os
import math
import opensim as osim

"""
This file performs the following problems using a
double pendulum model:

    1. predict an optimal trajectory (and controls),
    2. track the states from the optimal trajectory, and
    3. track the marker trajectories from the optimal trajectory.

"""

visualize = True
# The following environment variable is set during automated testing.
if os.getenv('OPENSIM_USE_VISUALIZER') == '0':
    visualize = False

# Create a model of a double pendulum.
# ------------------------------------
def createDoublePendulumModel():
    model = osim.Model()
    model.setName("double_pendulum")

    # Create two links, each with a mass of 1 kg, center of mass at the body's
    # origin, and moments and products of inertia of zero.
    b0 = osim.Body("b0", 1, osim.Vec3(0), osim.Inertia(1))
    model.addBody(b0)
    b1 = osim.Body("b1", 1, osim.Vec3(0), osim.Inertia(1))
    model.addBody(b1)

    # Add markers to body origin locations.
    m0 = osim.Marker("m0", b0, osim.Vec3(0))
    m1 = osim.Marker("m1", b1, osim.Vec3(0))
    model.addMarker(m0)
    model.addMarker(m1)

    # Connect the bodies with pin joints. Assume each body is 1 m long.
    j0 = osim.PinJoint("j0", model.getGround(), osim.Vec3(0), osim.Vec3(0),
        b0, osim.Vec3(-1, 0, 0), osim.Vec3(0))
    q0 = j0.updCoordinate()
    q0.setName("q0")
    j1 = osim.PinJoint("j1",
        b0, osim.Vec3(0), osim.Vec3(0), b1, osim.Vec3(-1, 0, 0), osim.Vec3(0))
    q1 = j1.updCoordinate()
    q1.setName("q1")
    model.addJoint(j0)
    model.addJoint(j1)

    tau0 = osim.CoordinateActuator()
    tau0.setCoordinate(j0.updCoordinate())
    tau0.setName("tau0")
    tau0.setOptimalForce(1)
    model.addComponent(tau0)

    tau1 = osim.CoordinateActuator()
    tau1.setCoordinate(j1.updCoordinate())
    tau1.setName("tau1")
    tau1.setOptimalForce(1)
    model.addComponent(tau1)

    # Add display geometry.
    bodyGeometry = osim.Ellipsoid(0.5, 0.1, 0.1)
    transform = osim.Transform(osim.Vec3(-0.5, 0, 0))
    b0Center = osim.PhysicalOffsetFrame("b0_center", b0, transform)
    b0.addComponent(b0Center)
    b0Center.attachGeometry(bodyGeometry.clone())
    b1Center = osim.PhysicalOffsetFrame("b1_center", b1, transform)
    b1.addComponent(b1Center)
    b1Center.attachGeometry(bodyGeometry.clone())

    model.finalizeConnections()
    model.printToXML("double_pendulum.osim")
    return model

    
def solvePrediction():
    # Predict the optimal trajectory for a minimum time swing-up.
    # In the diagram below, + represents the origin, and ---o represents a link
    # in the double pendulum.
    #
    #                              o
    #                              |
    #                              o
    #                              |
    #         +---o---o            +
    #
    #       iniital pose      final pose
    #
    study = osim.MocoStudy()
    study.setName("double_pendulum_predict")

    problem = study.updProblem()

    # Model (dynamics).
    problem.setModel(createDoublePendulumModel())

    # Bounds.
    problem.setTimeBounds(0, [0, 5])
    # Arguments are name, [lower bound, upper bound],
    #                     initial [lower bound, upper bound],
    #                     final [lower bound, upper bound].
    problem.setStateInfo("/jointset/j0/q0/value", [-10, 10], 0)
    problem.setStateInfo("/jointset/j0/q0/speed", [-50, 50], 0, 0)
    problem.setStateInfo("/jointset/j1/q1/value", [-10, 10], 0)
    problem.setStateInfo("/jointset/j1/q1/speed", [-50, 50], 0, 0)
    problem.setControlInfo("/tau0", [-100, 100])
    problem.setControlInfo("/tau1", [-100, 100])

    # Cost: minimize final time and error from desired
    #       end effector position.
    ftCost = osim.MocoFinalTimeGoal()
    ftCost.setWeight(0.001)
    problem.addGoal(ftCost)

    finalCost = osim.MocoMarkerFinalGoal()
    finalCost.setName("final")
    finalCost.setWeight(1000.0)
    finalCost.setPointName("/markerset/m1")
    finalCost.setReferenceLocation(osim.Vec3(0, 2, 0))
    problem.addGoal(finalCost)


    # Configure the solver.
    solver = study.initTropterSolver()
    solver.set_num_mesh_intervals(100)
    solver.set_verbosity(2)
    solver.set_optim_solver("ipopt")

    guess = solver.createGuess()
    guess.setNumTimes(2)
    guess.setTime([0, 1])
    guess.setState("/jointset/j0/q0/value", [0, -math.pi])
    guess.setState("/jointset/j1/q1/value", [0, 2*math.pi])
    guess.setState("/jointset/j0/q0/speed", [0, 0])
    guess.setState("/jointset/j1/q1/speed", [0, 0])
    guess.setControl("/tau0", [0, 0])
    guess.setControl("/tau1", [0, 0])
    guess.resampleWithNumTimes(10)
    solver.setGuess(guess)

    # Save the problem to a setup file for reference.
    study.printToXML("examplePredictAndTrack_predict.omoco")

    # Solve the problem.
    solution = study.solve()
    solution.write("examplePredictAndTrack_predict_solution.sto")

    if visualize:
        study.visualize(solution)
    return solution
    

def computeMarkersReference(predictedSolution):
    model = createDoublePendulumModel()
    model.initSystem()
    states = predictedSolution.exportToStatesTable()

    statesTraj = osim.StatesTrajectory.createFromStatesTable(model, states)
    
    markerTrajectories = osim.TimeSeriesTableVec3()
    markerTrajectories.setColumnLabels(["/markerset/m0", "/markerset/m1"])

    for state in statesTraj:
        model.realizePosition(state)
        m0 = model.getComponent("markerset/m0")
        m1 = model.getComponent("markerset/m1")
        markerTrajectories.appendRow(state.getTime(),
            osim.RowVectorVec3([m0.getLocationInGround(state),
                                  m1.getLocationInGround(state)]))
                                  
    # Assign a weight to each marker.
    markerWeights = osim.SetMarkerWeights()
    markerWeights.cloneAndAppend(osim.MarkerWeight("/markerset/m0", 1))
    markerWeights.cloneAndAppend(osim.MarkerWeight("/markerset/m1", 5))
    
    return osim.MarkersReference(markerTrajectories, markerWeights)

    
def solveStateTracking(stateRef):
    # Predict the optimal trajectory for a minimum time swing-up.
    study = osim.MocoStudy()
    study.setName("double_pendulum_track")

    problem = study.updProblem()

    # Model (dynamics).
    problem.setModel(createDoublePendulumModel())

    # Bounds.
    # Arguments are name, [lower bound, upper bound],
    #                     initial [lower bound, upper bound],
    #                     final [lower bound, upper bound].
    finalTime = stateRef.getIndependentColumn()[-1]
    problem.setTimeBounds(0, finalTime)
    problem.setStateInfo("/jointset/j0/q0/value", [-10, 10], 0)
    problem.setStateInfo("/jointset/j0/q0/speed", [-50, 50], 0)
    problem.setStateInfo("/jointset/j1/q1/value", [-10, 10], 0)
    problem.setStateInfo("/jointset/j1/q1/speed", [-50, 50], 0)
    problem.setControlInfo("/tau0", [-150, 150])
    problem.setControlInfo("/tau1", [-150, 150])

    # Cost: track provided state data.
    stateTracking = osim.MocoStateTrackingGoal()
    stateTracking.setReference(osim.TableProcessor(stateRef))
    problem.addGoal(stateTracking)

    effort = osim.MocoControlGoal()
    effort.setName("effort")
    effort.setWeight(0.001)
    # TODO problem.addGoal(effort)

    # Configure the solver.
    solver = study.initTropterSolver()
    solver.set_num_mesh_intervals(50)
    solver.set_verbosity(2)
    solver.set_optim_solver("ipopt")
    solver.set_optim_jacobian_approximation("exact")
    solver.set_optim_hessian_approximation("exact")
    solver.set_exact_hessian_block_sparsity_mode("dense")

    # Save the problem to a setup file for reference.
    study.printToXML("examplePredictAndTrack_track_states.omoco")

    # Solve the problem.
    solution = study.solve()

    solution.write("examplePredictAndTrack_track_states_solution.sto")

    if visualize:
        study.visualize(solution)
    return solution

    
def solveMarkerTracking(markersRef, guess):
    # Predict the optimal trajectory for a minimum time swing-up.
    study = osim.MocoStudy()
    study.setName("double_pendulum_track")

    problem = study.updProblem()

    # Model (dynamics).
    problem.setModel(createDoublePendulumModel())

    # Bounds.
    # Arguments are name, [lower bound, upper bound],
    #                     initial [lower bound, upper bound],
    #                     final [lower bound, upper bound].
    finalTime = markersRef.getMarkerTable().getIndependentColumn()[-1]
    problem.setTimeBounds(0, finalTime)
    problem.setStateInfo("/jointset/j0/q0/value", [-10, 10], 0)
    problem.setStateInfo("/jointset/j0/q0/speed", [-50, 50], 0)
    problem.setStateInfo("/jointset/j1/q1/value", [-10, 10], 0)
    problem.setStateInfo("/jointset/j1/q1/speed", [-50, 50], 0)
    problem.setControlInfo("/tau0", [-100, 100])
    problem.setControlInfo("/tau1", [-100, 100])

    # Cost: track provided marker data.
    markerTracking = osim.MocoMarkerTrackingGoal()
    markerTracking.setMarkersReference(markersRef)
    problem.addGoal(markerTracking)
    
    effort = osim.MocoControlGoal()
    effort.setName("effort")
    effort.setWeight(0.0001)
    # problem.addGoal(effort)

    # Configure the solver.
    solver = study.initTropterSolver()
    solver.set_num_mesh_intervals(50)
    solver.set_verbosity(2)
    solver.set_optim_solver("ipopt")
    solver.set_optim_jacobian_approximation("exact")
    solver.set_optim_hessian_approximation("exact")
    solver.set_exact_hessian_block_sparsity_mode("dense")
    
    solver.setGuess(guess)

    # Save the problem to a setup file for reference.
    study.printToXML("examplePredictAndTrack_track_markers.omoco")

    # Solve the problem.
    solution = study.solve()

    solution.write("examplePredictAndTrack_track_markers_solution.sto")

    if visualize:
        study.visualize(solution)
    return solution
    

    
    
optimalTrajectory = solvePrediction()

markersRef = computeMarkersReference(optimalTrajectory)

trackedSolution = solveStateTracking(optimalTrajectory.exportToStatesTable())

trackedSolution2 = solveMarkerTracking(markersRef, trackedSolution)


















[((1740, 1775), 'os.getenv', 'os.getenv', (['"""OPENSIM_USE_VISUALIZER"""'], {}), "('OPENSIM_USE_VISUALIZER')\n", (1749, 1775), False, 'import os\n'), ((1930, 1942), 'opensim.Model', 'osim.Model', ([], {}), '()\n', (1940, 1942), True, 'import opensim as osim\n'), ((2924, 2949), 'opensim.CoordinateActuator', 'osim.CoordinateActuator', ([], {}), '()\n', (2947, 2949), True, 'import opensim as osim\n'), ((3087, 3112), 'opensim.CoordinateActuator', 'osim.CoordinateActuator', ([], {}), '()\n', (3110, 3112), True, 'import opensim as osim\n'), ((3286, 3315), 'opensim.Ellipsoid', 'osim.Ellipsoid', (['(0.5)', '(0.1)', '(0.1)'], {}), '(0.5, 0.1, 0.1)\n', (3300, 3315), True, 'import opensim as osim\n'), ((3385, 3437), 'opensim.PhysicalOffsetFrame', 'osim.PhysicalOffsetFrame', (['"""b0_center"""', 'b0', 'transform'], {}), "('b0_center', b0, transform)\n", (3409, 3437), True, 'import opensim as osim\n'), ((3533, 3585), 'opensim.PhysicalOffsetFrame', 'osim.PhysicalOffsetFrame', (['"""b1_center"""', 'b1', 'transform'], {}), "('b1_center', b1, transform)\n", (3557, 3585), True, 'import opensim as osim\n'), ((4222, 4238), 'opensim.MocoStudy', 'osim.MocoStudy', ([], {}), '()\n', (4236, 4238), True, 'import opensim as osim\n'), ((5085, 5109), 'opensim.MocoFinalTimeGoal', 'osim.MocoFinalTimeGoal', ([], {}), '()\n', (5107, 5109), True, 'import opensim as osim\n'), ((5183, 5209), 'opensim.MocoMarkerFinalGoal', 'osim.MocoMarkerFinalGoal', ([], {}), '()\n', (5207, 5209), True, 'import opensim as osim\n'), ((6520, 6578), 'opensim.StatesTrajectory.createFromStatesTable', 'osim.StatesTrajectory.createFromStatesTable', (['model', 'states'], {}), '(model, states)\n', (6563, 6578), True, 'import opensim as osim\n'), ((6609, 6635), 'opensim.TimeSeriesTableVec3', 'osim.TimeSeriesTableVec3', ([], {}), '()\n', (6633, 6635), True, 'import opensim as osim\n'), ((7151, 7174), 'opensim.SetMarkerWeights', 'osim.SetMarkerWeights', ([], {}), '()\n', (7172, 7174), True, 'import opensim as osim\n'), ((7335, 7391), 'opensim.MarkersReference', 'osim.MarkersReference', (['markerTrajectories', 'markerWeights'], {}), '(markerTrajectories, markerWeights)\n', (7356, 7391), True, 'import opensim as osim\n'), ((7510, 7526), 'opensim.MocoStudy', 'osim.MocoStudy', ([], {}), '()\n', (7524, 7526), True, 'import opensim as osim\n'), ((8376, 8404), 'opensim.MocoStateTrackingGoal', 'osim.MocoStateTrackingGoal', ([], {}), '()\n', (8402, 8404), True, 'import opensim as osim\n'), ((8516, 8538), 'opensim.MocoControlGoal', 'osim.MocoControlGoal', ([], {}), '()\n', (8536, 8538), True, 'import opensim as osim\n'), ((9414, 9430), 'opensim.MocoStudy', 'osim.MocoStudy', ([], {}), '()\n', (9428, 9430), True, 'import opensim as osim\n'), ((10301, 10330), 'opensim.MocoMarkerTrackingGoal', 'osim.MocoMarkerTrackingGoal', ([], {}), '()\n', (10328, 10330), True, 'import opensim as osim\n'), ((10436, 10458), 'opensim.MocoControlGoal', 'osim.MocoControlGoal', ([], {}), '()\n', (10456, 10458), True, 'import opensim as osim\n'), ((2147, 2159), 'opensim.Vec3', 'osim.Vec3', (['(0)'], {}), '(0)\n', (2156, 2159), True, 'import opensim as osim\n'), ((2161, 2176), 'opensim.Inertia', 'osim.Inertia', (['(1)'], {}), '(1)\n', (2173, 2176), True, 'import opensim as osim\n'), ((2228, 2240), 'opensim.Vec3', 'osim.Vec3', (['(0)'], {}), '(0)\n', (2237, 2240), True, 'import opensim as osim\n'), ((2242, 2257), 'opensim.Inertia', 'osim.Inertia', (['(1)'], {}), '(1)\n', (2254, 2257), True, 'import opensim as osim\n'), ((2357, 2369), 'opensim.Vec3', 'osim.Vec3', (['(0)'], {}), '(0)\n', (2366, 2369), True, 'import opensim as osim\n'), ((2402, 2414), 'opensim.Vec3', 'osim.Vec3', (['(0)'], {}), '(0)\n', (2411, 2414), True, 'import opensim as osim\n'), ((2585, 2597), 'opensim.Vec3', 'osim.Vec3', (['(0)'], {}), '(0)\n', (2594, 2597), True, 'import opensim as osim\n'), ((2599, 2611), 'opensim.Vec3', 'osim.Vec3', (['(0)'], {}), '(0)\n', (2608, 2611), True, 'import opensim as osim\n'), ((2625, 2644), 'opensim.Vec3', 'osim.Vec3', (['(-1)', '(0)', '(0)'], {}), '(-1, 0, 0)\n', (2634, 2644), True, 'import opensim as osim\n'), ((2646, 2658), 'opensim.Vec3', 'osim.Vec3', (['(0)'], {}), '(0)\n', (2655, 2658), True, 'import opensim as osim\n'), ((2750, 2762), 'opensim.Vec3', 'osim.Vec3', (['(0)'], {}), '(0)\n', (2759, 2762), True, 'import opensim as osim\n'), ((2764, 2776), 'opensim.Vec3', 'osim.Vec3', (['(0)'], {}), '(0)\n', (2773, 2776), True, 'import opensim as osim\n'), ((2782, 2801), 'opensim.Vec3', 'osim.Vec3', (['(-1)', '(0)', '(0)'], {}), '(-1, 0, 0)\n', (2791, 2801), True, 'import opensim as osim\n'), ((2803, 2815), 'opensim.Vec3', 'osim.Vec3', (['(0)'], {}), '(0)\n', (2812, 2815), True, 'import opensim as osim\n'), ((3347, 3368), 'opensim.Vec3', 'osim.Vec3', (['(-0.5)', '(0)', '(0)'], {}), '(-0.5, 0, 0)\n', (3356, 3368), True, 'import opensim as osim\n'), ((5352, 5370), 'opensim.Vec3', 'osim.Vec3', (['(0)', '(2)', '(0)'], {}), '(0, 2, 0)\n', (5361, 5370), True, 'import opensim as osim\n'), ((7208, 7245), 'opensim.MarkerWeight', 'osim.MarkerWeight', (['"""/markerset/m0"""', '(1)'], {}), "('/markerset/m0', 1)\n", (7225, 7245), True, 'import opensim as osim\n'), ((7280, 7317), 'opensim.MarkerWeight', 'osim.MarkerWeight', (['"""/markerset/m1"""', '(5)'], {}), "('/markerset/m1', 5)\n", (7297, 7317), True, 'import opensim as osim\n'), ((8436, 8465), 'opensim.TableProcessor', 'osim.TableProcessor', (['stateRef'], {}), '(stateRef)\n', (8455, 8465), True, 'import opensim as osim\n')]