set RGB colors for key groups with support for layers (qmk)

keymap.c

enum keycode_groups {
    KC_GROUP_ALPHA,
    KC_GROUP_SYM,
    KC_GROUP_NAV,
    KC_GROUP_MOD,
    KC_GROUP_OTHER,

    KC_GROUP_COUNT
};

uint8_t HUE_MODIFIER = 255 / KC_GROUP_COUNT;

uint8_t get_keycode_group(uint16_t keycode) {
    if (keycode >= KC_A && keycode <= KC_Z) {
        return KC_GROUP_ALPHA;
    }

    if ((keycode >= KC_1 && keycode <= KC_0) || (keycode >= KC_MINUS && keycode <= KC_SLASH)) {
        return KC_GROUP_SYM;
    }

    if ((keycode >= KC_INSERT && keycode <= KC_UP)) {
        return KC_GROUP_NAV;
    }

    if (IS_QK_MODS(keycode) || IS_QK_ONE_SHOT_MOD(keycode)) {
        return KC_GROUP_MOD;
    }

    return KC_GROUP_OTHER;
}

// get_highest_layer_with_keymap_key gets the highest layer
// which has a non KC_TRNS value configured at a particular
// position.
uint8_t get_highest_layer_with_keymap_key(keypos_t key) {
    for (uint8_t l = get_highest_layer(layer_state); l > 0; --l) {
        if (IS_LAYER_ON(l) && keymap_key_to_keycode(l, key) > KC_TRNS) {
            return l;
        }
    }

    return 0;
}

// get_layer_depth computes the difference in layers between
// the current highest active layer and any other layer, without
// counting inactive layers in between.
uint8_t get_layer_depth(uint8_t layer) {
    uint8_t depth = 0;
    for (uint8_t l = get_highest_layer(layer_state); l > layer; --l) {
        if (IS_LAYER_ON(l)) {
            depth++;
        }
    }

    return depth;
}

// set_color sets the color for an led based on it's position.
// It gets the keycode at that position (in the layer obtained by
// calling get_highest_layer_with_keymap_key) and assigns a color
// based on the group it is associated to in get_keycode_group.
// An hsv color is then set in such a way that all groups have
// evenly spaced hues, and the value is lowered based on how
// many active layers are beneath the current active layer (get_layer_depth).
// This results in:
// - similar keys have the same color
// - all colors are evenly spaced out
// - keys from lower layers appear "dimmed" a bit
//   - (the more layers are active, the more dimmed some keys from the first few layers will appear)
// - you may adjust the colors by binding keys to:
//   - RGB_SAI, RGB_SAD: increase/decrease saturation
//   - RGB_VAI, RGB_VAD: increase/decrease overall brightness (works well with dimmed layers!)
//   - RGB_HUI, RGB_HUD: increase/decrease hue value (essentially, rotate the base color for each group)
void set_color(uint8_t led, keypos_t keypos) {
    uint8_t  layer   = get_highest_layer_with_keymap_key(keypos);
    uint16_t keycode = keymap_key_to_keycode(layer, keypos);

    if (keycode <= KC_TRNS) {
        return rgb_matrix_set_color(led, 0, 0, 0);
    }

    uint8_t group = get_keycode_group(keycode);

    HSV hsv = rgb_matrix_config.hsv;
    hsv.h += group * HUE_MODIFIER;
    hsv.v /= 1 + get_layer_depth(layer);

    RGB rgb = hsv_to_rgb(hsv);
    return rgb_matrix_set_color(led, rgb.r, rgb.g, rgb.b);
}

bool rgb_matrix_indicators_advanced_user(uint8_t led_min, uint8_t led_max) {
    if (keyboard_config.disable_layer_led || rgb_matrix_get_mode() != RGB_MATRIX_SOLID_COLOR || !HAS_FLAGS(rgb_matrix_config.flags, LED_FLAG_ALL)) {
        return false;
    }

    for (uint8_t row = 0; row < MATRIX_ROWS; ++row) {
        for (uint8_t col = 0; col < MATRIX_COLS; ++col) {
            uint8_t led = g_led_config.matrix_co[row][col];
            if (led < led_min || led > led_max || led == NO_LED) {
                continue;
            }

            set_color(led, MAKE_KEYPOS(row, col));
        }
    }

    return rgb_matrix_check_finished_leds(led_max);
}