keyboard_heatmap.py

import os
import csv
from collections import defaultdict

import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

FONT_SIZE = 10
FIGSIZE = (12, 6)
CSV_PATH = "keyboard_freq.csv"
PERCENT_THRESHOLD = 0.01

COLOR_WHITE = "#FFFFFF"
COLOR_L4 = "#DDDDDD"
COLOR_L3 = "#BBBBBB"
COLOR_L2 = "#888888"
COLOR_L1 = "#000000"

KEY_ALIAS = {
    "[space]": "SPACE",
    "[shift]": "SHIFT", "[shift_r]": "RSHIFT",
    "[ctrl_l]": "CTRL",
    "[alt_l]": "ALT", "[alt_gr]": "RALT",
    "[cmd]": "WIN",
    "[esc]": "ESC",
    "[enter]": "ENTER",
    "[tab]": "TAB",
    "[backspace]": "BACKSPACE",
    "[insert]": "INS",
    "[delete]": "DEL",
    "[page_up]": "PGUP",
    "[page_down]": "PGDN",
    "[up]": "↑", "[down]": "↓", "[left]": "←", "[right]": "→",
    "[caps_lock]": "CAPS LOCK",
}
KEY_ALIAS.update({f"[f{i}]": f"F{i}" for i in range(1, 12 + 1)})

CHAR_MAPPING = {
    "`": "` ~", "~": "` ~",
    "1": "1 !", "!": "1 !",
    "2": "2 @", "@": "2 @",
    "3": "3 #", "#": "3 #",
    "4": "4 $", "$": "4 $",
    "5": "5 %", "%": "5 %",
    "6": "6 ^", "^": "6 ^",
    "7": "7 &", "&": "7 &",
    "8": "8 *", "*": "8 *",
    "9": "9 (", "(": "9 (",
    "0": "0 )", ")": "0 )",
    "-": "- _", "_": "- _",
    "=": "= +", "+": "= +",
    "[": "[ {", "{": "[ {",
    "]": "] }", "}": "] }",
    "\\": "\\ |", "|": "\\ |",
    ";": "; :", ":": "; :",
    "'": "' \"", "\"": "' \"",
    ",": ", <", "<": ", <",
    ".": ". >", ">": ". >",
    "/": "/ ?", "?": "/ ?",
}

IGNORE_KEYS = {"\x16", "[f23]"}


def normalize_key_name(name):
    if name in IGNORE_KEYS:
        return None
    if name in KEY_ALIAS:
        return KEY_ALIAS[name]
    if len(name) == 1:
        if name.isalpha():
            return name.upper()
        if name in CHAR_MAPPING:
            return CHAR_MAPPING[name]
    return name


def load_freq(path):
    freq = {}
    with open(path, "r", encoding="utf-8") as f:
        for row in csv.DictReader(f):
            k = row["key"]
            c = int(row["count"])
            freq[k] = freq.get(k, 0) + c
    return freq


def build_display_freq(raw_freq):
    disp = defaultdict(int)
    for k, c in raw_freq.items():
        label = normalize_key_name(k)
        if label is not None:
            disp[label] += c
    return dict(disp)


main_base = [
    [("ESC", 1),
     ("F1", 1), ("F2", 1), ("F3", 1), ("F4", 1),
     ("F5", 1), ("F6", 1), ("F7", 1), ("F8", 1),
     ("F9", 1), ("F10", 1), ("F11", 1), ("F12", 1)],
    [("` ~", 1), ("1 !", 1), ("2 @", 1), ("3 #", 1), ("4 $", 1),
     ("5 %", 1), ("6 ^", 1), ("7 &", 1), ("8 *", 1), ("9 (", 1),
     ("0 )", 1), ("- _", 1), ("= +", 1),
     ("BACKSPACE", 2)],
    [("TAB", 1.5),
     ("Q", 1), ("W", 1), ("E", 1), ("R", 1), ("T", 1),
     ("Y", 1), ("U", 1), ("I", 1), ("O", 1), ("P", 1),
     ("[ {", 1), ("] }", 1),
     ("\\ |", 1.5)],
    [("CAPS LOCK", 1.75),
     ("A", 1), ("S", 1), ("D", 1), ("F", 1), ("G", 1),
     ("H", 1), ("J", 1), ("K", 1), ("L", 1),
     ("; :", 1), ("' \"", 1),
     ("ENTER", 2)],
    [("SHIFT", 2.25),
     ("Z", 1), ("X", 1), ("C", 1), ("V", 1), ("B", 1),
     ("N", 1), ("M", 1),
     (", <", 1), (". >", 1), ("/ ?", 1),
     ("RSHIFT", 2.75)],
    [("CTRL", 1.25), ("WIN", 1.25), ("ALT", 1.25),
     ("SPACE", 6.25),
     ("RALT", 1.25), ("RWIN", 1.25), ("MENU", 1.25), ("RCTRL", 1.25)],
]

