import numpy as np
import mujoco_py as mj
from mujoco_py_renderer import SimulationError, XMLError, MujocoPyRenderer
from mujoco_py import (MjSim, load_model_from_xml,functions,
                       load_model_from_path, MjSimState,
                       ignore_mujoco_warnings,
                       load_model_from_mjb)


from matplotlib import pyplot as plt
import time

xml = """
<mujoco model="example">
    <compiler coordinate="global"/>
    <default>
        <geom rgba=".8 .6 .4 1"/>
    </default>
    <asset>
        <texture type="skybox" builtin="gradient" rgb1="1 1 1" rgb2=".6 .8 1" 
                 width="256" height="256"/>
    </asset>
    <worldbody>
        <light pos="0 1 1" dir="0 -1 -1" diffuse="1 1 1"/>
        <geom name="floor" pos="0 0 0" rgba="0.8 0.9 0.8 1" size="10 10 10" type="plane"/>
        <body>
			<site name="world" size="0.1" pos="0 0 0" />
			
            <geom name="first_pole" type="capsule" fromto="0 0 0  0 0 0.5" size="0.04"/>
            <joint name='a' type="hinge" pos="0 0 0" axis="0 0 1" />
            <body name="second_pole">
				<inertial pos="0 0 0" mass="0.00000001" diaginertia="1e-008 1e-008 1e-008" />
                <geom type="capsule" fromto="0 0 0.5  0.5 0 0.5" size="0.04" name="second_pole"/>
                <joint name='b' type="hinge" pos="0 0 0.5" axis="0 1 0"/>      
				<body name='third_pole'>
					<inertial pos="0 0 0" mass="0.00000001" diaginertia="1e-008 1e-008 1e-008" />
					<geom type="capsule" fromto="0.5 0 0.5  1 0 0.5" size="0.04" name="third_pole"/>
					<joint name='c' type="hinge" pos="0.5 0 0.5" axis="0 1 0"/>    
					<site name="target" size="0.1" pos="1 0 0.5" />
					<body name="mass">
						<inertial pos="1 0 0.5" mass="1e-2" diaginertia="1e-008 1e-008 1e-008" />
						<geom type="sphere" pos="1 0 0.5" size="0.2" name="mass"/>
					</body>
				</body>
            </body>
        </body>
    </worldbody>
	<actuator>
		<motor joint="a"/>
		<motor joint="b"/>
		<motor joint="c"/>
	</actuator>	
	

</mujoco>
"""

model = load_model_from_xml(xml)

sim = MjSim(model)
viewer = MujocoPyRenderer(sim)


sim.reset()
    # After reset jacobians are all zeros
sim.forward()
target_jacp = np.zeros(3 * sim.model.nv)
target_jacr= np.zeros(3 * sim.model.nv)



F=np.array([0,0,-9.81*1e-2,0,0,0]).T

#np.testing.assert_allclose(target_jacp, np.zeros(3 * sim.model.nv))
    # After first forward, jacobians are real
#sim.forward()
K_diag=2000
C_diag=100

A_diag=1e-3

K=np.identity(3)*K_diag
C=np.identity(3)*C_diag
A=np.identity(3)*A_diag

#K_diag=0.3
#C_diag=0.05


for i in range(3):
    K[i, i]=K_diag
    C[i,i]=C_diag
    A[i, i] = A_diag


x_intial=sim.data.site_xpos[1]
print(x_intial)
x_desired=np.array([0,1,0.3])

v_intial=sim.data.site_xvelp[1]
v_desired=np.array([0,0,0])

a_desired=np.array([0,0,0])
a_intial=np.array([0,0,0])


dt=sim.model.opt.timestep
#sim.data.get_site_jacp('target', jacp=target_jacp)
    # Should be unchanged after steps (zero action)
