from robodk import robolink    # RoboDK API
from robodk import robomath    # Robot toolbox
from robodk import robolinkutils
import time

class RobotTest : # Class to set a new position for a robot and increment it there repeatedly
    def __init__(self, RobotName) :

        self.RDK = robolink.Robolink()
        
        self.Robot = self.RDK.Item(RobotName,2)
        self.RobotFrame = self.Robot.Parent()
      
        self.Robot.setFrame(self.RobotFrame)

        self.RobotPose = self.Robot.Pose()

        self.RobotPos = [10,0,-140]

        self.RobotJoints = self.Robot.Joints().list()

        self.isMovingForward = True
        self.dS = 20

    def TickL(self) : # Update Linear Position
        if self.RobotPos[1] > 200 and self.isMovingForward :
            self.isMovingForward = False
            self.dS = -20
        elif self.RobotPos[1] < -200 and not self.isMovingForward :
            self.isMovingForward = True
            self.dS = 20
            
        self.RobotPos[1] += self.dS

    def TickJ(self) : # Update Joint Position
        if self.RobotJoints[0] > 200 and self.isMovingForward :
            self.isMovingForward = False
            self.dS = -20
        elif self.RobotJoints[0] < -200 and not self.isMovingForward :
            self.isMovingForward = True
            self.dS = 20
            
        self.RobotJoints[0] += self.dS

    def MoveL(self) : # Move to Updated Position
        self.Robot.setPose(self.RobotPose.setPos(self.RobotPos))
        self.TickL()
        
    def MoveJ(self) : # Move to Updated Position
        self.Robot.setJoints(self.RobotJoints)
        self.TickJ()  

    def IK(self) : # Run IK on Robot to test API latency
        self.Robot.SolveIK(self.RobotPose)

class StressTest : # Class to Aquire a list of robots and move them one by one.
    def __init__(self) :
        self.RDK = robolink.Robolink()
        self.Frame = self.RDK.Item('Track Frame', 3)
        self.Frameb = self.RDK.Item('Track Frame 2', 3)
        self.Links = self.Frame.Childs()
        self.Linksb = self.Frameb.Childs()

        self.RobotList = []
        self.RobotListb = []

        self.Robotb = None
        for Link in self.Links :
            if Link.Type() == 3:
                self.RobotList.append(Link.Childs()[0].Childs()[0])

        for Link in self.Linksb :
            if Link.Type() == 3:
                self.RobotListb.append(Link.Childs()[0].Childs()[0])


        self.TesterList = []
        self.TesterListb = []
        for i, Robot in enumerate(self.RobotList) :
            self.TesterList.append(RobotTest(str(Robot.Name())))
            self.TesterListb.append(RobotTest(str(self.RobotListb[i].Name())))

   
    def MoveJ(self, N) :
        while N > 0 :
            for i, Robot in enumerate(self.TesterList) :
                Robot.MoveJ()
                Thread1 = Thread( target = self.TesterListb[i].IK() )  # Comment out this line to reduce latency

            N -= 1

    def MoveL(self, N) : # Solve IK
        while N > 0 :
            for i, Robot in enumerate(self.TesterList) :
                Robot.MoveL()
            
            for Robot in self.TesterListb :
                Robot.IK()

            N -= 1