import zlib
import struct
import math
import os

SIZE = 16
CENTER = (SIZE - 1) / 2.0

# Double-headed arrow parameters, in pixels, along/perp to the arrow's axis.
RADIUS = 7.0
SHAFT_LENGTH = 2.2
SHAFT_HALF_WIDTH = 1.0
HEAD_HALF_WIDTH = 3.4


def arrow_mask(axis_degrees):
    theta = math.radians(axis_degrees)
    ux, uy = math.cos(theta), math.sin(theta)
    mask = [[False] * SIZE for _ in range(SIZE)]

    for y in range(SIZE):
        for x in range(SIZE):
            # Sample the pixel center.
            dx = (x + 0.5) - (CENTER + 0.5)
            dy = (y + 0.5) - (CENTER + 0.5)

            along = dx * ux + dy * uy
            perp = -dx * uy + dy * ux

            absAlong = abs(along)

            if absAlong <= SHAFT_LENGTH:
                inside = abs(perp) <= SHAFT_HALF_WIDTH
            elif absAlong <= RADIUS:
                span = RADIUS - SHAFT_LENGTH
                headHalfWidth = HEAD_HALF_WIDTH * (RADIUS - absAlong) / span
                inside = abs(perp) <= headHalfWidth
            else:
                inside = False

            mask[y][x] = inside

    return mask


def dilate(mask):
    returnMe = [[False] * SIZE for _ in range(SIZE)]

    for y in range(SIZE):
        for x in range(SIZE):
            if mask[y][x]:
                continue

            for offsetY in (-1, 0, 1):
                for offsetX in (-1, 0, 1):
                    neighborX, neighborY = x + offsetX, y + offsetY

                    if 0 <= neighborX < SIZE and 0 <= neighborY < SIZE and mask[neighborY][neighborX]:
                        returnMe[y][x] = True

    return returnMe


def render_rgba(axis_degrees):
    mask = arrow_mask(axis_degrees)
    outline = dilate(mask)

    rgba = bytearray(SIZE * SIZE * 4)

    for y in range(SIZE):
        for x in range(SIZE):
            offset = (y * SIZE + x) * 4

            if mask[y][x]:
                rgba[offset:offset + 4] = bytes((255, 255, 255, 255))
            elif outline[y][x]:
                rgba[offset:offset + 4] = bytes((0, 0, 0, 255))
            else:
                rgba[offset:offset + 4] = bytes((0, 0, 0, 0))

    return rgba


def write_png(path, rgba, width, height):
    def chunk(tag, data):
        return (struct.pack(">I", len(data)) + tag + data +
                struct.pack(">I", zlib.crc32(tag + data) & 0xffffffff))

    signature = b"\x89PNG\r\n\x1a\n"
    ihdr = struct.pack(">IIBBBBB", width, height, 8, 6, 0, 0, 0)

    raw = bytearray()
    stride = width * 4

    for y in range(height):
        raw.append(0)
        raw.extend(rgba[y * stride:(y + 1) * stride])

    idat = zlib.compress(bytes(raw), 9)

    with open(path, "wb") as file:
        file.write(signature)
        file.write(chunk(b"IHDR", ihdr))
        file.write(chunk(b"IDAT", idat))
        file.write(chunk(b"IEND", b""))


# Sheet cell coordinates (axis degrees, column, row).  Cells are a plain SIZE x SIZE square with no border between
# them, unlike UiPacksheetMaterials' grid, since nothing here samples across cell boundaries.
CURSORS = {
    "Horizontal": (0, 0, 0),
    "Vertical": (90, 1, 0),
    "DiagonalNwSe": (-45, 0, 1),
    "DiagonalNeSw": (45, 1, 1),
}

SHEET_COLUMNS = 2
SHEET_ROWS = 2
SHEET_FILE_NAME = "ResizeCursors.png"


def compose_sheet():
    sheetWidth = SHEET_COLUMNS * SIZE
    sheetHeight = SHEET_ROWS * SIZE
    sheet = bytearray(sheetWidth * sheetHeight * 4)

    for axisDegrees, column, row in CURSORS.values():
        # Screen Y grows downward, so a positive on-screen angle needs its sign flipped for math's CCW convention.
        rgba = render_rgba(-axisDegrees)

        originX = column * SIZE
        originY = row * SIZE

        for y in range(SIZE):
            sourceOffset = y * SIZE * 4
            destinationOffset = ((originY + y) * sheetWidth + originX) * 4
            sheet[destinationOffset:destinationOffset + SIZE * 4] = rgba[sourceOffset:sourceOffset + SIZE * 4]

    return sheet, sheetWidth, sheetHeight


def main():
    outputDirectory = os.path.dirname(os.path.abspath(__file__))
    sheet, sheetWidth, sheetHeight = compose_sheet()
    write_png(os.path.join(outputDirectory, SHEET_FILE_NAME), sheet, sheetWidth, sheetHeight)
    print(f"wrote {SHEET_FILE_NAME}")


if __name__ == "__main__":
    main()