# in klippy/stepper.py:
# change:
# for mname in ['stepper_enable', 'force_move', 'motion_report']:
# to:
# for mname in ['stepper_enable', 'force_move', 'motion_report', 'z_gantry_calibration']:

# Add z_gantry_calibration.py to klippy/extras with following content:
import math
import chelper

STALL_TIME = 0.100

# Calculate a move's accel_t, cruise_t, and cruise_v
def calc_move_time(dist, speed, accel):
    axis_r = 1.
    if dist < 0.:
        axis_r = -1.
        dist = -dist
    if not accel or not dist:
        return axis_r, 0., dist / speed, speed
    max_cruise_v2 = dist * accel
    if max_cruise_v2 < speed**2:
        speed = math.sqrt(max_cruise_v2)
    accel_t = speed / accel
    accel_decel_d = accel_t * speed
    cruise_t = (dist - accel_decel_d) / speed
    return axis_r, accel_t, cruise_t, speed

class ZGantryCalibHelper:
    def __init__(self, config):
        self.printer = config.get_printer()

        self.initial_z_height = config.getfloat('initial_z_height')
        self.initial_move_speed = config.getfloat('initial_move_speed')
        self.align_step_size = config.getfloat('align_step_size')
        self.align_step_speed = config.getfloat('align_step_speed')
        self.align_step_accel = config.getfloat('align_step_accel')
        self.backtrack_distance = config.getfloat('backtrack_distance')

        ppins = self.printer.lookup_object('pins')
        self.z_pin = ppins.setup_pin('endstop', config.get('z_pin'))
        self.z1_pin = ppins.setup_pin('endstop', config.get('z1_pin'))

        self.steppers = {}

        ffi_main, ffi_lib = chelper.get_ffi()
        self.trapq = ffi_main.gc(ffi_lib.trapq_alloc(), ffi_lib.trapq_free)
        self.trapq_append = ffi_lib.trapq_append
        self.trapq_finalize_moves = ffi_lib.trapq_finalize_moves
        self.stepper_kinematics = ffi_main.gc(
            ffi_lib.cartesian_stepper_alloc(b'x'), ffi_lib.free)

    def manual_move(self, stepper, dist, speed, accel=0.):
        toolhead = self.printer.lookup_object('toolhead')
        toolhead.flush_step_generation()
        prev_sk = stepper.set_stepper_kinematics(self.stepper_kinematics)
        prev_trapq = stepper.set_trapq(self.trapq)
        stepper.set_position((0., 0., 0.))
        axis_r, accel_t, cruise_t, cruise_v = calc_move_time(dist, speed, accel)
        print_time = toolhead.get_last_move_time()
        self.trapq_append(self.trapq, print_time, accel_t, cruise_t, accel_t,
                          0., 0., 0., axis_r, 0., 0., 0., cruise_v, accel)
        print_time = print_time + accel_t + cruise_t + accel_t
        stepper.generate_steps(print_time)
        self.trapq_finalize_moves(self.trapq, print_time + 99999.9,
                                  print_time + 99999.9)
        stepper.set_trapq(prev_trapq)
        stepper.set_stepper_kinematics(prev_sk)
        toolhead.note_mcu_movequeue_activity(print_time)
        toolhead.dwell(accel_t + cruise_t + accel_t)
        toolhead.flush_step_generation()

    def register_stepper(self, config, mcu_stepper):
        self.steppers[mcu_stepper.get_name()] = mcu_stepper

    def do_z_calib(self):
        toolhead = self.printer.lookup_object('toolhead')
        gcode = self.printer.lookup_object('gcode')
        initial_pos = toolhead.get_position()
        initial_pos[2] = self.initial_z_height

        gcode = self.printer.lookup_object('gcode')

        gcode.respond_info('Alignment started')
        gcode.respond_info('Gantry moving to initial position: Z' + str(self.initial_z_height))
        toolhead.move(initial_pos, self.initial_move_speed)

        print_time = toolhead.get_last_move_time()
        z_state = self.z_pin.query_endstop(print_time)
        z1_state = self.z1_pin.query_endstop(print_time)

        while not z_state or not z1_state:
            if not z_state:
                self.manual_move(self.steppers['stepper_z'], self.align_step_size, self.align_step_speed, self.align_step_accel)
            if not z1_state:
                self.manual_move(self.steppers['stepper_z1'], self.align_step_size, self.align_step_speed, self.align_step_accel)

            print_time = toolhead.get_last_move_time()
            z_state = self.z_pin.query_endstop(print_time)
            z1_state = self.z1_pin.query_endstop(print_time)

        initial_pos[2] = initial_pos[2] - self.backtrack_distance
        toolhead.move(initial_pos, self.initial_move_speed)
        gcode.respond_info('Alignment finished, Z homing required')

class ZGantryCalib:
    cmd_Z_GANTRY_CALIB_help = 'Auto Z gantry alignment'

    def __init__(self, config):
        self.printer = config.get_printer()
        self.plugin_helper = ZGantryCalibHelper(config)
        gcode = self.printer.lookup_object('gcode')

        # make your gcode recognisable
        gcode.register_command('Z_GANTRY_CALIBRATE', self.cmd_Z_GANTRY_CALIB, self.cmd_Z_GANTRY_CALIB_help)

    def register_stepper(self, config, mcu_stepper):
        self.plugin_helper.register_stepper(config, mcu_stepper)

    def cmd_Z_GANTRY_CALIB(self, gcmd):
        self.plugin_helper.do_z_calib()

def load_config(config):
    return ZGantryCalib(config)

# Sample printer.cfg entry:
[z_gantry_calibration]
z_pin: ^PC13
z1_pin: ^PC2
initial_z_height: 200 # It will lift up gantry to this position before alignment for faster operation
initial_move_speed: 10 # initial_z_height will use this speed
align_step_size: 0.5 # Single step of single stepper during alignment operation
align_step_speed: 5
align_step_accel: 100
backtrack_distance: 30 # Move gantry by this distance after alignment