# Measure lost XY steps on a CoreXY printer by re-homing.
#
# Homing resets Klipper's idea of the position, so the toolhead position
# after G28 is always the endstop position and cannot show lost steps. The
# stepper drivers' step counters (MCU positions) are not reset by homing:
# when no steps were lost, the endstop triggers at the same step count as
# last time; every lost step shows up as a difference. This module records
# the step counts at each X and Y endstop trigger and reports the drift in
# mm against a stored reference.
#
# Install: copy to ~/klipper/klippy/extras/xy_drift.py and add [xy_drift]
#
# Commands:
#   XY_DRIFT_STORE   use the last X and Y homing as the reference
#   XY_DRIFT_QUERY   print the drift of the last homing against the reference
#
# Status (printer.xy_drift):
#   ref_set          a reference is stored (cleared by a Klipper restart,
#                    which also resets the step counters)
#   homed_x, homed_y the axis was homed since the reference was stored
#   x, y             drift in mm of the last homing against the reference
#                    (positive = Klipper had moved further than the head did)

import logging


class XYDrift:
    def __init__(self, config):
        self.printer = config.get_printer()
        self.sx = self.sy = None
        self.last = {}          # axis -> kinematic position (mm) at the trigger
        self.ref = {}
        self.since_ref = set()
        self.printer.register_event_handler('klippy:connect',
                                            self._handle_connect)
        self.printer.register_event_handler('homing:home_rails_end',
                                            self._handle_home_rails_end)
        gcode = self.printer.lookup_object('gcode')
        gcode.register_command('XY_DRIFT_STORE', self.cmd_STORE,
                               desc="Store the last XY homing as drift reference")
        gcode.register_command('XY_DRIFT_QUERY', self.cmd_QUERY,
                               desc="Report XY drift against the reference")

    def _handle_connect(self):
        kin = self.printer.lookup_object('toolhead').get_kinematics()
        steppers = {s.get_name(): s for s in kin.get_steppers()}
        self.sx = steppers.get('stepper_x')
        self.sy = steppers.get('stepper_y')
        if self.sx is None or self.sy is None:
            raise self.printer.config_error(
                "xy_drift needs stepper_x and stepper_y (CoreXY)")

    def _handle_home_rails_end(self, homing_state, rails):
        # This event also comes from other modules (Beacon sends it with no
        # rails and a state that raises on get_trigger_position), so only
        # handle X/Y homing and never let an error escape: an exception
        # here shuts Klipper down.
        try:
            names = [rail.get_name() for rail in rails]
            axes = [axis for axis, st in (('x', self.sx), ('y', self.sy))
                    if st is not None and st.get_name() in names]
            if not axes:
                return
            a = (homing_state.get_trigger_position(self.sx.get_name())
                 * self.sx.get_step_dist())
            b = (homing_state.get_trigger_position(self.sy.get_name())
                 * self.sy.get_step_dist())
        except Exception:
            logging.exception("xy_drift: could not read trigger positions")
            return
        # CoreXY: x = (a + b) / 2, y = (a - b) / 2
        pos = {'x': 0.5 * (a + b), 'y': 0.5 * (a - b)}
        for axis in axes:
            self.last[axis] = pos[axis]
            self.since_ref.add(axis)

    def _drift(self, axis):
        if axis in self.ref and axis in self.last:
            return round(self.last[axis] - self.ref[axis], 4)
        return 0.

    def cmd_STORE(self, gcmd):
        if 'x' not in self.last or 'y' not in self.last:
            raise gcmd.error("XY_DRIFT_STORE: home X and Y first")
        self.ref = dict(self.last)
        self.since_ref = set()
        gcmd.respond_info("XY_DRIFT: reference stored")

    def cmd_QUERY(self, gcmd):
        if not self.ref:
            gcmd.respond_info("XY_DRIFT: no reference stored")
            return
        gcmd.respond_info("XY_DRIFT: X=%.3f mm Y=%.3f mm" % (
            self._drift('x'), self._drift('y')))

    def get_status(self, eventtime):
        return {'ref_set': bool(self.ref),
                'homed_x': 'x' in self.since_ref,
                'homed_y': 'y' in self.since_ref,
                'x': self._drift('x'), 'y': self._drift('y')}


def load_config(config):
    return XYDrift(config)