# Copyright (C) 2022 Justin Schuh <code@justinschuh.com>
#
# This file may be distributed under the terms of the GNU GPLv3 license.
#
# Credit to original inspiration:
# https://klipper.discourse.group/t/saving-and-adjusting-per-build-surface-z-offsets/696 

[gcode_macro _apply_bed_surface_offset]
gcode:
  {% set surfaces = printer.save_variables.variables.bed_surfaces %}
  {% if surfaces.active %}
    SET_SURFACE_ACTIVE SURFACE={surfaces.active}
  {% endif %}

[gcode_macro _init_surfaces]
gcode:
  {% set km = printer["gcode_macro _km_globals"] %}
  {% if "bed_surfaces" in printer.save_variables.variables %}
    {% set old_surfaces = printer.save_variables.variables.bed_surfaces %}
  {% else %}
    {% set old_surfaces = { 'active' : '', 'available' : {} } %}
  {% endif %}
  {% set settings = printer.configfile.settings %}
  {% set new_probe_z = (settings.probe | default(settings.bltouch) |
                        default(settings.smart_effector) | default({})
                        ).z_offset|default(0.0)|float %}
  {% set new_endstop_z = (settings.stepper_z | default({})).position_endstop |
                        default(0.0)|float %}
  {% if 'endstop_z' not in old_surfaces %}
    {% set dummy = old_surfaces.__setitem__('endstop_z', new_endstop_z) %}
  {% endif %}
  {% if 'probe_z' not in old_surfaces %}
    {% set dummy = old_surfaces.__setitem__('probe_z', new_probe_z) %}
  {% endif %}
  {% set surfaces = { 'active' : '', 'available' : {},
                      'endstop_z' : old_surfaces.endstop_z,
                      'probe_z' : old_surfaces.probe_z } %}
  {% for s in km.bed_surfaces %}
    {% set s = s.split()|join(' ')|lower %}
    {% if s|length > km.bed_surface_max_name_length or
          s|list|select("in", " \r\n\"\'")|list %}
      {action_raise_error('Invalid surface name "%s". Name must be %d or fewer '
        'characters and must not include space or quotation characters'
        | format(s, km.bed_surface_max_name_length))}
    {% endif %}
    {% if s in old_surfaces.available %}
      {% set dummy = surfaces.available.__setitem__(s,
                       old_surfaces.available[s]) %}
    {% else %}
      {% set dummy = surfaces.available.__setitem__(s, {'offset' : 0.0}) %}
    {% endif %}
  {% endfor %}
  {% if old_surfaces.active in surfaces.available %}
    {% set dummy = surfaces.__setitem__('active', old_surfaces.active) %}
  {% elif km.bed_surfaces %}
    {% set dummy = surfaces.__setitem__('active', km.bed_surfaces[0]|lower) %}
  {% endif %}
  SAVE_VARIABLE VARIABLE=bed_surfaces VALUE="{surfaces}"
  _APPLY_BED_SURFACE_OFFSET
  {% if new_probe_z != surfaces.probe_z or
        new_endstop_z != surfaces.endstop_z %}
    { action_respond_info(
      'Z probe offset or endstop position changed. Run ADJUST_SURFACE_OFFSETS '
      'to adjust the offset for all saved surfaces by the change differential, '
      'or run ADJUST_SURFACE_OFFSETS IGNORE=1 to hide this message without '
      'making changes.') }
  {% endif %}

[gcode_macro adjust_surface_offsets]
description: Adjusts surface offsets to account for changes in the Z endstop
  position or probe Z offset.
  Usage: ADJUST_SURFACE_OFFSETS [IGNORE]
gcode:
  {% set surfaces = printer.save_variables.variables.bed_surfaces %}
  {% set settings = printer.configfile.settings %}
  {% set new_probe_z = (settings.probe | default(settings.bltouch) |
                        default(settings.smart_effector) | default({})
                        ).z_offset|default(0.0)|float %}
  {% set new_endstop_z = (settings.stepper_z | default({})).position_endstop |
                          default(0.0)|float %}
  {% set diff = (surfaces.probe_z - new_probe_z +
                 surfaces.endstop_z - new_endstop_z)|round(6) %}
  {% if not params.IGNORE|default(0)|int %}
    {% for s in surfaces.available %}
      {% set offset = (surfaces.available[s].offset - diff)|round(6) %}
      {% set dummy = surfaces.available.__setitem__(s, {'offset' : offset}) %}
    {% endfor %}
    { action_respond_info("All bed surfaces now adjusted by %1.4f"|
                          format(diff))}
  {% elif diff != 0 %}
    { action_respond_info("Status cleared without adjustment") }
  {% endif %}
  {% set dummy = surfaces.__setitem__('endstop_z', new_endstop_z| round(6)) %}
  {% set dummy = surfaces.__setitem__('probe_z', new_probe_z|round(6)) %}
  SAVE_VARIABLE VARIABLE=bed_surfaces VALUE="{surfaces}"

