# Copyright (C) 2023 Justin Schuh <code@justinschuh.com>
#
# This file may be distributed under the terms of the GNU GPLv3 license.
#
# Credit to original inspiration:
# https://gist.github.com/ChipCE/95fdbd3c2f3a064397f9610f915f7d02

[gcode_macro bed_mesh_calibrate_fast]
description: Wraps BED_MESH_CALIBRATE, scaling probe count to specified area.
  Usage: See Klipper documentation.
gcode:
  # Abort on a bad config.
  BED_MESH_CHECK ABORT=1

  # Find the real bed_mesh_calibrate command.
  {% set calibrate_cmd =
           (printer["gcode_macro list_macros"].macros.bed_mesh_calibrate|
              default(["bed_mesh_calibrate"],True))[-1] %}

  {% set km = printer["gcode_macro _km_globals"] %}
  {% set probe_mesh_padding = km.probe_mesh_padding %}
  {% set probe_min_count = km.probe_min_count %}
  {% set probe_count_scale = km.probe_count_scale %}
  {% set bed_mesh = printer.configfile.settings.bed_mesh %}

  # TODO: Handle the math for a delta bed.
  {%if "mesh_radius" not in bed_mesh and
       "MESH_RADIUS" not in params %}
    {% set safe_min_x = bed_mesh.mesh_min[0] %}
    {% set safe_min_y = bed_mesh.mesh_min[1] %}
    {% set safe_max_x = bed_mesh.mesh_max[0] %}
    {% set safe_max_y = bed_mesh.mesh_max[1] %}

    # Always bound MESH_MIN and MESH_MAX.
    {% if "MESH_MIN" in params %}
      {% set mesh_min_x = (params.MESH_MIN.split(",")[0]|float -
                           probe_mesh_padding, safe_min_x)|max %}
      {% set mesh_min_y = (params.MESH_MIN.split(",")[1]|float -
                           probe_mesh_padding, safe_min_y)|max %}
    {% else %}
      {% set mesh_min_x = safe_min_x %}
      {% set mesh_min_y = safe_min_y %}
    {% endif %}
    {% if "MESH_MAX" in params %}
      {% set mesh_max_x = (params.MESH_MAX.split(",")[0]|float +
                           probe_mesh_padding, safe_max_x)|min %}
      {% set mesh_max_y = (params.MESH_MAX.split(",")[1]|float +
                           probe_mesh_padding, safe_max_y)|min %}
    {% else %}
      {% set mesh_max_x = safe_max_x %}
      {% set mesh_max_y = safe_max_y %}
    {% endif %}

    {% set probe_count = bed_mesh.probe_count if not params.PROBE_COUNT else
                         params.PROBE_COUNT.split(",")|map('int')|list %}
    # Don't scale the probe count if one was explicitly provided.
    {% if "PROBE_COUNT" not in params %}
      {% set max_x_probes = probe_count[0] %}
      {% set max_y_probes = probe_count[-1] %}

      {% set x_probes = (max_x_probes * (mesh_max_x - mesh_min_x) /
                         (safe_max_x - safe_min_x) * probe_count_scale)
                        | round(0) | int %}
      {% set x_probes = ((x_probes, probe_min_count)|max, max_x_probes)|min %}

      {% set y_probes = (max_y_probes * (mesh_max_y - mesh_min_y ) /
                         (safe_max_y - safe_min_y) * probe_count_scale )
                        | round(0) | int %}
      {% set y_probes = ((y_probes, probe_min_count)|max, max_y_probes)|min %}
      # Add probes for bicubic if one axis has too many probes for lagrange.
      {% if x_probes > 6 and y_probes < 4 %}
        {% set y_probes = 4 %}
      {% elif y_probes > 6 and x_probes < 4 %}
        {% set x_probes = 4 %}
      {% endif %}
      {% set probe_count = [x_probes,y_probes] %}
    {% elif probe_count|length == 1 %}
      {% set dummy = probe_count.append(probe_count[0]) %}
    {% endif %}

    # If the config includes a relative_reference_index then we need to find the
    # point in the new mesh that's closest to the index point in the mesh that
    # the config would have generated.
    # TODO: Could also adjust the mesh parameters in here to ensure it includes
    # the original index point, but that would be extra work and would cause
    # slower probes if the mesh needs to be expanded to include the point.
    {% if "relative_reference_index" in bed_mesh %}
      {% set row = (bed_mesh.relative_reference_index / bed_mesh.probe_count[0]
                    )|int%}
      {% set rrf_x = (((safe_max_x - safe_min_x) /
                       (bed_mesh.probe_count[0] - 1))|round(2, 'floor')) *
                     (bed_mesh.relative_reference_index %
                      bed_mesh.probe_count[-1]) %}
      {% if row % 2 %}
        {% set rrf_x = safe_max_x - rrf_x %}
      {% else %}
        {% set rrf_x = safe_min_x + rrf_x %}
      {% endif %}
      {% set rrf_y = (((safe_max_y - safe_min_y) /
                       (bed_mesh.probe_count[-1] - 1))|round(2, 'floor')) *
                     row + safe_min_x %}
      {% set x_dist = (mesh_max_x - mesh_min_x) / (probe_count[0] - 1) %}
      {% set y_dist = (mesh_max_y - mesh_min_y) / (probe_count[1] - 1) %}
      {% set rrf = {'x':0, 'y':0, 'dist':safe_max_x**2+safe_max_y**2,'pos':0} %}
      {% for row in range(probe_count[1])%}
        {% for col in range(probe_count[0])%}
          {% if row % 2 %}
            {% set x = mesh_max_x - col * x_dist %}
          {% else %}
            {% set x = mesh_min_x + col * x_dist %}
          {% endif %}
          {% set y = mesh_min_y + row * y_dist %}
          {% set dist = ((x - rrf_x)**2 + (y - rrf_y)**2)**0.5 %}
          {% if dist < rrf.dist %}
            {% set dummy = rrf.__setitem__("dist", dist) %}
            {% set dummy = rrf.__setitem__("x", x) %}
            {% set dummy = rrf.__setitem__("y", y) %}
            {% set dummy = rrf.__setitem__("pos", row * probe_count[1] + col) %}
          {% endif %}
        {% endfor %}
      {% endfor %}
      {% if rrf.x != rrf_x or rrf.y != rrf_y %}
        {action_respond_info("relative_reference_index remapped to"
                             " %d (%.2f,%.2f) from %d (%.2f,%.2f)" %
                             (rrf.pos, rrf.x, rrf.y,
                              bed_mesh.relative_reference_index, rrf_x, rrf_y))}
      {% endif %}
      {% set dummy = params.__setitem__("RELATIVE_REFERENCE_INDEX", rrf.pos) %}
    {% endif %}

    {% set dummy = params.__setitem__("MESH_MIN", mesh_min_x~","~mesh_min_y) %}
    {% set dummy = params.__setitem__("MESH_MAX", mesh_max_x~","~mesh_max_y) %}
    {% set dummy = params.__setitem__("PROBE_COUNT", probe_count|join(',')) %}
    # Force bicubic if we've exceeded the max for lagrange.
    {% if probe_count[0] > 6 or probe_count[1]|default(0) > 6 %}
      {% set dummy = params.__setitem__("ALGORITHM", "bicubic") %}
    {% endif %}
    # Warn on bad parameters that were fixed.
    {% if "MESH_MIN" in params or "MESH_MAX" in params %}
      BED_MESH_CHECK {rawparams}
    {% endif %}
  {% else %}
    # Mesh limits may be out of bounds, so strip them from the fallback path.
    {% set dummy = params.__delitem__("MESH_MIN") %}
    {% set dummy = params.__delitem__("MESH_MAX") %}
  {% endif %}
  # Abort on bad parameters.
  {% if "MESH_MIN" in params or "MESH_MAX" in params %}
    BED_MESH_CHECK ABORT=1{%for k in params%}{' '~k~'="'~params[k]~'"'}{%
                            endfor%}
  {% endif %}
  _KM_PRINT_STATUS ACTION=PUSH_STATUS
  _KM_PRINT_STATUS ACTION=CHANGE STATUS=meshing
  {calibrate_cmd}{%for k in params%}{' '~k~'="'~params[k]~'"'}{%endfor%}
  _KM_PRINT_STATUS ACTION=CHANGE STATUS=pop_status