cluster_rows = [
    [("PRTSC", 1), ("LOCK", 1), ("PAUSE", 1)],
    [("INS", 1), ("HOME", 1), ("PGUP", 1)],
    [("DEL", 1), ("END", 1), ("PGDN", 1)],
    [("", 1), ("", 1), ("", 1)],
    [("", 1), ("↑", 1), ("", 1)],
    [("←", 1), ("↓", 1), ("→", 1)],
]

row_widths = [sum(w for _, w in row) for row in main_base]
main_width = max(row_widths)
cluster_width = 3

main_scaled = []
for row, w_row in zip(main_base, row_widths):
    s = main_width / w_row
    main_scaled.append([(label, w * s) for label, w in row])


def freq_to_gray(pct, max_pct, threshold_pct):
    if max_pct <= 0 or pct < threshold_pct:
        return COLOR_WHITE
    norm = 1.0 if max_pct == threshold_pct else (pct - threshold_pct) / (max_pct - threshold_pct)
    norm = min(max(norm, 0.0), 1.0)
    if norm < 0.25:
        return COLOR_L4
    if norm < 0.5:
        return COLOR_L3
    if norm < 0.75:
        return COLOR_L2
    return COLOR_L1


def pick_text_color(bg_hex):
    c = bg_hex.lstrip("#")
    if len(c) != 6:
        return "black"
    r, g, b = int(c[0:2], 16), int(c[2:4], 16), int(c[4:6], 16)
    y = 0.2126 * r + 0.7152 * g + 0.0722 * b
    return "black" if y > 170 else "white"


def draw_key(ax, x, y, w, h, label, freq_pct, max_pct, threshold_pct):
    pct = freq_pct.get(label, 0.0)
    face = freq_to_gray(pct, max_pct, threshold_pct)
    rect = Rectangle((x, y), w, h, linewidth=1, edgecolor="black", facecolor=face)
    ax.add_patch(rect)
    ax.text(x + w / 2, y + h / 2, label,
            ha="center", va="center",
            fontsize=FONT_SIZE, color=pick_text_color(face))


def draw_keyboard(ax, main_rows, cluster_rows, freq_pct, key_h=1.0):
    layout_labels = [label for rows in (main_rows, cluster_rows)
                     for row in rows for label, _ in row if label]
    layout_set = set(layout_labels)
    extra_keys = [k for k in freq_pct if k not in layout_set]
    if extra_keys:
        print("WARNING: the following keys exist in CSV after mapping, but do not exist in the current layout:")
        for k in sorted(extra_keys):
            print("   ", repr(k))

    max_pct = max(freq_pct.values()) if freq_pct else 0.0
    threshold_pct = PERCENT_THRESHOLD
    print(f"max_pct = {max_pct:.6f}, threshold_pct = {threshold_pct:.6f}")

    for i in range(len(main_rows)):
        y = -i * key_h
        x = 0.0
        for label, w in main_rows[i]:
            if label:
                draw_key(ax, x, y, w, key_h, label, freq_pct, max_pct, threshold_pct)
            x += w
        x = main_width
        for label, w in cluster_rows[i]:
            if label:
                draw_key(ax, x, y, w, key_h, label, freq_pct, max_pct, threshold_pct)
            x += w


def main():
    raw = load_freq(CSV_PATH)
    disp = build_display_freq(raw)
    total = sum(disp.values())
    freq_pct = {k: v / total for k, v in disp.items()} if total else {}
    fig, ax = plt.subplots(figsize=FIGSIZE)
    ax.set_facecolor("#FFFFFF")
    draw_keyboard(ax, main_scaled, cluster_rows, freq_pct, key_h=1.0)
    ax.set_xlim(0, main_width + cluster_width)
    ax.set_ylim(-len(main_scaled), 1)
    ax.axis("off")
    plt.tight_layout()
    plt.savefig("keyboard_heatmap.png", dpi=300)
    print("Saved:", os.path.abspath("keyboard_heatmap.png"))
    plt.show()

if __name__ == "__main__":
    main()