import torch
import numpy as np
import itertools
from typing import Optional


def generate_truth_table(
    k: int, tables_count: int, device: Optional[str]
) -> torch.Tensor:
    """This function generate truth tables with size of k * (2**k) * tables_count

    Args:
        k (int): truth table power
        tables_count (int): number of truth table repetition
        device (str): target device of the result

    Returns:
        torch.Tensor: 2d torch tensor with k*tables_count rows and (2**k) columns
    """

    table = torch.from_numpy(np.array(list(itertools.product([-1, 1], repeat=k)))).T
    return torch.vstack([table] * tables_count).to(device)


def reduce_truth_table(k: int, table: torch.Tensor, device: str) -> torch.Tensor:
    """This function reduce truth table generated by generate_truth_table by multiply each consecutive k lines

    Args:
        k (int): truth table power
        table (torch.Tensor): Nd array where truth table represented in last 2d
        device (str): target device of the results

    Returns:
        torch.Tensor: Nd torch tensor -dims count not changing- reduced number of rows in the table by k
    """
    result = torch.ones_like(table)
    result = result[..., 0::k, :]
    for i in range(k):
        result = result * table[..., i::k, :]
    return result.to(device)