graph=[]
for _ in range(100000):
    F[:3]=np.dot(K,x_desired-x_intial)+np.dot(C,v_desired-v_intial)+np.dot(A,a_desired-a_intial)
    H = np.zeros(sim.model.nv* sim.model.nv)
    functions.mj_fullM(sim.model, H, sim.data.qM)

    sim.data.get_site_jacp('target', jacp=target_jacp)
    sim.data.get_site_jacr('target', jacr=target_jacr)
    J_L = target_jacp.reshape((3, sim.model.nv))
    J_A = target_jacr.reshape((3, sim.model.nv))
    J = np.concatenate((J_L, J_A), axis=0)
    H_L =np.dot(np.linalg.pinv(J_L.T),np.dot(H.reshape(sim.model.nv, sim.model.nv), np.linalg.pinv(J_L)))
    H_all=np.dot(np.linalg.pinv(J.T),np.dot(H.reshape(sim.model.nv, sim.model.nv), np.linalg.pinv(J)))
    #F_a=np.dot(A,0.3-sim.data.qacc)
    #action = np.dot(J_L.T, np.dot(H_L, F[:3]))+sim.data.qfrc_bias
    action = sim.data.qfrc_bias+np.dot(H.reshape(3,3),np.dot(J_L.T,F[:3]))
    #print(action)
    #action =  np.dot(J.T, F)
    sim.data.ctrl[:] = action
    sim.step()
    sim.forward()
    #print(np.max(action))
    #print(sim.data.qacc)
    viewer.render()
    x_intial = sim.data.site_xpos[1]
    a_intial=(v_intial-sim.data.site_xvelp[1])/dt
    print(a_intial)
    v_intial = sim.data.site_xvelp[1]
    normal=np.linalg.norm(x_intial-x_desired)
    #print(normal)
    if normal<0.1:
        print("in")
        if x_desired[0]==0:
            x_desired = np.array([-1, 0, 0.5])
        elif x_desired[0]==1:
            x_desired = np.array([0, 1, 0.3])
        elif x_desired[0] == -1:
            x_desired = np.array([1, 0, 0.5])


    graph.append(np.abs(x_intial-x_desired))
 #   sim.forward()


