|
import math |
|
|
|
|
|
# essentially inverts the array to map val -> index |
|
def color_map(colors: list) -> dict: |
|
color_map = {} |
|
for i, v in enumerate(colors): |
|
color_map[v] = i |
|
|
|
assert len(color_map) == len(colors) # no dupes allowed |
|
return color_map |
|
|
|
|
|
# maps observed colors to integers; not preferred; color_map is preferred |
|
def color_map_from_bitmap(grid, zero="ZERO_NOT_SPECIFIED") -> dict: |
|
colors = {} |
|
if zero != "ZERO_NOT_SPECIFIED": |
|
colors[zero] = 0 |
|
|
|
for x, col in enumerate(grid): |
|
assert len(col) == len(grid) # must be square |
|
for y, val in enumerate(col): |
|
if val not in colors: |
|
colors[val] = len(colors) |
|
return colors |
|
|
|
|
|
def bpp(n_colors: int) -> int: |
|
return int(math.ceil(math.log2(n_colors))) |
|
|
|
|
|
def encode_socket(grid, color_map) -> int: |
|
n_colors = len(color_map) |
|
size = len(grid) |
|
bits_per_pixel = bpp(n_colors) |
|
|
|
socket = 0 |
|
for x, col in enumerate(grid): |
|
assert len(col) == size # must be square |
|
for y, val in enumerate(col): |
|
color = color_map[val] |
|
pixel_idx = bits_per_pixel * (x + y * size) |
|
socket = socket | (color << pixel_idx) |
|
|
|
return socket |
|
|
|
|
|
def decode_socket(socket: int, size: int, colors: list): |
|
n_colors = len(color_map) |
|
bits_per_pixel = bpp(n_colors) |
|
max_col = int(math.pow(2, bits_per_pixel)) - 1 |
|
|
|
grid = [] |
|
for x in range(size): |
|
grid.append([]) |
|
for y in range(size): |
|
pixel_idx = bits_per_pixel * (x + y * size) |
|
color_idx = max_col & (socket >> pixel_idx) |
|
print(format(max_col, f'#010b'), format(socket, f'#010b'), format((socket >> pixel_idx), f'#010b')) |
|
color = colors[color_idx] |
|
grid[x].append(color) |
|
assert len(grid[x]) == size |
|
assert len(grid) == size |
|
|
|
return grid |
|
|
|
|
|
if __name__ == "__main__": |
|
colors = [None, "Red", "Blue", "Green"] |
|
grid = [ |
|
[None, "Red"], |
|
["Green", None], |
|
] |
|
color_map = color_map(colors) |
|
sock = encode_socket(grid, color_map) |
|
# print(sock, format(sock, f'#010b')) |
|
|
|
recovered = decode_socket(sock, len(grid), colors) |
|
print(recovered) |