from pydantic import BaseModel
from secrets import choice

from src.models import KeypadSize
from src.utils import list_to_matrix, secure_fisher_yates_shuffle, matrix_to_list, matrix_transpose


class UserInterface(BaseModel):
    attr_indices: list[int]
    keypad_size: KeypadSize

    @staticmethod
    def new(keypad_size: KeypadSize):
        # Todo: this a hack do a proper random interface
        interface = UserInterface(
            attr_indices=list(range(keypad_size.numb_of_attrs)),
            keypad_size=keypad_size
        )
        interface.disperse_interface()
        for _ in range(10):
            interface.shuffle_interface()
        return interface

    def disperse_interface(self):
        user_interface_matrix = list_to_matrix(self.attr_indices, self.keypad_size.attrs_per_key)
        shuffled_keys = secure_fisher_yates_shuffle(user_interface_matrix)
        dispersed_interface = self._random_attribute_rotation(shuffled_keys, list(range(self.keypad_size.attrs_per_key)))
        self.attr_indices = matrix_to_list(dispersed_interface)

    def shuffle_interface(self):
        """just like dispersion but only half the sets are rotated"""
        numb_of_selected_sets = self.keypad_size.attrs_per_key // 2
        # randomly shuffle half the sets. if attrs_per_key is odd, randomly add one 50% of the time
        numb_of_selected_sets += choice([0, 1]) if (self.keypad_size.attrs_per_key & 1) == 1 else 0
        selected_sets = secure_fisher_yates_shuffle(list(range(self.keypad_size.attrs_per_key)))[:numb_of_selected_sets]
        user_interface_matrix = list_to_matrix(self.attr_indices, self.keypad_size.attrs_per_key)
        shuffled_keys = secure_fisher_yates_shuffle(user_interface_matrix)
        interface_by_sets = []
        for idx, attrs in enumerate(matrix_transpose(shuffled_keys)):
            if idx in selected_sets:
                interface_by_sets.append(secure_fisher_yates_shuffle(attrs))
            else:
                interface_by_sets.append(attrs)
        self.attr_indices = matrix_to_list(matrix_transpose(interface_by_sets))

    def _random_attribute_rotation(self, user_interface: list[list[int]], selected_sets: list[int]) -> list[list[int]]:
        attr_rotation = secure_fisher_yates_shuffle(list(range(self.keypad_size.numb_of_keys)))[:self.keypad_size.attrs_per_key]
        transposed_user_interface = matrix_transpose(user_interface)
        assert (len(attr_rotation) == len(transposed_user_interface))
        for idx, attr_set in enumerate(transposed_user_interface):
            if idx not in selected_sets:
                continue
            rotation = attr_rotation[idx]
            transposed_user_interface[idx] = attr_set[rotation:] + attr_set[:rotation]
        return matrix_transpose(transposed_user_interface)

    def attribute_adjacency_graph(self) -> dict[int, set[int]]:
        user_interface_keypad = list_to_matrix(self.attr_indices, self.keypad_size.attrs_per_key)
        graph = {}
        for key in user_interface_keypad:
            for attr in key:
                graph[attr] = set(key)
                graph[attr].remove(attr)

        return graph

    def get_key_attr_idxs(self, key_numb: int) -> list[int]:
        assert (0 <= key_numb < self.keypad_size.numb_of_keys)
        keypad_attr_idx = list_to_matrix(self.attr_indices, self.keypad_size.attrs_per_key)
        return keypad_attr_idx[key_numb]