################################################
###### VIBRATIONS AND SPEED OPTIMIZATIONS ######
################################################
# Written by Frix_x#0161 #
# @version: 2.2

# CHANGELOG:
#   v2.2: allow custom accel values and set sane defaults (3k) to avoid using very high accels that produce bad results
#   v2.1: allow decimal entries for speed and increment and added the E axis as an option to be neasured
#   v2.0: added the possibility to measure mutliple axis
#   v1.0: first speed and vibrations optimization macro


[gcode_macro VIBRATIONS_CALIBRATION]
gcode:
    {% set size = params.SIZE|default(60)|int %} # size of the area where the movements are done
    {% set direction = params.DIRECTION|default('XY') %} # can be set to either XY, AB, ABXY, A, B, X, Y, Z
    {% set z_height = params.Z_HEIGHT|default(20)|int %} # z height to put the toolhead before starting the movements
    {% set verbose = params.VERBOSE|default(true) %} # Wether to log the current speed in the console

    {% set min_speed = params.MIN_SPEED|default(20)|float * 60 %} # minimum feedrate for the movements
    {% set max_speed = params.MAX_SPEED|default(200)|float * 60 %} # maximum feedrate for the movements
    {% set speed_increment = params.SPEED_INCREMENT|default(2)|float * 60 %} # feedrate increment between each move
    {% set feedrate_travel = params.TRAVEL_SPEED|default(200)|int * 60 %} # travel feedrate between moves
    {% set accel = params.ACCEL|default(3000)|int %} # accel value used to move on the pattern
    {% set accel_chip = params.ACCEL_CHIP|default("adxl345") %} # ADXL chip name in the config

    {% set mid_x = printer.toolhead.axis_maximum.x|float / 2 %}
    {% set mid_y = printer.toolhead.axis_maximum.y|float / 2 %}
    {% set nb_samples = ((max_speed - min_speed) / speed_increment + 1) | int %}

    {% set accel = [accel, printer.configfile.settings.printer.max_accel]|min %}
    {% set old_accel = printer.toolhead.max_accel %}
    {% set old_accel_to_decel = printer.toolhead.max_accel_to_decel %}
    {% set old_sqv = printer.toolhead.square_corner_velocity %}

    {% set direction_factor = {
        'XY'  : {
            'start' : {'x': -0.5, 'y': -0.5 },
            'move_factors' : {
                '0' : {'x': 0.5, 'y': -0.5, 'z': 0.0 },
                '1' : {'x': 0.5, 'y': 0.5, 'z': 0.0 },
                '2' : {'x': -0.5, 'y': 0.5, 'z': 0.0 },
                '3' : {'x': -0.5, 'y': -0.5, 'z': 0.0 }
                }
            },
        'AB' : {
            'start' : {'x': 0.0, 'y': 0.0 },
            'move_factors' : {
                '0' : {'x': 0.5, 'y': -0.5, 'z': 0.0 },
                '1' : {'x': -0.5, 'y': 0.5, 'z': 0.0 },
                '2' : {'x': 0.0, 'y': 0.0, 'z': 0.0 },
                '3' : {'x': 0.5, 'y': 0.5, 'z': 0.0 },
                '4' : {'x': -0.5, 'y': -0.5, 'z': 0.0 },
                '5' : {'x': 0.0, 'y': 0.0, 'z': 0.0 }
                }
            },
        'ABXY' : {
            'start' : {'x': -0.5, 'y': 0.5 },
            'move_factors' : {
                '0' : {'x': -0.5, 'y': -0.5, 'z': 0.0 },
                '1' : {'x': 0.5, 'y': -0.5, 'z': 0.0 },
                '2' : {'x': -0.5, 'y': 0.5, 'z': 0.0 },
                '3' : {'x': 0.5, 'y': 0.5, 'z': 0.0 },
                '4' : {'x': -0.5, 'y': -0.5, 'z': 0.0 },
                '5' : {'x': -0.5, 'y': 0.5, 'z': 0.0 }
                }
            },
        'B'  : {
            'start' : {'x': 0.5, 'y': 0.5 },
            'move_factors' : {
                '0' : {'x': -0.5, 'y': -0.5, 'z': 0.0 },
                '1' : {'x': 0.5, 'y': 0.5, 'z': 0.0 }
                }
            },
        'A'  : {
            'start' : {'x': -0.5, 'y': 0.5 },
            'move_factors' : {
                '0' : {'x': 0.5, 'y': -0.5, 'z': 0.0 },
                '1' : {'x': -0.5, 'y': 0.5, 'z': 0.0 }
                }
            },
        'X'  : {
            'start' : {'x': -0.5, 'y': 0.0 },
            'move_factors' : {
                '0' : {'x': 0.5, 'y': 0.0, 'z': 0.0 },
                '1' : {'x': -0.5, 'y': 0.0, 'z': 0.0 }
                }
            },
        'Y'  : {
            'start' : {'x': 0.0, 'y': 0.5 },
            'move_factors' : {
                '0' : {'x': 0.0, 'y': -0.5, 'z': 0.0 },
                '1' : {'x': 0.0, 'y': 0.5, 'z': 0.0 }
                }
            },
        'Z'  : {
            'start' : {'x': 0.0, 'y': 0.0 },
            'move_factors' : {
                '0' : {'x': 0.0, 'y': 0.0, 'z': 1.0 },
                '1' : {'x': 0.0, 'y': 0.0, 'z': 0.0 }
                }
            },
        'E'  : {
            'start' : {'x': 0.0, 'y': 0.0 },
            'move_factor' : 0.05
            }
        }
    %}

    #
    # STARTING...
    #
    {% if not 'xyz' in printer.toolhead.homed_axes %}
        { action_raise_error("Must Home printer first!") }
    {% endif %}

    {% if params.SPEED_INCREMENT|default(2)|float * 100 != (params.SPEED_INCREMENT|default(2)|float * 100)|int %}
        { action_raise_error("Only 2 decimal digits are allowed for SPEED_INCREMENT") }
    {% endif %}

    {% if (size / (max_speed / 60)) < 0.25 and direction != 'E' %}
        { action_raise_error("SIZE is too small for this MAX_SPEED. Increase SIZE or decrease MAX_SPEED!") }
    {% endif %}

    {% if not (direction in direction_factor) %}
        { action_raise_error("DIRECTION is not valid. Only XY, AB, ABXY, A, B, X, Y, Z or E is allowed!") }
    {% endif %}

    {action_respond_info("")}
    {action_respond_info("Starting speed and vibration calibration")}
    {action_respond_info("This operation can not be interrupted by normal means. Hit the \"emergency stop\" button to stop it if needed")}
    {action_respond_info("")}

    SAVE_GCODE_STATE NAME=STATE_VIBRATIONS_CALIBRATION

    M83
    G90

    # Set the wanted acceleration values (not too high to avoid oscillation, not too low to be able to reach constant speed on each segments)
    SET_VELOCITY_LIMIT ACCEL={accel} ACCEL_TO_DECEL={accel} SQUARE_CORNER_VELOCITY={[(accel / 1000), 5.0]|max}
    
    # Going to the start position
    G1 Z{z_height}
    G1 X{mid_x + (size * direction_factor[direction].start.x) } Y{mid_y + (size * direction_factor[direction].start.y)} F{feedrate_travel}

    # vibration pattern for each frequency
    {% for curr_sample in range(0, nb_samples) %}
        {% set curr_speed = min_speed + curr_sample * speed_increment %}
        {% if verbose %}
            RESPOND MSG="{"Current speed: %.2f mm/s" % (curr_speed / 60)|float}"
        {% endif %}

        ACCELEROMETER_MEASURE CHIP={accel_chip}
        {% if direction == 'E' %}
            G0 E{curr_speed*direction_factor[direction].move_factor} F{curr_speed}
        {% else %}
            {% for key, factor in direction_factor[direction].move_factors|dictsort %}
                G1 X{mid_x + (size * factor.x) } Y{mid_y + (size * factor.y)} Z{z_height + (size * factor.z)} F{curr_speed}
            {% endfor %}
        {% endif %}
        ACCELEROMETER_MEASURE CHIP={accel_chip} NAME=sp{("%.2f" % (curr_speed / 60)|float)|replace('.','_')}n1

        G4 P300
        M400
    {% endfor %}

    {% if verbose %}
        RESPOND MSG="Graphs generation... Please wait a minute or two and look in the configured folder."
    {% endif %}
    RUN_SHELL_COMMAND CMD=plot_graph PARAMS="VIBRATIONS {direction}"

    # Restore the previous acceleration values
    SET_VELOCITY_LIMIT ACCEL={old_accel} ACCEL_TO_DECEL={old_accel_to_decel} SQUARE_CORNER_VELOCITY={old_sqv}

    RESTORE_GCODE_STATE NAME=STATE_VIBRATIONS_CALIBRATION