[gcode_macro bed_mesh_check]
description: Warns if bed_mesh config may generate an invalid mesh.
  Usage: BED_MESH_CHECK [MESH_MIN=<x,y>] [MESH_MAX=<x,y>] [ABORT=<0|1>]
gcode:
  {% if printer.bed_mesh is defined %}
    {% set action = action_respond_info if params.ABORT|default(0)|int == 0 else
                    action_raise_error %}
    {% set settings = printer.configfile.settings %}
    {% set x_min = settings.stepper_x.position_min %}
    {% set y_min = settings.stepper_y.position_min %}
    {% set x_max = settings.stepper_x.position_max %}
    {% set y_max = settings.stepper_y.position_max %}

    {% set label = "[bed_mesh] config" %}
    {% if "MESH_MIN" in params %}
      {% set label = "BED_MESH_CALIBRATE params" %}
      {% set mesh_min_x = params.MESH_MIN.split(",")[0]|float %}
      {% set mesh_min_y = params.MESH_MIN.split(",")[1]|float %}
    {% else %}
      {% set mesh_min_x = settings.bed_mesh.mesh_min[0] %}
      {% set mesh_min_y = settings.bed_mesh.mesh_min[1] %}
    {% endif %}
    {% if "MESH_MAX" in params %}
      {% set label = "BED_MESH_CALIBRATE params" %}
      {% set mesh_max_x = params.MESH_MAX.split(",")[0]|float %}
      {% set mesh_max_y = params.MESH_MAX.split(",")[1]|float %}
    {% else %}
      {% set mesh_max_x = settings.bed_mesh.mesh_max[0] %}
      {% set mesh_max_y = settings.bed_mesh.mesh_max[1] %}
    {% endif %}

    {% if "bltouch" in settings %}
        {% set x_offset = settings.bltouch.x_offset %}
        {% set y_offset = settings.bltouch.y_offset %}
        {% set probe = "bltouch" %}
    {% elif "probe" in settings %}
        {% set x_offset = settings.probe.x_offset %}
        {% set y_offset = settings.probe.y_offset %}
        {% set probe = "probe" %}
    {% else %}
        {% set x_offset = 0.0 %}
        {% set y_offset = 0.0 %}
    {% endif %}

    {% set output = [] %}
    {% set warn =
         "* mesh_%s (%f, %f) adjusted by " ~ probe ~
         ".%s_offset (%f) can move out of range for "
         "stepper_%s.position_%s (%f)." %}
    {% if x_offset > 0 and (mesh_min_x - x_offset) < x_min %}
      {% set dummy = output.append(warn % ('min', mesh_min_x, mesh_min_y,
                                           'x', x_offset, 'x', 'min', x_min)) %}
    {% elif x_offset < 0 and (mesh_max_x - x_offset) > x_max %}
      {% set dummy = output.append(warn % ('max', mesh_max_x, mesh_max_y,
                                           'x', x_offset, 'x', 'max', x_max)) %}
    {% endif %}
    {% if y_offset > 0 and (mesh_min_y - y_offset) < y_min %}
      {% set dummy = output.append(warn % ('min', mesh_min_x, mesh_min_y,
                                           'y', y_offset, 'y', 'min', y_min)) %}
    {% elif y_offset < 0 and (mesh_max_y - y_offset) > y_max %}
      {% set dummy = output.append(warn % ('max', mesh_max_x, mesh_max_y,
                                           'y', y_offset, 'y', 'max', y_max)) %}
    {% endif %}

    {% if output %}
      { action(
        "Warning: The following issue(s) were detected in your " ~ label ~
        ":\n" ~ output|join("\n")) }
    {% endif %}
  {% endif %}