print("the desired is {} and the intial is{}".format(x_desired,x_intial))
plt.plot(graph)
plt.show()
[((58, 8, 58, 32), 'mujoco_py.load_model_from_xml', 'load_model_from_xml', ({(58, 28, 58, 31): 'xml'}, {}), '(xml)', False, 'from mujoco_py import MjSim, load_model_from_xml, functions, load_model_from_path, MjSimState, ignore_mujoco_warnings, load_model_from_mjb\n'), ((60, 6, 60, 18), 'mujoco_py.MjSim', 'MjSim', ({(60, 12, 60, 17): 'model'}, {}), '(model)', False, 'from mujoco_py import MjSim, load_model_from_xml, functions, load_model_from_path, MjSimState, ignore_mujoco_warnings, load_model_from_mjb\n'), ((61, 9, 61, 30), 'mujoco_py_renderer.MujocoPyRenderer', 'MujocoPyRenderer', ({(61, 26, 61, 29): 'sim'}, {}), '(sim)', False, 'from mujoco_py_renderer import SimulationError, XMLError, MujocoPyRenderer\n'), ((67, 14, 67, 40), 'numpy.zeros', 'np.zeros', ({(67, 23, 67, 39): '3 * sim.model.nv'}, {}), '(3 * sim.model.nv)', True, 'import numpy as np\n'), ((68, 13, 68, 39), 'numpy.zeros', 'np.zeros', ({(68, 22, 68, 38): '3 * sim.model.nv'}, {}), '(3 * sim.model.nv)', True, 'import numpy as np\n'), ((98, 10, 98, 29), 'numpy.array', 'np.array', ({(98, 19, 98, 28): '[0, 1, 0.3]'}, {}), '([0, 1, 0.3])', True, 'import numpy as np\n'), ((101, 10, 101, 27), 'numpy.array', 'np.array', ({(101, 19, 101, 26): '[0, 0, 0]'}, {}), '([0, 0, 0])', True, 'import numpy as np\n'), ((103, 10, 103, 27), 'numpy.array', 'np.array', ({(103, 19, 103, 26): '[0, 0, 0]'}, {}), '([0, 0, 0])', True, 'import numpy as np\n'), ((104, 9, 104, 26), 'numpy.array', 'np.array', ({(104, 18, 104, 25): '[0, 0, 0]'}, {}), '([0, 0, 0])', True, 'import numpy as np\n'), ((155, 0, 155, 15), 'matplotlib.pyplot.plot', 'plt.plot', ({(155, 9, 155, 14): 'graph'}, {}), '(graph)', True, 'from matplotlib import pyplot as plt\n'), ((156, 0, 156, 10), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'from matplotlib import pyplot as plt\n'), ((72, 2, 72, 34), 'numpy.array', 'np.array', ({(72, 11, 72, 33): '[0, 0, -9.81 * 0.01, 0, 0, 0]'}, {}), '([0, 0, -9.81 * 0.01, 0, 0, 0])', True, 'import numpy as np\n'), ((82, 2, 82, 16), 'numpy.identity', 'np.identity', ({(82, 14, 82, 15): '(3)'}, {}), '(3)', True, 'import numpy as np\n'), ((83, 2, 83, 16), 'numpy.identity', 'np.identity', ({(83, 14, 83, 15): '(3)'}, {}), '(3)', True, 'import numpy as np\n'), ((84, 2, 84, 16), 'numpy.identity', 'np.identity', ({(84, 14, 84, 15): '(3)'}, {}), '(3)', True, 'import numpy as np\n'), ((113, 8, 113, 44), 'numpy.zeros', 'np.zeros', ({(113, 17, 113, 43): 'sim.model.nv * sim.model.nv'}, {}), '(sim.model.nv * sim.model.nv)', True, 'import numpy as np\n'), ((114, 4, 114, 49), 'mujoco_py.functions.mj_fullM', 'functions.mj_fullM', ({(114, 23, 114, 32): 'sim.model', (114, 34, 114, 35): 'H', (114, 37, 114, 48): 'sim.data.qM'}, {}), '(sim.model, H, sim.data.qM)', False, 'from mujoco_py import MjSim, load_model_from_xml, functions, load_model_from_path, MjSimState, ignore_mujoco_warnings, load_model_from_mjb\n'), ((120, 8, 120, 42), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((138, 11, 138, 45), 'numpy.linalg.norm', 'np.linalg.norm', ({(138, 26, 138, 44): 'x_intial - x_desired'}, {}), '(x_intial - x_desired)', True, 'import numpy as np\n'), ((112, 68, 112, 96), 'numpy.dot', 'np.dot', ({(112, 75, 112, 76): 'A', (112, 77, 112, 95): '(a_desired - a_intial)'}, {}), '(A, a_desired - a_intial)', True, 'import numpy as np\n'), ((121, 16, 121, 37), 'numpy.linalg.pinv', 'np.linalg.pinv', ({(121, 31, 121, 36): 'J_L.T'}, {}), '(J_L.T)', True, 'import numpy as np\n'), ((122, 17, 122, 36), 'numpy.linalg.pinv', 'np.linalg.pinv', ({(122, 32, 122, 35): 'J.T'}, {}), '(J.T)', True, 'import numpy as np\n'), ((150, 17, 150, 43), 'numpy.abs', 'np.abs', ({(150, 24, 150, 42): '(x_intial - x_desired)'}, {}), '(x_intial - x_desired)', True, 'import numpy as np\n'), ((112, 10, 112, 38), 'numpy.dot', 'np.dot', ({(112, 17, 112, 18): 'K', (112, 19, 112, 37): '(x_desired - x_intial)'}, {}), '(K, x_desired - x_intial)', True, 'import numpy as np\n'), ((112, 39, 112, 67), 'numpy.dot', 'np.dot', ({(112, 46, 112, 47): 'C', (112, 48, 112, 66): '(v_desired - v_intial)'}, {}), '(C, v_desired - v_intial)', True, 'import numpy as np\n'), ((121, 84, 121, 103), 'numpy.linalg.pinv', 'np.linalg.pinv', ({(121, 99, 121, 102): 'J_L'}, {}), '(J_L)', True, 'import numpy as np\n'), ((122, 83, 122, 100), 'numpy.linalg.pinv', 'np.linalg.pinv', ({(122, 98, 122, 99): 'J'}, {}), '(J)', True, 'import numpy as np\n'), ((125, 54, 125, 73), 'numpy.dot', 'np.dot', ({(125, 61, 125, 66): 'J_L.T', (125, 67, 125, 72): 'F[:3]'}, {}), '(J_L.T, F[:3])', True, 'import numpy as np\n'), ((143, 24, 143, 46), 'numpy.array', 'np.array', ({(143, 33, 143, 45): '[-1, 0, 0.5]'}, {}), '([-1, 0, 0.5])', True, 'import numpy as np\n'), ((145, 24, 145, 45), 'numpy.array', 'np.array', ({(145, 33, 145, 44): '[0, 1, 0.3]'}, {}), '([0, 1, 0.3])', True, 'import numpy as np\n'), ((147, 24, 147, 45), 'numpy.array', 'np.array', ({(147, 33, 147, 44): '[1, 0, 0.5]'}, {}), '([1, 0, 0.5])', True, 'import numpy as np\n')]