Fix RGB heatmap to use XY positions and use correct led limits (#17184)

* Fix RGB heatmap to use XY positions

* lower effect area limit and make configurable

* tidy up macro

* Fix triggering in both directions.

* add docs

* fix bug when decreasing value

* performance tweak

3 files changed, 67 insertions, 45 deletions

diff --git a/docs/feature_rgb_matrix.md b/docs/feature_rgb_matrix.md
index 2e88ae4727..9eaf63d710 100644
--- a/docs/feature_rgb_matrix.md
+++ b/docs/feature_rgb_matrix.md
@@ -667,7 +667,22 @@ In order to change the delay of temperature decrease define `RGB_MATRIX_TYPING_H
 #define RGB_MATRIX_TYPING_HEATMAP_DECREASE_DELAY_MS 50
 ```
 
-Heatmap effect may not light up the correct adjacent LEDs for certain key matrix layout such as split keyboards. The following define will limit the effect to pressed keys only:
+As heatmap uses the physical position of the leds set in the g_led_config, you may need to tweak the following options to get the best effect for your keyboard. Note the size of this grid is `224x64`.
+
+Limit the distance the effect spreads to surrounding keys. 
+
+```c
+#define RGB_MATRIX_TYPING_HEATMAP_SPREAD 40
+```
+
+Limit how hot surrounding keys get from each press.
+
+```c
+#define RGB_MATRIX_TYPING_HEATMAP_AREA_LIMIT 16
+```
+
+Remove the spread effect entirely.
+
 ```c
 #define RGB_MATRIX_TYPING_HEATMAP_SLIM
 ```
diff --git a/quantum/rgb_matrix/animations/typing_heatmap_anim.h b/quantum/rgb_matrix/animations/typing_heatmap_anim.h
index 4b17c4c3ed..cfc3fc015b 100644
--- a/quantum/rgb_matrix/animations/typing_heatmap_anim.h
+++ b/quantum/rgb_matrix/animations/typing_heatmap_anim.h
@@ -6,30 +6,35 @@ RGB_MATRIX_EFFECT(TYPING_HEATMAP)
 #            define RGB_MATRIX_TYPING_HEATMAP_DECREASE_DELAY_MS 25
 #        endif
 
+#        ifndef RGB_MATRIX_TYPING_HEATMAP_SPREAD
+#            define RGB_MATRIX_TYPING_HEATMAP_SPREAD 40
+#        endif
+
+#        ifndef RGB_MATRIX_TYPING_HEATMAP_AREA_LIMIT
+#            define RGB_MATRIX_TYPING_HEATMAP_AREA_LIMIT 16
+#        endif
 void process_rgb_matrix_typing_heatmap(uint8_t row, uint8_t col) {
 #        ifdef RGB_MATRIX_TYPING_HEATMAP_SLIM
     // Limit effect to pressed keys
     g_rgb_frame_buffer[row][col] = qadd8(g_rgb_frame_buffer[row][col], 32);
 #        else
-    uint8_t m_row = row - 1;
-    uint8_t p_row = row + 1;
-    uint8_t m_col = col - 1;
-    uint8_t p_col = col + 1;
-
-    if (m_col < col) g_rgb_frame_buffer[row][m_col] = qadd8(g_rgb_frame_buffer[row][m_col], 16);
-    g_rgb_frame_buffer[row][col] = qadd8(g_rgb_frame_buffer[row][col], 32);
-    if (p_col < MATRIX_COLS) g_rgb_frame_buffer[row][p_col] = qadd8(g_rgb_frame_buffer[row][p_col], 16);
-
-    if (p_row < MATRIX_ROWS) {
-        if (m_col < col) g_rgb_frame_buffer[p_row][m_col] = qadd8(g_rgb_frame_buffer[p_row][m_col], 13);
-        g_rgb_frame_buffer[p_row][col] = qadd8(g_rgb_frame_buffer[p_row][col], 16);
-        if (p_col < MATRIX_COLS) g_rgb_frame_buffer[p_row][p_col] = qadd8(g_rgb_frame_buffer[p_row][p_col], 13);
-    }
-
-    if (m_row < row) {
-        if (m_col < col) g_rgb_frame_buffer[m_row][m_col] = qadd8(g_rgb_frame_buffer[m_row][m_col], 13);
-        g_rgb_frame_buffer[m_row][col] = qadd8(g_rgb_frame_buffer[m_row][col], 16);
-        if (p_col < MATRIX_COLS) g_rgb_frame_buffer[m_row][p_col] = qadd8(g_rgb_frame_buffer[m_row][p_col], 13);
+    for (uint8_t i_row = 0; i_row < MATRIX_ROWS; i_row++) {
+        for (uint8_t i_col = 0; i_col < MATRIX_COLS; i_col++) {
+            if (i_row == row && i_col == col) {
+                g_rgb_frame_buffer[row][col] = qadd8(g_rgb_frame_buffer[row][col], 32);
+            } else {
+#            define LED_DISTANCE(led_a, led_b) sqrt16(((int8_t)(led_a.x - led_b.x) * (int8_t)(led_a.x - led_b.x)) + ((int8_t)(led_a.y - led_b.y) * (int8_t)(led_a.y - led_b.y)))
+                uint8_t distance = LED_DISTANCE(g_led_config.point[g_led_config.matrix_co[row][col]], g_led_config.point[g_led_config.matrix_co[i_row][i_col]]);
+#            undef LED_DISTANCE
+                if (distance <= RGB_MATRIX_TYPING_HEATMAP_SPREAD) {
+                    uint8_t amount = qsub8(RGB_MATRIX_TYPING_HEATMAP_SPREAD, distance);
+                    if (amount > RGB_MATRIX_TYPING_HEATMAP_AREA_LIMIT) {
+                        amount = RGB_MATRIX_TYPING_HEATMAP_AREA_LIMIT;
+                    }
+                    g_rgb_frame_buffer[i_row][i_col] = qadd8(g_rgb_frame_buffer[i_row][i_col], amount);
+                }
+            }
+        }
     }
 #        endif
 }
@@ -40,10 +45,7 @@ static uint16_t heatmap_decrease_timer;
 static bool decrease_heatmap_values;
 
 bool TYPING_HEATMAP(effect_params_t* params) {
-    // Modified version of RGB_MATRIX_USE_LIMITS to work off of matrix row / col size
-    uint8_t led_min = RGB_MATRIX_LED_PROCESS_LIMIT * params->iter;
-    uint8_t led_max = led_min + RGB_MATRIX_LED_PROCESS_LIMIT;
-    if (led_max > sizeof(g_rgb_frame_buffer)) led_max = sizeof(g_rgb_frame_buffer);
+    RGB_MATRIX_USE_LIMITS(led_min, led_max);
 
     if (params->init) {
         rgb_matrix_set_color_all(0, 0, 0);
@@ -63,28 +65,26 @@ bool TYPING_HEATMAP(effect_params_t* params) {
     }
 
     // Render heatmap & decrease
-    for (int i = led_min; i < led_max; i++) {
-        uint8_t row = i % MATRIX_ROWS;
-        uint8_t col = i / MATRIX_ROWS;
-        uint8_t val = g_rgb_frame_buffer[row][col];
-
-        // set the pixel colour
-        uint8_t led[LED_HITS_TO_REMEMBER];
-        uint8_t led_count = rgb_matrix_map_row_column_to_led(row, col, led);
-        for (uint8_t j = 0; j < led_count; ++j) {
-            if (!HAS_ANY_FLAGS(g_led_config.flags[led[j]], params->flags)) continue;
-
-            HSV hsv = {170 - qsub8(val, 85), rgb_matrix_config.hsv.s, scale8((qadd8(170, val) - 170) * 3, rgb_matrix_config.hsv.v)};
-            RGB rgb = rgb_matrix_hsv_to_rgb(hsv);
-            rgb_matrix_set_color(led[j], rgb.r, rgb.g, rgb.b);
-        }
-
-        if (decrease_heatmap_values) {
-            g_rgb_frame_buffer[row][col] = qsub8(val, 1);
+    uint8_t count = 0;
+    for (uint8_t row = 0; row < MATRIX_ROWS && count < RGB_MATRIX_LED_PROCESS_LIMIT; row++) {
+        for (uint8_t col = 0; col < MATRIX_COLS && RGB_MATRIX_LED_PROCESS_LIMIT; col++) {
+            if (g_led_config.matrix_co[row][col] >= led_min && g_led_config.matrix_co[row][col] < led_max) {
+                count++;
+                uint8_t val = g_rgb_frame_buffer[row][col];
+                if (!HAS_ANY_FLAGS(g_led_config.flags[g_led_config.matrix_co[row][col]], params->flags)) continue;
+
+                HSV hsv = {170 - qsub8(val, 85), rgb_matrix_config.hsv.s, scale8((qadd8(170, val) - 170) * 3, rgb_matrix_config.hsv.v)};
+                RGB rgb = rgb_matrix_hsv_to_rgb(hsv);
+                rgb_matrix_set_color(g_led_config.matrix_co[row][col], rgb.r, rgb.g, rgb.b);
+
+                if (decrease_heatmap_values) {
+                    g_rgb_frame_buffer[row][col] = qsub8(val, 1);
+                }
+            }
         }
     }
 
-    return led_max < sizeof(g_rgb_frame_buffer);
+    return rgb_matrix_check_finished_leds(led_max);
 }
 
 #    endif // RGB_MATRIX_CUSTOM_EFFECT_IMPLS
diff --git a/quantum/rgb_matrix/rgb_matrix.c b/quantum/rgb_matrix/rgb_matrix.c
index f721dfc7f2..a51e379025 100644
--- a/quantum/rgb_matrix/rgb_matrix.c
+++ b/quantum/rgb_matrix/rgb_matrix.c
@@ -249,8 +249,15 @@ void process_rgb_matrix(uint8_t row, uint8_t col, bool pressed) {
 #endif // RGB_MATRIX_KEYREACTIVE_ENABLED
 
 #if defined(RGB_MATRIX_FRAMEBUFFER_EFFECTS) && defined(ENABLE_RGB_MATRIX_TYPING_HEATMAP)
-    if (rgb_matrix_config.mode == RGB_MATRIX_TYPING_HEATMAP) {
-        process_rgb_matrix_typing_heatmap(row, col);
+#    if defined(RGB_MATRIX_KEYRELEASES)
+    if (!pressed)
+#    else
+    if (pressed)
+#    endif // defined(RGB_MATRIX_KEYRELEASES)
+    {
+        if (rgb_matrix_config.mode == RGB_MATRIX_TYPING_HEATMAP) {
+            process_rgb_matrix_typing_heatmap(row, col);
+        }
     }
 #endif // defined(RGB_MATRIX_FRAMEBUFFER_EFFECTS) && defined(ENABLE_RGB_MATRIX_TYPING_HEATMAP)
 }