# This example shows how to modify program instructions
# This example iterates over the selected program changing the speeds, and removing any pause instructions and adding custom program calls
# See also:
# https://robodk.com/doc/en/PythonAPI/robolink.html#robolink.Robolink.AddProgram

from robolink import *    # API to communicate with RoboDK
from robodk import *      # basic matrix operations
import json

#----------------------------------------

# Set the amplitude, in mm
Amplitude = 5

# Set the path frequency (cycles/mm)
# Frequency = 0.5
Frequency = 1/4 # -> 4 mm is one cycle

#------------------------------------
# Set the spacing, in mm
Spacing = 0.2
#Spacing = (1/Frequency)/4  #4 samples per cycle

# triangular zig zag # triangular zig zag (4 samples/cycle)
#Samples_x_Cycle = 10
#Spacing = 1/(Frequency*Samples_x_Cycle) 

#------------------------------------

# Use wave from the start
APPLY_WAVE = False

# Program that triggers arc start
START_KEYWORD = 'ArcStart'

# Program that triggers arc end
END_KEYWORD = 'ArcEnd'

# Sine wave offset in mm
Offset = 0

#---------------------------------
# To have a trapezoidal type of wave:
# Use 4 samples per wave
# Use 1/8 of the sine wave as offset
#Spacing = (1/Frequency)/4  # each point is 0.25 mm appart
#Offset = -(1/Frequency)/8

# To have a triangular type of wave:
# Use 4 samples per wave
# Use 0 of the sine wave as offset
#Spacing = (1/Frequency)/4  # each point is 0.25 mm appart
#Offset = 0

# To have a sinusoidal type of wave:
# Use 8 samples per wave or More
# Use 0 of the sine wave as offset
#Spacing = (1/Frequency)/8  # each point is 0.25 mm appart
#Offset = 0


# Side motion: set to True to have the sine wave on the XY plane of the tool (default)
# set to Falso to move up and down the Z axis
SIDE_MOTION = True

#----------------------------------------
#----------------------------------------
#----------------------------------------
#----------------------------------------
#----------------------------------------

Xcount = Offset
WaveApplied = False

# Do not change this variable
Remaining = 0

# Start the RoboDK API
RDK = Robolink()

# Calculate the samples per cycle
samples_x_cycle = 1/(Spacing * Frequency)

print("Samples per cycle: " + str(samples_x_cycle))
if samples_x_cycle < 4:
    RDK.ShowMessage("Warning: Samples per cycle is %.1f.<br>It is recommended to decrease the frequency of the sine wave or reduce the spacing. To have 4 samples per cycle or more." % (samples_x_cycle))

# Angle tolerance to consider 2 vectors parallel
ANG_TOL = 0.1 * pi/180 # in radians

def PoseSplit(pose1, pose2, delta_mm, offst=0, dtot=0):
    pose_delta = invH(pose1) * pose2
    angle = pose_angle(pose_delta)*180/pi
    dist = norm(pose_delta.Pos())
    
    d = delta_mm - offst
    
    x_list = []
    pose_list = []
    if d > dist or angle > 179.999:
        return [], [], d-dist, dtot

    x,y,z,w,p,r = Pose_2_UR(pose_delta)        
    
    steps = max(1,int(dist/delta_mm))

    xd = x/dist
    yd = y/dist
    zd = z/dist
    wd = w/dist
    pd = p/dist
    rd = r/dist
    while d < dist + 1e-3:
        dtot += delta_mm
        x_list.append(dtot)
        pose_list.append( pose1 * UR_2_Pose([xd*d,yd*d,zd*d,wd*d,pd*d,rd*d]) )
        d += delta_mm
    
    offst = d-dist
    return pose_list, x_list, offst, dtot

START_KEYWORD = START_KEYWORD.lower()
END_KEYWORD = END_KEYWORD.lower()




# Ask the user to select a program:
prog = RDK.ItemUserPick("Select a Program to create a sine wave", ITEM_TYPE_PROGRAM)
if not prog.Valid():
    print("Operation cancelled or no programs available")
    quit()

# Ask the user to enter a function call that will be added after each movement:
print("Program selected: " + prog.Name())

prog2 = RDK.AddProgram(prog.Name() + "Sine", prog.getLink(ITEM_TYPE_ROBOT))
prog2.ShowInstructions(False)
RDK.Render(False)

