from roboticstoolbox import models
import matplotlib.pyplot as plt
import numpy as np
from spatialmath import SE3
puma = models.DH.Puma560()
puma.tool = SE3.Tz(0.05)
T = SE3.Tx(np.arange(0.35, 0.85, 0.01)) * SE3.Tz(puma.links[0].d) * SE3.Ry(90, "deg")
sol = puma.ikine_a(T)
w = [0,0,-9.81,0,0,0]
cap = puma.payload_capacity(sol.q, [56, 97, 52,10,10,10], w)
fig, (ax1, ax2) = plt.subplots(2, 1)
ax1.plot(T.x, [np.min(c) for c in cap])
ax1.set_ylabel("payload capacity (kg)")
ax1.grid(True)
ax2.plot(T.x, [np.argmin(c) for c in cap])
ax2.set_xlabel("x-position of end-effector (m)")
ax2.set_ylabel("index of limiting joint")
ax2.grid(True)