#!/usr/bin/env python3.8

# This test shared with RoboDK in the forum thread:
# https://robodk.com/forum/Thread-Targets-don-t-update-their-joints-when-the-targets-move

import os
from robodk import robolink
from robodk import robomath
import time

def joints_similar(a, b):
    a = a.tolist()
    b = b.tolist()
    for j in range(len(a)):
        if abs(a[j] - b[j]) > 0.01:
            return False
    return True

rdk = robolink.Robolink()
print(f"RoboDK version: {rdk.Version()}")

station_name = "RecalculateBug"

station = rdk.Item(station_name, robolink.ITEM_TYPE_STATION)
if station.Valid():
    print(f"'{station_name}' station already open.  Closing and recreating")
    rdk.setActiveStation(station)
    rdk.CloseStation()

station = rdk.AddStation(station_name)
assert station.Valid()

rdk.setActiveStation(station)

print("Make sure you point the code at a valid UR10e.robot file")
robot_filename = "/tmp/UR10e.robot"
robot = rdk.AddFile(robot_filename, station)
assert robot.Valid()

# Setting these for later when we run the post
robot.setParam("PostProcessor", "CSV")

frame = rdk.AddFrame("frame")
target = rdk.AddTarget("target", itemparent=frame)
#target.setAsJointTarget()

program = rdk.AddProgram("program")
program.MoveJ(target)

original_target_joints = target.Joints().copy()
print(f"Target PoseAbs: {target.PoseAbs().Pos()}")
print(f"Target joints: {target.Joints()}")

print(f"\nMoving target's parent reference frame")
frame.setPose(robomath.Pose(100, 1000, 0, 0, 0, 0))

if joints_similar(original_target_joints, target.Joints()):
    print(f"Bug: target joints remained the same even though target PoseAbs changed: {target.PoseAbs().Pos()}")
else:
    print(f"Good: Interesting, joints updated automatically.  Version 5.5.4 doesn't do that for me.")

target.setParam('Recalculate')

if joints_similar(original_target_joints, target.Joints()):
    print(f"Bug: Joints didn't update even after target.setParam('Recalculate')")
    assert False, "Aborting so we don't confuse ourselves"
else:
    print(f"Good: target joints are correct after target.setParam('Recalculate')")
    different_target_joints = target.Joints().copy()

print(f"\nMoving reference frame back to its original pose")
frame.setPose(robomath.Pose(0, 1000, 0, 0, 0, 0))

if joints_similar(original_target_joints, target.Joints()):
    print(f"Good: Interesting, joints updated automatically again.  Version 5.5.4 doesn't do that for me.")
else:
    print(f"target joints didn't update, so setParam('Recalculate') seems to be a command rather than a persistent setting")

target.setParam('Recalculate')

if joints_similar(original_target_joints, target.Joints()):
    print(f"Good: after another recalculate the target is back to its original joints")
else:
    print(f"WTF, joints didn't return to original pose")
    assert False, "Aborting so we don't confuse ourselves"

print(f"\nMoving target's parent reference frame again")
frame.setPose(robomath.Pose(100, 1000, 0, 0, 0, 0))

if joints_similar(different_target_joints, target.Joints()):
    print(f"Good: joints automatically updated")
else:
    print(f"Bug: once again joints haven't updated, no surprise there.")

print("Trying program.setParam('RecalculateTargets')")
program.setParam("RecalculateTargets")

if joints_similar(different_target_joints, target.Joints()):
    print(f"Good: joints updated")
else:
    print(f"Bug: Joints didn't change even after program.setParam('RecalculateTargets')")
    assert False, "Aborting so we don't confuse ourselves"

print(f"\nMoving reference frame back to its original pose")
frame.setPose(robomath.Pose(0, 1000, 0, 0, 0, 0))

if joints_similar(original_target_joints, target.Joints()):
    print(f"Good: Joints updated automatically again.")
else:
    print(f"target joints didn't update, so program.setParam('RecalculateTargets') seems to be a command rather than a persistent setting")
    print("Recalculating again to update the joints")
    program.setParam("RecalculateTargets")
    assert joints_similar(original_target_joints, target.Joints())

do_interactive_test = True
if do_interactive_test:
    print(f"\nMoving target's parent reference frame again")
    frame.setPose(robomath.Pose(100, 1000, 0, 0, 0, 0))

    if joints_similar(original_target_joints, target.Joints()):
        print(f"Bug: once again joints haven't updated, no surprise there.")
    else:
        print(f"Good: joints automatically updated")

    input("Please right click 'program' in the RoboDK GUI and select 'Recalculate targets', then hit enter")

    if joints_similar(original_target_joints, target.Joints()):
        print(f"Bug: Joints didn't change even after you manually requested recalculation")
    else:
        print(f"Good: joints updated")

    print(f"\nMoving reference frame back to its original pose and recalculating to ensure update")
    frame.setPose(robomath.Pose(0, 1000, 0, 0, 0, 0))
    program.setParam("RecalculateTargets")
    assert joints_similar(original_target_joints, target.Joints())

print(f"\nMoving target's parent reference frame again")
frame.setPose(robomath.Pose(100, 1000, 0, 0, 0, 0))

if joints_similar(original_target_joints, target.Joints()):
    print(f"Bug: once again joints haven't updated, no surprise there.")
else:
    print(f"Good: joints automatically updated")

success, _, _ = program.MakeProgram("/tmp/", run_mode=robolink.RUNMODE_MAKE_ROBOTPROG)
# assert success  # Fails since I don't want to install Pyside2, but the .csv gets created

if joints_similar(original_target_joints, target.Joints()):
    print(f"Joints didn't change even after running the postprocessor.  Did you remember to set Tools...Options...Program...Recalculate targets before generation?")
else:
    print(f"Good: joints are correct after running the postprocessor")
    print(f"Joints in /tmp/program.csv are expected to be very close to: {target.Joints().tolist()}")
    print(f"Here are the contents.  Do they match?")
    os.system("cat /tmp/program.csv")


print(f"\nMoving reference frame back to its original pose")
frame.setPose(robomath.Pose(0, 1000, 0, 0, 0, 0))
program.setParam("RecalculateTargets")
assert joints_similar(original_target_joints, target.Joints())