lastpose = None

# Iterate for all instructions (index start is 0, you can also use -1 for the last instruction)
ins_count = prog.InstructionCount()
print("Instructions: " + str(ins_count))
for ins_id in range(ins_count):
    ins_dict = prog.setParam(ins_id)    
    print("Instruction: " + str(ins_id))    
    # Note: The type is unique for each instruction and can't be changed.
    #    However, setting the Type value to -1 will delete the instruction (same as InstructionDelete())
    if ins_dict['Type'] == INS_TYPE_MOVE:
        ins_name, ins_type, mv_type, isjointtarget, pose, joints = prog.Instruction(ins_id)
        if APPLY_WAVE and mv_type == MOVE_TYPE_LINEAR and not isjointtarget:
            # split pose
            if lastpose is not None:
                # apply sine wave from lastpose to pose
                v_path = subs3(pose.Pos(), lastpose.Pos())
                
                # Make sure we have a movement
                if norm(v_path) < 0.001: 
                    print("Moving to same pose")
                else:
                    # Only split if the move is not parallel to the Z axis (we could be doing an approach)
                    ang1 = angle3(v_path, pose.VZ())
                    ang2 = angle3(v_path, lastpose.VZ())
                    if ang1 > ANG_TOL and abs(pi-ang1) > ANG_TOL and ang2 > ANG_TOL and abs(pi-ang2) > ANG_TOL:
                        hs, xs, Remaining, Xcount = PoseSplit(lastpose, pose, Spacing, Remaining, Xcount)
                        print("Splitting move %s in %i steps" % (ins_dict['Name'], len(hs)))
                        for h,x in zip(hs,xs):
                            # Calculate v_path in relative coordinate system (vx), relative to the TCP
                            vz = [0,0,1] # h.VZ()
                            vx = h[:3,:3].tr()*v_path # vector rotation
                            angxz = angle3(vx,vz)
                            if angxz > ANG_TOL and abs(pi-angxz) > ANG_TOL:
                                vy = cross(vz, vx) #amplitude direction
                                vy = normalize3(vy)
                                a = -Amplitude * sin(x*Frequency*2*pi)
                                if SIDE_MOTION:
                                    tool_move = mult3(vy, a)
                                else:
                                    tool_move = mult3(vz, a)
                                    
                                newpose = h*transl(tool_move)
                                prog2.MoveL(newpose)                                
                                
                            else:
                                print("Warning! Tool Z axis is parallel to the movement!")
                            
                        lastpose = pose
                        continue
                    
                    else:
                        print("Move is parallel to Z axis")
                

        lastpose = pose
        #prog2.setInstruction(ins_id, ins_name, ins_type, move_type, isjointtarget, pose, joints)
        
        # Add a normal move instruction
        tgt = pose
        if isjointtarget:
            tgt = joints.list()
            
        if mv_type == MOVE_TYPE_LINEAR:
            prog2.MoveL(tgt)
        elif mv_type == MOVE_TYPE_JOINT:
            prog2.MoveJ(tgt)
        else:
            RDK.ShowMessage("Circular movements not supported for splitting")
            quit()

        continue
        
    elif ins_dict['Type'] == INS_TYPE_CODE:        
        # Add a new program call
        code_lower = ins_dict['Code'].lower()
        if START_KEYWORD in code_lower:
            APPLY_WAVE = True
            WaveApplied = True
        elif END_KEYWORD in code_lower:
            APPLY_WAVE = False
            # Reset sine wave to start again
            Xcount = Offset
    
    elif ins_dict['Type'] == INS_TYPE_CHANGESPEED:        
        # Add a new program call
        speed_mms = ins_dict['Speed']
        if speed_mms < 100.0:
            APPLY_WAVE = True
            WaveApplied = True
        elif speed_mms >= 1000.0:
            APPLY_WAVE = False
            # Reset sine wave to start again
            Xcount = Offset

    # Add copy of instruction to new program
    prog2.setParam("last", ins_dict)

if not WaveApplied:
    RDK.ShowMessage("Add speed instructions to indicate the weld area (a speed of less than 100 mm/s is considered as weld area). You can also trigger program calls to ArcStart and ArcEnd to indicate the weld area.")
    quit()
  
# Update joint values for all targets
prog2.setParam("RecalculateTargets")    

# Remove selection automatically created for new instructions
RDK.setSelection([])