import os

# Robot parameters
ROBOT_POST = 'Igus_Rebel'
ROBOT_NAME = 'igus REBEL-6DOF'
ROBOT_TYPE = 'REBEL-6DOF-00'
PROG_EXT = '.xml'

def pose_2_str(pose):
    [x, y, z, rx, ry, rz] = pose_2_xyzrpw(pose)
    return f'x="{x:.4f}" y="{y:.4f}" z="{z:.4f}" a="{rx:.4f}" b="{ry:.4f}" c="{rz:.4f}"'

def joints_2_str(joints):
    return ' '.join([f'a{i+1}="{j:.2f}"' for i, j in enumerate(joints)])

def pose_2_xyzrpw(pose):
    from robodk import pose_2_xyzrpw
    return pose_2_xyzrpw(pose)

class RobotPost(object):
    def __init__(self, robotpost=None, robotname=None, robot_axes=6, **kwargs):
        self.ROBOT_POST = robotpost
        self.ROBOT_NAME = robotname
        self.PROG = []
        self.LINE_COUNT = 0

    def ProgStart(self, progname):
        # Ensure no extra spaces or new lines before the XML declaration
        self.PROG.append('<?xml version="1.0" encoding="utf-8"?>')
        self.addline('<Program>')
        self.addline(f'<Header RobotName="{self.ROBOT_NAME}" RobotType="{ROBOT_TYPE}" GripperType="" Software="iRC V902-14-001" VelocitySetting="2" />')

    def ProgFinish(self, progname):
        self.addline('</Program>')

    def ProgSave(self, folder, progname, ask_user=False, show_result=False):
        if not os.path.exists(folder):
            os.makedirs(folder)
        
        filesave = os.path.join(folder, progname + PROG_EXT)
        with open(filesave, "w") as fid:
            for line in self.PROG:
                fid.write(line + '\n')
        print(f'SAVED: {filesave}')

    def MoveJ(self, pose, joints, conf_RLF=None):
        self.LINE_COUNT += 1
        xml = f'<Joint AbortCondition="False" Nr="{self.LINE_COUNT}" Source="Numerical" velPercent="40" acc="40" smooth="0" '
        xml += joints_2_str(joints)
        xml += ' e1="0" e2="0" e3="0" Descr="" />'
        self.addline(xml)

    def MoveL(self, pose, joints, conf_RLF=None):
        self.LINE_COUNT += 1
        xml = f'<Linear AbortCondition="False" Nr="{self.LINE_COUNT}" Source="Numerical" vel="100" acc="40" smooth="0" '
        xml += pose_2_str(pose)
        xml += ' e1="0" e2="0" e3="0" Descr="" />'
        self.addline(xml)

    def MoveC(self, pose1, joints1, pose2, joints2, conf_RLF_1=None, conf_RLF_2=None):
        self.LINE_COUNT += 1
        p1 = pose_2_xyzrpw(pose1)
        p2 = pose_2_xyzrpw(pose2)
        xml = f'<Circular Nr="{self.LINE_COUNT}" vel="100" acc="40" smooth="0" AbortCondition="False" Source="Constant" '
        xml += f'p1x="{p1[0]}" p1y="{p1[1]}" p1z="{p1[2]}" '
        xml += f'p2x="{p2[0]}" p2y="{p2[1]}" p2z="{p2[2]}" '
        xml += pose_2_str(pose2)
        xml += ' e1="0" e2="0" e3="0" useAngle="False" angle="0" Descr="" />'
        self.addline(xml)

    def setDO(self, io_var, io_value):
        self.LINE_COUNT += 1
        xml = f'<Output Nr="{self.LINE_COUNT}" Channel="DOut{io_var}" State="{str(io_value).lower()}" Descr="" />'
        self.addline(xml)

    def waitDI(self, io_var, io_value, timeout_ms=-1):
        self.LINE_COUNT += 1
        xml = f'<If Nr="{self.LINE_COUNT}" Condition="DIn{io_var}" Descr="" />'
        self.addline(xml)

    def Relative(self, pose, vel=None, acc=None, smooth=None):
        self.LINE_COUNT += 1
        xml = f'<Relative AbortCondition="False" Nr="{self.LINE_COUNT}" acc="{acc}" smooth="{smooth}" MoType="cartbase" vel="{vel}" '
        xml += pose_2_str(pose)
        xml += ' e1="0" e2="0" e3="0" Descr="" />'
        self.addline(xml)

    def setFrame(self, pose, frame_id=None, frame_name=None):
        self.LINE_COUNT += 1
        xml = f'<BaseFrame Nr="{self.LINE_COUNT}" '
        xml += pose_2_str(pose)
        if frame_name:
            xml += f' Descr="{frame_name}"'
        xml += ' />'
        self.addline(xml)

    def setTool(self, pose, tool_id=None, tool_name=None):
        self.LINE_COUNT += 1
        xml = f'<ToolFrame Nr="{self.LINE_COUNT}" '
        xml += pose_2_str(pose)
        if tool_name:
            xml += f' Descr="{tool_name}"'
        xml += ' />'
        self.addline(xml)

    def addline(self, newline):
        self.PROG.append(newline)

    def RunMessage(self, message, iscomment=False):
        if iscomment:
            self.addline(f'<!-- {message} -->')
        else:
            print(f'Message: {message}')

    def RunCode(self, code, is_function_call=False):
        if is_function_call:
            self.LINE_COUNT += 1
            self.addline(f'<Sub Nr="{self.LINE_COUNT}" File="{code}.xml" Descr="" />')
        else:
            self.addline(f'<!-- {code} -->')

def Pose(xyzrpw):
    return xyzrpw

def loadRobotPost():
    return RobotPost(ROBOT_POST, ROBOT_NAME)

if __name__ == "__main__":
    print("This is a post processor file for RoboDK and should not be executed directly.")
    print("Load this file from RoboDK:")
    print("> RoboDK > Program > Add/Edit Post Processor > Select this file")