# https://youtu.be/wTJI_WuZSwE

from functools import reduce, cache
from typing import List, Tuple, Callable, Iterable
from math import sqrt, log2
from random import Random

def get_flip_pos_2x2(board: List[int], key: int) -> int:
    assert all(x in [0,1] for x in board)
    assert len(board) == 4
    s = sum(board)
    if s==0:
        return key
    if s==1:
        pos = board.index(1)
        if key == 0:
            return pos # empty
        if key == 1:
            return [1,0,3,2][pos] # horizontal
        if key == 2:
            return [2,3,0,1][pos] # vertical
        if key == 3:
            return [3,2,1,0][pos] # diagonal
        assert False
    if s==2:
        # ensure ``sum(board) in [1,3]``, pointing to key
        if board[key] == 0:  # sum(flipped board) == 3
            possibilities = [i for i in range(len(board)) if i != key and board[i]==0]
        else: # sum(flipped board) == 3
            possibilities = [i for i in range(len(board)) if i != key and board[i]==1]
        assert len(possibilities) == 1
        return possibilities[0]
    if s==3:
        return get_flip_pos_2x2(all_flipped(board), key)
    if s==4:
        return get_flip_pos_2x2(all_flipped(board), key)
    assert False

def all_flipped(board: List[int]) -> List[int]:
    return [1-x for x in board]

def guess_key_2x2(board: List[int]) -> int:
    s = sum(board)
    if s==0:
        return 0
    if s==1:
        return board.index(1)
    if s==2:
        if board==[1,1,0,0] or board==[0,0,1,1]:
            return 1 # horizontal
        if board==[1,0,1,0] or board==[0,1,0,1]:
            return 2 # vertical
        if board==[1,0,0,1] or board==[0,1,1,0]:
            return 3 # diagonal
    if s==3:
        return guess_key_2x2(all_flipped(board))
    if s==4:
        return guess_key_2x2(all_flipped(board))
    assert False


def flip_mutate(board: List[int], flip_pos: int):
    board[flip_pos] = 1-board[flip_pos]

def flip_copy(board: List[int], flip_pos: int) -> List[int]:
    board = board[:]
    flip_mutate(board, flip_pos)
    return board

def print_board(board: List[int]):
    sidelength = sqrt(len(board))
    assert sidelength == int(sidelength)
    sidelength = int(sidelength)
    for row_i in range(sidelength):
        print(''.join(str(x) for x in board[row_i*sidelength:(row_i+1)*sidelength]))

def generate_boards(board_squre_count: int, sample_count: int) -> Iterable[int]:
    if sample_count >= 2**board_squre_count:
        print('doing exhaustive verification.')
        yield from range(2**board_squre_count)
    print(f'not doing exhaustive. need sample_count>={2**board_squre_count}, got sample_count={sample_count}')
    rng = Random(4)
    for i in range(sample_count):
        yield rng.getrandbits(board_squre_count)


def validate(board_squre_count, get_flip_pos: Callable, guess_key: Callable, sample_count=2**16) -> bool:
    success_boards = set()

    for sample_i, board_int in enumerate(generate_boards(board_squre_count, sample_count)):
        if board_int in success_boards:
            continue

        if sample_i %100 == 0:
            print(sample_i/sample_count)
        for key in range(board_squre_count):
            # potential optimization: could re-use board and unflip guess
            board = [(board_int>>i)&1 for i in reversed(range(board_squre_count))]

            flip_pos = get_flip_pos(board, key)
            flip_mutate(board, flip_pos)
            guess = guess_key(board)
            if guess != key:
                print(f'Found counter example. Got guess {guess}, want {key}')
                flip_mutate(board, flip_pos)
                print('before flip:')
                print_board(board)
                flip_mutate(board, flip_pos)
                print('after flip:')
                print_board(board)
                return False

        success_boards.add(board_int)
        if len(success_boards) == 2**board_squre_count:
            # All boards have been validated
            print('exhaustive, yay')
            return True
    print(f'non exhaustive finish. covered {len(success_boards)} cases ({len(success_boards) / 2**board_squre_count:%} of cases) successfully')
    return True

# validate(4, get_flip_pos_2x2, guess_key_2x2)

# Scaler possibilities:
# Total key index % 4 if we split into 4-squares
# is each reg
# XOR each 4x4 section into a bit to determine which quadrant
# layer each 4x4 section with XOR, to determine which quadrant in the above 4x4
# layer each 2x2 section with XOR, to determine which quadrant in 2x2

def guess_key_pow2board(board: List[int]) -> int:
    if len(board) == 4:
        return guess_key_2x2(board)

    sidelength = sqrt(len(board))
    assert sidelength == int(sidelength)
    sidelength = int(sidelength)
    assert log2(sidelength) == int(log2(sidelength))
    quarter = len(board) / 4
    assert quarter == int(quarter)
    quarter = int(quarter)

    # Not actually quadrants but good enough
    quadrants = [ board[i*quarter:(i+1)*quarter] for i in range(0,4) ]

    quarter_i = guess_key_2x2([ reduce(lambda a,b: a^b, q) for q in quadrants ])
    sub_i = guess_key_pow2board([a^b^c^d for a,c,b,d in zip(*quadrants)])
    return quarter * quarter_i + sub_i

def make_get_flip_pos(guess_key: Callable) -> Callable:

    # Flip the thing that works
    def get_flip_pos(board, key) -> int:
        possibilities = [
            flip_pos for flip_pos in range(len(board))
            if guess_key(flip_copy(board, flip_pos)) == key
        ]
        assert len(possibilities) == 1, f"possibilities={possibilities} want length 1"
        return possibilities[0]
    return get_flip_pos

def get_flip_pos_pow2board(board: List[int], key: int) -> int:
    if len(board) == 4:
        return get_flip_pos_2x2(board, key)
    quarter = len(board) / 4
    assert quarter == int(quarter)
    quarter = int(quarter)

    # Not actually quadrants but good enough
    quadrants = [ board[i*quarter:(i+1)*quarter] for i in range(0,4) ]
    quarter_i = get_flip_pos_2x2([ reduce(lambda a,b: a^b, q) for q in quadrants ], key//4)
    sub_i = get_flip_pos_pow2board([a^b^c^d for a,c,b,d in zip(*quadrants)], key%4)
    return quarter * quarter_i + sub_i


# validate(4, make_get_flip_pos(guess_key_2x2), guess_key_2x2)
# validate(16, make_get_flip_pos(guess_key_pow2board), guess_key_pow2board)
# validate(16, get_flip_pos_pow2board, guess_key_pow2board)
validate(64, get_flip_pos_pow2board, guess_key_pow2board)
# validate(64, make_get_flip_pos(guess_key_pow2board), guess_key_pow2board)

