From df35aa554748aab041f1c017297da8b01858a176 Mon Sep 17 00:00:00 2001
From: Kuba Birecki <kuba.birecki@crystalplanet.io>
Date: Sun, 12 Dec 2021 08:36:56 +0100
Subject: [PATCH] Implement helpers for key pixel position and pixel distance

---
 app/Kconfig                         |  4 +++
 app/dts/bindings/zmk,animation.yaml |  9 +++++++
 app/include/zmk/animation.h         | 12 +++++++++
 app/src/animation/animation.c       | 38 +++++++++++++++++++++++++++++
 4 files changed, 63 insertions(+)

diff --git a/app/Kconfig b/app/Kconfig
index 9762d6c8..354c0b36 100644
--- a/app/Kconfig
+++ b/app/Kconfig
@@ -312,6 +312,10 @@ config ZMK_ANIMATION_FPS
 	range 0 60
 	default 30
 
+config ZMK_ANIMATION_PIXEL_DISTANCE
+	bool "Generate a lookup table for distances between pixels"
+	default y
+
 #ZMK_ANIMATION
 endif
 
diff --git a/app/dts/bindings/zmk,animation.yaml b/app/dts/bindings/zmk,animation.yaml
index 664de459..a2ad35b7 100644
--- a/app/dts/bindings/zmk,animation.yaml
+++ b/app/dts/bindings/zmk,animation.yaml
@@ -26,3 +26,12 @@ properties:
       This field contains the pixel configuration for the entire board.
       The order of this array determines in what order pixels are sent to the driver device API.
       If multiple driving devices are used, their chain-length property determines the size of the buffer for each device.
+
+  key-pixels:
+    type: array
+    required: false
+    description: |
+      Use this field to specify the pixel index corresponding to each key
+      following the order used in your keymap.
+      When left unspecified, the driver assumes that for every key, the pixel has a matching id.
+      So for N keys, the first N pixels are exactly in the same order as keys in your keymap.
diff --git a/app/include/zmk/animation.h b/app/include/zmk/animation.h
index a2dae1c1..aeede8ea 100644
--- a/app/include/zmk/animation.h
+++ b/app/include/zmk/animation.h
@@ -21,6 +21,18 @@ struct zmk_color_hsl {
     uint8_t l;
 };
 
+#if DT_NODE_HAS_PROP(DT_INST(0, animation), key_position)
+size_t zmk_animation_get_pixel_by_key_position(size_t key_position);
+#else
+static inline size_t zmk_animation_get_pixel_by_key_position(size_t key_position) {
+    return key_position;
+}
+#endif
+
+#if defined(CONFIG_ZMK_ANIMATION_PIXEL_DISTANCE) && (CONFIG_ZMK_ANIMATION_PIXEL_DISTANCE == 1)
+uint16_t zmk_animation_get_pixel_distance(size_t pixel_idx, size_t other_pixel_idx);
+#endif
+
 /**
  * Converts color from HSL to RGB.
  *
diff --git a/app/src/animation/animation.c b/app/src/animation/animation.c
index e4c3df58..7fcece19 100644
--- a/app/src/animation/animation.c
+++ b/app/src/animation/animation.c
@@ -11,6 +11,9 @@
 #include <init.h>
 #include <kernel.h>
 
+#include <stdlib.h>
+#include <math.h>
+
 #include <logging/log.h>
 
 #include <drivers/animation.h>
@@ -82,6 +85,31 @@ static const size_t pixels_size = DT_INST_PROP_LEN(0, pixels);
  */
 static struct led_rgb px_buffer[DT_INST_PROP_LEN(0, pixels)];
 
+/**
+ * Conditional implementation of zmk_animation_get_pixel_by_key_position
+ * if key-pixels is set.
+ */
+#if DT_INST_NODE_HAS_PROP(0, key_position)
+static const uint8_t pixels_by_key_position[] = DT_INST_PROP(0, key_pixels);
+
+size_t zmk_animation_get_pixel_by_key_position(size_t key_position) {
+    return pixels_by_key_position[key_position];
+}
+#endif
+
+#if defined(CONFIG_ZMK_ANIMATION_PIXEL_DISTANCE) && (CONFIG_ZMK_ANIMATION_PIXEL_DISTANCE == 1)
+
+/**
+ * Lookup table for distance between any two pixels.
+ */
+static uint16_t pixel_distance[DT_INST_PROP_LEN(0, pixels)][DT_INST_PROP_LEN(0, pixels)];
+
+uint16_t zmk_animation_get_pixel_distance(size_t pixel_idx, size_t other_pixel_idx) {
+    return pixel_distance[pixel_idx][other_pixel_idx];
+}
+
+#endif
+
 static void zmk_animation_tick(struct k_work *work) {
     for (size_t i = 0; i < animations_size; ++i) {
         animation_on_before_frame(animations[i]);
@@ -119,6 +147,16 @@ static void zmk_animation_tick_handler(struct k_timer *timer) { k_work_submit(&a
 K_TIMER_DEFINE(animation_tick, zmk_animation_tick_handler, NULL);
 
 static int zmk_animation_init(const struct device *dev) {
+#if defined(CONFIG_ZMK_ANIMATION_PIXEL_DISTANCE) && (CONFIG_ZMK_ANIMATION_PIXEL_DISTANCE == 1)
+    // Prefill the pixel distance lookup table
+    for (size_t i = 0; i < pixels_size; ++i) {
+        for (size_t j = 0; j < pixels_size; ++j) {
+            pixel_distance[i][j] = sqrt(pow(pixels[i].position_x - pixels[j].position_x, 2) +
+                                        pow(pixels[i].position_y - pixels[j].position_y, 2));
+        }
+    }
+#endif
+
     LOG_INF("ZMK Animation Ready");
 
     k_timer_start(&animation_tick, K_NO_WAIT, K_MSEC(1000 / CONFIG_ZMK_ANIMATION_FPS));