[gcode_macro set_surface_offset]
description: Sets the offset for a surface and moves the toolhead (if homed).
  Usage: SET_SURFACE_OFFSET [OFFSET=<offset>] [SURFACE=<surface>]
gcode:
  {% set surfaces = printer.save_variables.variables.bed_surfaces %}
  {% set SURFACE = params.SURFACE|default(surfaces.active)|lower %}
  {% if SURFACE not in surfaces.available %}
    { action_raise_error("Bed surface %s does not exist." | format(SURFACE)) }
  {% endif %}
  {% set active = surfaces.available[SURFACE] %}
  # If no offset is provided just print out the current offset.
  {% set OFFSET = params.OFFSET|default(active.offset)|float %}
  {% if OFFSET != active.offset %}
    {% set dummy = surfaces.available[SURFACE].__setitem__("offset", OFFSET) %}
    SAVE_VARIABLE VARIABLE=bed_surfaces VALUE="{surfaces}"
    {% if SURFACE == surfaces.active %}
      _km_set_gcode_offset_base Z="{OFFSET}" MOVE={
        1 if printer.toolhead.homed_axes == 'xyz' else 0}
    {% endif %}
  {% endif %}
  { action_respond_info("Bed surface: %s Offset: %.3f"
                        | format(SURFACE, OFFSET)) }
  # Dummy argument block for Mainsail
  {% set dummy = None if True else "
  {% set dummy = params.SURFACE|default(active surface) %}
  {% set dummy = params.OFFSET|default(none)|float %}
  " %} # End argument block for Mainsail

[gcode_macro set_surface_active]
description: Sets the active bed surface and moves the toolhead (if homed). If
  no SURFACE argument is present the available surfaces are listed and the
  active one is preceded by a "*".
  Usage: SET_SURFACE_ACTIVE [SURFACE=<surface>]
gcode:
  {% set surfaces = printer.save_variables.variables.bed_surfaces %}
  {% if "SURFACE" in params %}
    {% set SURFACE = params.SURFACE|lower %}
    {% if SURFACE not in surfaces.available %}
      { action_raise_error("Bed surface %s does not exist." | format(SURFACE)) }
    {% endif %}
    {% if SURFACE != surfaces.active %}
      {% set dummy = surfaces.__setitem__("active", SURFACE) %}
      SAVE_VARIABLE VARIABLE=bed_surfaces VALUE="{surfaces}"
    {% endif %}
    {% if surfaces.available[SURFACE].offset !=
          printer.gcode_move.homing_origin.z %}
      _km_set_gcode_offset_base Z="{surfaces.available[SURFACE].offset
        }" MOVE={1 if printer.toolhead.homed_axes == 'xyz' else 0}
    {% endif %}
    {action_respond_info("Active bed surface: %s; offset: %.3f"
                        | format(SURFACE, surfaces.available[SURFACE].offset))}
  {% else %}
    {% set output = [] %}
    {% for s in surfaces.available|list|sort %}
      {% set dummy = output.append("%s&nbsp;%s - offset: %.3f" 
                      | format("*" if s == surfaces.active else "&nbsp;",
                                s, surfaces.available[s].offset)) %}
    {% endfor %}
    {action_respond_info(output|join('\n'))}
  {% endif %}
  # Dummy argument block for Mainsail
  {% set dummy = None if True else "
  {% set dummy = params.SURFACE|default(none) %}
  " %} # End argument block for Mainsail

[gcode_macro set_gcode_offset]
description: Wraps SET_GCODE_OFFSET to update the current bed sheet offset.
  Usage: SET_GCODE_OFFSET [X=<pos>|X_ADJUST=<adjust>]
                          [Y=<pos>|Y_ADJUST=<adjust>]
                          [Z=<pos>|Z_ADJUST=<adjust>]
                          [MOVE=1 [MOVE_SPEED=<speed>]]
rename_existing: _KM_SET_GCODE_OFFSET_BASE
gcode:
  {% set surfaces = printer.save_variables.variables.bed_surfaces %}
  {% if surfaces.active and
        not printer["gcode_macro _km_save_state"].is_ephemeral %}
    {% set Z = params.Z|default(0.0)|float|round(6) %}
    {% set Z_ADJUST = params.Z_ADJUST|default(0.0)|float %}
    {% if 'Z' in params and
                 Z != surfaces.available[surfaces.active].offset %}
      {% set dummy = surfaces.available[surfaces.active].__setitem__("offset",
                       Z) %}
      SAVE_VARIABLE VARIABLE=bed_surfaces VALUE="{surfaces}"
    {% elif Z_ADJUST != 0.0 %}
      {% set dummy = surfaces.available[surfaces.active].__setitem__(
        "offset", (Z_ADJUST + printer.gcode_move.homing_origin.z)|round(6)) %}
      SAVE_VARIABLE VARIABLE=bed_surfaces VALUE="{surfaces}"
    {% endif %}
  {% endif %}
  _km_set_gcode_offset_base{% for k in params%}{' '~k~'="'~params[k]~'"'
    }{% endfor %}

[gcode_macro make_surface_mesh]
description: Generates and saves a mesh to automatically load in PRINT_START.
  Usage: MAKE_SURFACE_MESH [SURFACE=<surface>] [BED=<probing temperature>]
                           [EXTRUDER=<probing temperature>]
                           [MESH_MULTIPLIER=<integer multiple of mesh density>]
                           [bed_mesh_calibrate parameters]
gcode:
  {% set surfaces = printer.save_variables.variables.bed_surfaces %}
  {% set SURFACE = params.SURFACE|default(surfaces.active)|lower %}
  {% if SURFACE not in surfaces.available %}
    { action_raise_error("Bed surface %s does not exist." | format(SURFACE)) }
  {% endif %}
  {% set dummy = params.__setitem__('PROFILE', SURFACE) %}
  {% set BED = params.BED|default(70) | int %}
  {% set km = printer["gcode_macro _km_globals"] %}
  {% set EXTRUDER = params.EXTRUDER|default(km.start_extruder_probing_temp) |
                    int %}
  {% set MESH_MULTIPLIER = (params.MESH_MULTIPLIER|default(2)|int, 1)|max %}

  M104 S{EXTRUDER}
  M140 S{BED}
  G28

  # Adjust offset before running mesh
  {% if surfaces.available[SURFACE].offset !=
        printer.gcode_move.homing_origin.z %}
    _km_set_gcode_offset_base Z="{surfaces.available[SURFACE].offset
      }" MOVE=1
  {% endif %}

  # If no offset is provided just print out the current offset.
  {% if BED > 0 %}
    {action_respond_info("Stabilizing bed at %iC" | format(BED,))}
    _KM_PARK_IF_NEEDED HEATER=heater_bed RANGE=0.5
    {% if BED < (printer.heater_bed.temperature - 0.2) %}
      M190 R{BED}
      G4 P{((km.start_bed_heat_delay * 10, 30000)|min,
            km.start_bed_heat_delay)|max}
    {% else %}
      G4 P{km.start_bed_heat_delay}
    {% endif %}
    M190 R{BED}
  {% endif %}

  {% if EXTRUDER > 0 %}
    _KM_PARK_IF_NEEDED HEATER={printer.toolhead.extruder} RANGE=2
    M109 R{km.start_extruder_probing_temp}
  {% endif %}

  {% if km.start_home_z_at_temp and not bed_at_target %}
    G28 Z # Re-home only the Z axis now that the bed has stabilized.
  {% endif %}

  # Scale the mesh grid size while preserving the original points and relative
  # reference index.
  {% set probe_count = printer.configfile.settings.bed_mesh.probe_count
                         if not params.PROBE_COUNT else
                       params.PROBE_COUNT.split(",")|map('int')|list %}
  {% set dummy = params.__setitem__('PROBE_COUNT',(
                   MESH_MULTIPLIER * (probe_count[0] - 1) + 1,
                   MESH_MULTIPLIER * (probe_count[-1] - 1) + 1,)|join(',')) %}

  BED_MESH_CLEAR
  BED_MESH_CALIBRATE_FAST {%
    for k in params|reject('in',['BED','SURFACE'])|list %}{
      ' ' ~ k ~ '=' ~ params[k] }{% endfor %}
  SAVE_CONFIG

  # Dummy argument block for Mainsail
  {% set dummy = None if True else "
  {% set dummy = params.MESH_MULTIPLIER|default(2) %}
  {% set dummy = params.SURFACE|default(active surface) %}
  " %} # End argument block for Mainsail

[gcode_macro load_surface_mesh]
description: Attempts to load a mesh associated with the specified surface.
  Usage: LOAD_SURFACE_MESH [SURFACE=<surface>]
gcode:
  {% set surfaces = printer.save_variables.variables.bed_surfaces %}
  {% set SURFACE = params.SURFACE|default(surfaces.active)|lower %}
  {% if SURFACE != surfaces.active %}
    SET_SURFACE_ACTIVE SURFACE={SURFACE}
  {% endif %}

  {% if SURFACE in printer.bed_mesh.profiles %}
    {% set mesh = printer.bed_mesh.profiles[SURFACE].mesh_params %}
    {% set default = printer.configfile.settings.bed_mesh %}
    # Ensure the saved mesh has at least the resolution of the default.
    {% if mesh.min_x <= (default.mesh_min[0] + 0.5) and
          mesh.min_y <= (default.mesh_min[1] + 0.5) and
          mesh.max_x >= (default.mesh_max[0] - 0.5) and
          mesh.max_y >= (default.mesh_max[1] - 0.5) and
          mesh.x_count >= default.probe_count[0] and
          mesh.y_count >= default.probe_count[-1] %}
      # Skip the bed level if we have a good profile.
      BED_MESH_PROFILE LOAD={SURFACE}
      {action_respond_info("Loaded mesh: %s" | format(SURFACE,))}
    {% endif %}
  {% endif %}