cs50p/project.py

"""Final project for CS50P - Davte. See README.md for more information."""
import math
import random
from typing import Tuple

import fpdf
import fpdf.table
from fpdf.fonts import FontFace
from fpdf.enums import TextEmphasis


BOLD = TextEmphasis.coerce('B')


def draw_table(pdf: fpdf.FPDF, state: str, key_position: int = -1, coin_to_flip: int = -1,
               highlight_parity_bits: bool = False) -> fpdf.table.Table:
    square_side = int(len(state)**0.5)
    with pdf.table(text_align='CENTER',
                   first_row_as_headings=False) as table:
        for i, v in enumerate(state):
            if i % square_side == 0:
                row = table.row()
            cell_value = {'0': 'H', '1': 'T'}[v]
            cell_style = FontFace()
            if i == coin_to_flip:
                cell_style.fill_color = (153, 0, 153)
                cell_style.emphasis = BOLD
            if i == key_position:
                cell_style.color = (255, 128, 0)
                cell_style.emphasis = BOLD
            elif highlight_parity_bits and i in (1, 2, 4, 8, 16, 32):
                cell_style.color = (119, 136, 153)
                cell_style.emphasis = BOLD
            row.cell(cell_value, style=cell_style)
    return table


def get_parity(n: str) -> int:
    num_bits = int(math.log2(len(n)))
    parity = 0
    for i in range(num_bits):
        block_parity = 0
        for j, val in enumerate(n):
            if j & (2**i) == 2**i:
                block_parity = block_parity ^ int(val)
        parity += block_parity * (2**i)
    return parity


def get_coin_to_flip(initial_state: str, key_position: int) -> int:
    current_value = get_parity(initial_state)
    return current_value ^ key_position


def store_pdf(file_name, state, key_position: int = -1, coin_to_flip: int = -1):
    pdf = fpdf.FPDF(orientation="P", unit="mm", format="A4")
    pdf.set_auto_page_break(False)
    pdf.add_page()
    pdf.set_font("Times", size=100 // math.log2(math.sqrt(len(state))))
    draw_table(pdf=pdf, state=state, key_position=key_position, coin_to_flip=coin_to_flip)
    pdf.output(file_name)


def solve(initial_state: str, coin_to_flip: int) -> str:
    """Return the chessboard in `initial_state` after flipping `coin_to_flip`."""
    result = list(initial_state)
    result[coin_to_flip] = '0' if result[coin_to_flip] == '1' else '1'
    return ''.join(result)


def is_power_of_two(n: int) -> bool:
    k = 1
    while k <= n:
        if k == n:
            return True
        k *= 2
    return False


def get_parameters(board_side: int = 8) -> Tuple[str, int, int]:
    """Generate a random chessboard and a random key position and solve the puzzle.

    Return a board of side `board_side`, a key position and the coin to flip.
    """
    if not is_power_of_two(board_side):
        raise ValueError("Board side must be a power of two!")
    random.seed()
    initial_state = ''.join(map(str, (random.randint(0, 1) for _ in range(board_side ** 2))))
    key_position = random.randint(0, board_side ** 2 - 1)
    coin_to_flip = get_coin_to_flip(initial_state, key_position)
    return initial_state, key_position, coin_to_flip


def main() -> None:
    board_side = 0
    while board_side < 2:
        try:
            board_side = input("Choose a side length for the chessboard (press enter for default value 8)\t\t")
            if not board_side:
                board_side = 8
            board_side = int(board_side)
            if not is_power_of_two(board_side):
                raise ValueError
        except (ValueError, TypeError):
            board_side = 0
            print(f"Invalid input `{board_side}`. Please enter a power of two.")
            continue
        except KeyboardInterrupt:
            print("\nExiting...")
            return
    print(f"Generating a random {board_side} x {board_side} chessboard...")
    initial_state, key_position, coin_to_flip = get_parameters(board_side=board_side)
    print("Show Player 1 the file `Key.pdf`.")
    store_pdf(file_name='Key.pdf', state=initial_state,
              key_position=key_position)
    final_state = solve(initial_state, coin_to_flip)
    print("Once Player 1 has flipped a coin, the chessboard should look like "
          "the one in `Problem.pdf`. Show it to Player 2.")
    store_pdf(file_name='Problem.pdf', state=final_state)
    print("You can use `Solution.pdf` to validate the answer of Player 2.")
    store_pdf(file_name='Solution.pdf', state=final_state,
              key_position=key_position, coin_to_flip=coin_to_flip)


if __name__ == '__main__':
    main()