# -------------------------------------------------------------------------- #
# 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)


















[((35, 3, 35, 38), 'os.getenv', 'os.getenv', ({(35, 13, 35, 37): '"""OPENSIM_USE_VISUALIZER"""'}, {}), "('OPENSIM_USE_VISUALIZER')", False, 'import os\n'), ((41, 12, 41, 24), 'opensim.Model', 'osim.Model', ({}, {}), '()', True, 'import opensim as osim\n'), ((69, 11, 69, 36), 'opensim.CoordinateActuator', 'osim.CoordinateActuator', ({}, {}), '()', True, 'import opensim as osim\n'), ((75, 11, 75, 36), 'opensim.CoordinateActuator', 'osim.CoordinateActuator', ({}, {}), '()', True, 'import opensim as osim\n'), ((82, 19, 82, 48), 'opensim.Ellipsoid', 'osim.Ellipsoid', ({(82, 34, 82, 37): '0.5', (82, 39, 82, 42): '0.1', (82, 44, 82, 47): '0.1'}, {}), '(0.5, 0.1, 0.1)', True, 'import opensim as osim\n'), ((84, 15, 84, 67), 'opensim.PhysicalOffsetFrame', 'osim.PhysicalOffsetFrame', ({(84, 40, 84, 51): '"""b0_center"""', (84, 53, 84, 55): 'b0', (84, 57, 84, 66): 'transform'}, {}), "('b0_center', b0, transform)", True, 'import opensim as osim\n'), ((87, 15, 87, 67), 'opensim.PhysicalOffsetFrame', 'osim.PhysicalOffsetFrame', ({(87, 40, 87, 51): '"""b1_center"""', (87, 53, 87, 55): 'b1', (87, 57, 87, 66): 'transform'}, {}), "('b1_center', b1, transform)", True, 'import opensim as osim\n'), ((109, 12, 109, 28), 'opensim.MocoStudy', 'osim.MocoStudy', ({}, {}), '()', True, 'import opensim as osim\n'), ((131, 13, 131, 37), 'opensim.MocoFinalTimeGoal', 'osim.MocoFinalTimeGoal', ({}, {}), '()', True, 'import opensim as osim\n'), ((135, 16, 135, 42), 'opensim.MocoMarkerFinalGoal', 'osim.MocoMarkerFinalGoal', ({}, {}), '()', True, 'import opensim as osim\n'), ((178, 17, 178, 75), 'opensim.StatesTrajectory.createFromStatesTable', 'osim.StatesTrajectory.createFromStatesTable', ({(178, 61, 178, 66): 'model', (178, 68, 178, 74): 'states'}, {}), '(model, states)', True, 'import opensim as osim\n'), ((180, 25, 180, 51), 'opensim.TimeSeriesTableVec3', 'osim.TimeSeriesTableVec3', ({}, {}), '()', True, 'import opensim as osim\n'), ((192, 20, 192, 43), 'opensim.SetMarkerWeights', 'osim.SetMarkerWeights', ({}, {}), '()', True, 'import opensim as osim\n'), ((196, 11, 196, 67), 'opensim.MarkersReference', 'osim.MarkersReference', ({(196, 33, 196, 51): 'markerTrajectories', (196, 53, 196, 66): 'markerWeights'}, {}), '(markerTrajectories, markerWeights)', True, 'import opensim as osim\n'), ((201, 12, 201, 28), 'opensim.MocoStudy', 'osim.MocoStudy', ({}, {}), '()', True, 'import opensim as osim\n'), ((223, 20, 223, 48), 'opensim.MocoStateTrackingGoal', 'osim.MocoStateTrackingGoal', ({}, {}), '()', True, 'import opensim as osim\n'), ((227, 13, 227, 35), 'opensim.MocoControlGoal', 'osim.MocoControlGoal', ({}, {}), '()', True, 'import opensim as osim\n'), ((256, 12, 256, 28), 'opensim.MocoStudy', 'osim.MocoStudy', ({}, {}), '()', True, 'import opensim as osim\n'), ((278, 21, 278, 50), 'opensim.MocoMarkerTrackingGoal', 'osim.MocoMarkerTrackingGoal', ({}, {}), '()', True, 'import opensim as osim\n'), ((282, 13, 282, 35), 'opensim.MocoControlGoal', 'osim.MocoControlGoal', ({}, {}), '()', True, 'import opensim as osim\n'), ((46, 28, 46, 40), 'opensim.Vec3', 'osim.Vec3', ({(46, 38, 46, 39): '0'}, {}), '(0)', True, 'import opensim as osim\n'), ((46, 42, 46, 57), 'opensim.Inertia', 'osim.Inertia', ({(46, 55, 46, 56): '1'}, {}), '(1)', True, 'import opensim as osim\n'), ((48, 28, 48, 40), 'opensim.Vec3', 'osim.Vec3', ({(48, 38, 48, 39): '0'}, {}), '(0)', True, 'import opensim as osim\n'), ((48, 42, 48, 57), 'opensim.Inertia', 'osim.Inertia', ({(48, 55, 48, 56): '1'}, {}), '(1)', True, 'import opensim as osim\n'), ((52, 31, 52, 43), 'opensim.Vec3', 'osim.Vec3', ({(52, 41, 52, 42): '0'}, {}), '(0)', True, 'import opensim as osim\n'), ((53, 31, 53, 43), 'opensim.Vec3', 'osim.Vec3', ({(53, 41, 53, 42): '0'}, {}), '(0)', True, 'import opensim as osim\n'), ((58, 48, 58, 60), 'opensim.Vec3', 'osim.Vec3', ({(58, 58, 58, 59): '0'}, {}), '(0)', True, 'import opensim as osim\n'), ((58, 62, 58, 74), 'opensim.Vec3', 'osim.Vec3', ({(58, 72, 58, 73): '0'}, {}), '(0)', True, 'import opensim as osim\n'), ((59, 12, 59, 31), 'opensim.Vec3', 'osim.Vec3', ({(59, 22, 59, 24): '-1', (59, 26, 59, 27): '0', (59, 29, 59, 30): '0'}, {}), '(-1, 0, 0)', True, 'import opensim as osim\n'), ((59, 33, 59, 45), 'opensim.Vec3', 'osim.Vec3', ({(59, 43, 59, 44): '0'}, {}), '(0)', True, 'import opensim as osim\n'), ((63, 12, 63, 24), 'opensim.Vec3', 'osim.Vec3', ({(63, 22, 63, 23): '0'}, {}), '(0)', True, 'import opensim as osim\n'), ((63, 26, 63, 38), 'opensim.Vec3', 'osim.Vec3', ({(63, 36, 63, 37): '0'}, {}), '(0)', True, 'import opensim as osim\n'), ((63, 44, 63, 63), 'opensim.Vec3', 'osim.Vec3', ({(63, 54, 63, 56): '-1', (63, 58, 63, 59): '0', (63, 61, 63, 62): '0'}, {}), '(-1, 0, 0)', True, 'import opensim as osim\n'), ((63, 65, 63, 77), 'opensim.Vec3', 'osim.Vec3', ({(63, 75, 63, 76): '0'}, {}), '(0)', True, 'import opensim as osim\n'), ((83, 31, 83, 52), 'opensim.Vec3', 'osim.Vec3', ({(83, 41, 83, 45): '-0.5', (83, 47, 83, 48): '0', (83, 50, 83, 51): '0'}, {}), '(-0.5, 0, 0)', True, 'import opensim as osim\n'), ((139, 35, 139, 53), 'opensim.Vec3', 'osim.Vec3', ({(139, 45, 139, 46): '(0)', (139, 48, 139, 49): '(2)', (139, 51, 139, 52): '(0)'}, {}), '(0, 2, 0)', True, 'import opensim as osim\n'), ((193, 33, 193, 70), 'opensim.MarkerWeight', 'osim.MarkerWeight', ({(193, 51, 193, 66): '"""/markerset/m0"""', (193, 68, 193, 69): '(1)'}, {}), "('/markerset/m0', 1)", True, 'import opensim as osim\n'), ((194, 33, 194, 70), 'opensim.MarkerWeight', 'osim.MarkerWeight', ({(194, 51, 194, 66): '"""/markerset/m1"""', (194, 68, 194, 69): '(5)'}, {}), "('/markerset/m1', 5)", True, 'import opensim as osim\n'), ((224, 31, 224, 60), 'opensim.TableProcessor', 'osim.TableProcessor', ({(224, 51, 224, 59): 'stateRef'}, {}), '(stateRef)', True, 'import opensim as osim\n')]