from robodk.robomath import *
from robodk.robodialogs import *
from robodk.robofileio import *
import subprocess
import os
import re

def pose_2_str(pose):
    """Converts a robot pose target to a string according to the modified syntax/format."""
    [x, y, z, r, p, w] = pose_2_xyzrpw(pose)
    return ('X%.3f Y%.3f Z%.3f A%.3f B%.3f C%.3f' % (z, x, y, 0.0, 0.0, p))

def joints_2_str(joints):
    """Converts a robot joint target to a string with adjusted rotational axes."""
    data = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L']
    joints[0] = 0.0  # Set A axis to 0
    joints[1] = 0.0  # Set B axis to 0
    str = ''
    for i in range(len(joints)):
        if i == 2:  # If we're at the C axis, use the value of the original B axis
            str += ('%s%.6f ' % ('C', joints[4]))
        else:
            str += ('%s%.6f ' % (data[i], joints[i]))
    return str[:-1]

def convert_to_inches(value):
    """Convert millimeters to inches."""
    return value / 25.4

class RobotPost(object):
    """RoboDK Post Processor for custom axis mapping, feed rates, and rotation adjustments."""
    
    PROG_EXT = 'gcode'

    def __init__(self, robotpost=None, robotname=None, robot_axes=6, **kwargs):
        self.ROBOT_POST = robotpost
        self.ROBOT_NAME = robotname
        self.nAxes = robot_axes
        self.PROG = []
        self.LOG = ''
        self.PROG_FILES = []  # Initialize PROG_FILES to store saved program paths
        self.E_value = 0.0

        # Set shift values for the desired starting point
        self.shift_x = -6.555   # Example offset for X
        self.shift_y = -4.8586  # Example offset for Y
        self.shift_z = 0.0      # Assuming Z doesn't need an offset

        # Feed rate values for G1, G2, G3
        self.feed_rate_G1 = 600    # Feed rate for G1 commands
        self.feed_rate_G2_G3 = 720 # Feed rate for G2/G3 commands
        self.current_feed_rate = None  # To track the current feed rate


    def convert_gcode_to_inches(self, gcode_path):
        """Convert X, Y, Z values in a G-code file from millimeters to inches and apply translation."""
        with open(gcode_path, "r", encoding="utf-8") as file:
            lines = file.readlines()

        with open(gcode_path, "w", encoding="utf-8") as file:
            for line in lines:
                # Convert X, Y, Z values to inches
                x_value = re.search(r'X([\d.-]+)', line)
                y_value = re.search(r'Y([\d.-]+)', line)
                z_value = re.search(r'Z([\d.-]+)', line)

                if x_value:
                    x_mm = float(x_value.group(1))
                    x_inch = convert_to_inches(x_mm)
                    x_inch_shifted, _, _ = self.apply_translation(x_inch, 0, 0, self.shift_x, 0, 0)
                    line = re.sub(r'X[\d.-]+', f'X{x_inch_shifted:.4f}', line)
                    
                if y_value:
                    y_mm = float(y_value.group(1))
                    y_inch = convert_to_inches(y_mm)
                    _, y_inch_shifted, _ = self.apply_translation(0, y_inch, 0, 0, self.shift_y, 0)
                    line = re.sub(r'Y[\d.-]+', f'Y{y_inch_shifted:.4f}', line)
                    
                if z_value:
                    z_mm = float(z_value.group(1))
                    z_inch = convert_to_inches(z_mm)
                    _, _, z_inch_shifted = self.apply_translation(0, 0, z_inch, 0, 0, self.shift_z)
                    line = re.sub(r'Z[\d.-]+', f'Z{z_inch_shifted:.4f}', line)

                file.write(line)

    def apply_translation(self, x, y, z, shift_x, shift_y, shift_z):
        """Apply a translation to the X, Y, and Z values."""
        return x + shift_x, y + shift_y, z + shift_z

    def clean_gcode_for_arcwelder(self, filepath):
        """Cleans up non-standard G-code commands for Arc Welder."""
        with open(filepath, "r", encoding="utf-8") as file:
            lines = file.readlines()

        with open(filepath, "w", encoding="utf-8") as file:
            for line in lines:
                if not line.startswith(('G1', 'G2', 'G3', ';')):
                    file.write(f"; {line}")  # Comment out non-standard lines
                else:
                    file.write(line)

    def run_arc_welder(self, input_gcode, output_gcode):
        """Run Arc Welder on the given G-code file."""
        arc_welder_command = (
            'C:\\RoboDK\\Posts\\ArcWelder\\bin\\ArcWelder.exe '
            f'"{input_gcode}" '
            f'"{output_gcode}" '
            '--resolution-mm=0.05 '
            '--path-tolerance-percent=0.01 '
            '--max-radius-mm=99999 '
            '--allow-3d-arcs '
            #'--allow-travel-arcs '
            #'--allow-dynamic-precision '
            '--max-gcode-length=0 '
            '--log-level=DEBUG'
        )
        
        try:
            result = subprocess.run(arc_welder_command, shell=True, check=True, capture_output=True, text=True)
            arc_output = result.stdout + '\n' + result.stderr
            with open(output_gcode, "a", encoding="utf-8") as fid:
                fid.write('; ArcWelder Statistics and Log\n')
                fid.write(arc_output + '\n')
        except subprocess.CalledProcessError as e:
            print(f"Error running Arc Welder: {e.output}")

    def add_feedrates_and_post_process(self, gcode_path):
        """Add feed rates and other post-processing commands to the G-code file."""
        previous_feed_rate = None  # Initialize the previous feed rate

        with open(gcode_path, "r", encoding="utf-8") as file:
            lines = file.readlines()

        with open(gcode_path, "w", encoding="utf-8") as file:
            for line in lines:
                # Check if it's a G1, G2, or G3 command
                if line.startswith('G1'):
                    # Check if we need to change the feed rate
                    if previous_feed_rate != self.feed_rate_G1:
                        line = line.strip() + f' F{self.feed_rate_G1}\n'
                        previous_feed_rate = self.feed_rate_G1  # Update the feed rate tracker

                elif line.startswith('G2') or line.startswith('G3'):
                    # Check if we need to change the feed rate
                    if previous_feed_rate != self.feed_rate_G2_G3:
                        line = line.strip() + f' F{self.feed_rate_G2_G3}\n'
                        previous_feed_rate = self.feed_rate_G2_G3  # Update the feed rate tracker

                # Insert any additional post-processing commands here
            
                file.write(line)

        print(f'Post-processed G-code saved to: {gcode_path}')

    def ProgStart(self, progname):
        self.addline('PROC %s()' % progname)

    def ProgFinish(self, progname):
        self.addline('ENDPROC')

    def ProgSave(self, folder, progname, ask_user=False, show_result=False):
        progname = progname + '.' + self.PROG_EXT

        # Save the initial G-code file
        filesave = os.path.join(folder, progname)
        with open(filesave, "w", encoding="utf-8") as fid:
            for line in self.PROG:
                fid.write(line + '\n')

        print(f'SAVED: {filesave}')
    
        # Convert G-code to inches and apply translation using class attributes
        self.convert_gcode_to_inches(filesave)
    
        # Clean up G-code for Arc Welder
        self.clean_gcode_for_arcwelder(filesave)
    
        # Run Arc Welder to optimize arcs
        arc_filesave = filesave.replace('.gcode', '_arc.gcode')
        self.run_arc_welder(filesave, arc_filesave)

        # Post-process the Arc Welder output to add feed rates
        self.add_feedrates_and_post_process(arc_filesave)

        # Save the final processed file
        self.PROG_FILES.append(arc_filesave)

    def MoveJ(self, pose, joints, conf_RLF=None):
        self.addline('MOVJ ' + joints_2_str(joints))

    def MoveL(self, pose, joints, conf_RLF=None):
        """Defines a linear movement with only X, Y, and Z values, including feed rate F."""
        [x, y, z, r, p, w] = pose_2_xyzrpw(pose)
    
        self.E_value += 5.0
    
        # Add the G1 command with the feed rate
        self.addline(f'G1 X{x:.3f} Y{y:.3f} Z{z:.3f}')

    def MoveC(self, pose1, joints1, pose2, joints2, conf_RLF_1=None, conf_RLF_2=None):
        """Defines a circular movement between two points using G2/G3, including feed rate F."""
        [x1, y1, z1, r1, p1, w1] = pose_2_xyzrpw(pose1)
        [x2, y2, z2, r2, p2, w2] = pose_2_xyzrpw(pose2)
    
        cx = (x1 + x2) / 2  # Center X (simple midpoint calculation)
        cy = (y1 + y2) / 2  # Center Y (simple midpoint calculation)
    
        I = cx - x1
        J = cy - y1
    
        # Determine whether it's a G2 or G3 command
        direction = 'G2' if w1 > w2 else 'G3'
        feed_rate = self.feed_rate_G2_G3  # Use the combined feed rate for G2/G3
    
        # Add the G2/G3 command with the feed rate
        self.addline(f'{direction} X{x2:.3f} Y{y2:.3f} Z{z2:.3f} I{I:.3f} J{J:.3f} F{feed_rate}')

    def setFrame(self, pose, frame_id=None, frame_name=None):
        """Sets the reference frame with the adjusted pose."""
        self.addline('BASE_FRAME ' + pose_2_str(pose))

    def setTool(self, pose, tool_id=None, tool_name=None):
        """Sets the tool frame with the adjusted pose."""
        self.addline('TOOL_FRAME ' + pose_2_str(pose))

    def Pause(self, time_ms):
        """Pauses the program."""
        if time_ms < 0:
            self.addline('PAUSE')
        else:
            self.addline('WAIT %.3f' % (time_ms * 0.001))

    def setSpeed(self, speed_mms):
        """Sets the robot speed."""
        self.addlog('setSpeed not defined (%.2f mms)' % speed_mms)

    def setAcceleration(self, accel_mmss):
        """Sets the robot acceleration."""
        self.addlog('setAcceleration not defined')

    def setSpeedJoints(self, speed_degs):
        """Sets the joint speed."""
        self.addlog('setSpeedJoints not defined')

    def setAccelerationJoints(self, accel_degss):
        """Sets the joint acceleration."""
        self.addlog('setAccelerationJoints not defined')

    def setZoneData(self, zone_mm):
        """Sets the smoothing radius (rounding, blending radius, etc.)."""
        self.addlog('setZoneData not defined (%.1f mm)' % zone_mm)

    def setDO(self, io_var, io_value):
        """Sets a digital output."""
        if type(io_var) != str:
            io_var = 'OUT[%s]' % str(io_var)
        if type(io_value) != str:
            io_value = 'TRUE' if io_value > 0 else 'FALSE'
        self.addline('%s=%s' % (io_var, io_value))

    def waitDI(self, io_var, io_value, timeout_ms=-1):
        """Waits for a digital input."""
        if type(io_var) != str:
            io_var = 'IN[%s]' % str(io_var)
        if type(io_value) != str:
            io_value = 'TRUE' if io_value > 0 else 'FALSE'
        if timeout_ms < 0:
            self.addline('WAIT FOR %s==%s' % (io_var, io_value))
        else:
            self.addline('WAIT FOR %s==%s TIMEOUT=%.1f' % (io_var, io_value, timeout_ms))

    def RunCode(self, code, is_function_call=False):
        """Executes a piece of code or a function call."""
        if is_function_call:
            code = code.replace(' ', '_')
            if not code.endswith(')'):
                code += '()'
            self.addline(code)
        else:
            self.addline(code)

    def RunMessage(self, message, iscomment=False):
        """Displays a message or a comment."""
        if iscomment:
            self.addline('% ' + message)
        else:
            self.addline('% Show message: ' + message)

    def addline(self, newline):
        """Adds a new line to the program."""
        self.PROG.append(newline)

    def addlog(self, newline):
        """Adds a new line to the log."""
        self.LOG += newline + '\n'

