import torch
import numpy as np

from src.debfly.product import matrix_product
from src.debfly.support import bfly_supp
from src.debfly.tree import Node, square_pot_butterfly_tree
from src.debfly.product import twiddle_list_to_dense

from scipy.stats import ortho_group, unitary_group


def _fixed_supp_random_matrix(supp, mean=1, std=0.5):
    m = torch.distributions.bernoulli.Bernoulli(probs=0.5 * torch.ones_like(supp))
    supp = supp * (m.sample() * 2 - 1)
    return supp * (torch.randn(supp.shape[0], supp.shape[1]) * std + mean)


def random_bfly_matrix(log_n, mean=1, std=0.5):
    """generate a butterfly matrix where the butterfly factors have random entries"""
    list_supp = [bfly_supp(i, i, log_n) for i in range(1, log_n + 1)]
    list_fs_random_mat = [
        _fixed_supp_random_matrix(supp, mean, std) for supp in list_supp
    ]
    return matrix_product(list_fs_random_mat)


def random_parker_butterfly(log_n, device="cpu", norm=1.0, complex_value=False):
    """generate a random orthogonal butterfly matrix"""
    tree = square_pot_butterfly_tree(log_n, "unbalanced")
    supscript_list, subscript_list = tree.supscript_subscript_of_leaves()
    t_list = []
    for supscript, subscript in zip(supscript_list, subscript_list):
        assert subscript[0] == subscript[1] == 2
        if not complex_value:
            random_orth = norm * ortho_group.rvs(2, supscript[0] * supscript[1])
            twiddle = torch.from_numpy(random_orth).to(device)
        else:
            random_unit = norm * unitary_group.rvs(2, supscript[0] * supscript[1])
            twiddle = torch.from_numpy(random_unit).to(device)
        twiddle = twiddle.reshape(
            supscript[0], supscript[1], subscript[0], subscript[1]
        )
        t_list.append(twiddle)
    return twiddle_list_to_dense(t_list).squeeze()


def random_debfly_matrix(tree, device="cpu", complex_value=False):
    """compute a random butterfly matrix that admits a factorization described by tree"""
    assert isinstance(tree, Node)
    supscript_list, subscript_list = tree.supscript_subscript_of_leaves()
    t_list = []
    for supscript, subscript in zip(supscript_list, subscript_list):
        if complex_value:
            twiddle = torch.randn(
                supscript[0],
                supscript[1],
                subscript[0],
                subscript[1],
                device=device,
                dtype=torch.complex128,
            )
            twiddle += 1.0j * torch.randn(
                supscript[0], supscript[1], subscript[0], subscript[1], device=device
            )
        else:
            twiddle = torch.randn(
                supscript[0],
                supscript[1],
                subscript[0],
                subscript[1],
                device=device,
                dtype=torch.float64,
            )
        t_list.append(twiddle)
    return twiddle_list_to_dense(t_list).squeeze()


def rank_one_family_partitioning(partitioning, dim):
    """
    Generate a family of rank-one vectors. The family is partitioned according to partitioning.
    In each subset of the partition, vectors are in the same span.
    Vectors are of dimension dim.
    """
    n_subsets = len(partitioning)
    card = len(partitioning[0])
    all_vectors = torch.empty(n_subsets * card, dim)
    for subset in partitioning:
        assert len(subset) == card
        random_vectors = torch.outer(torch.randn(card), torch.randn(dim))
        idx = torch.LongTensor(subset)
        all_vectors[idx, :] = random_vectors
    return all_vectors


def random_haar_butterfly(log_n):
    if log_n == 0:
        return torch.tensor([[1.0]])
    theta = 2 * np.pi * np.random.rand()
    rotation = torch.tensor(
        [[np.cos(theta), np.sin(theta)], [-np.sin(theta), np.cos(theta)]]
    )
    return torch.kron(rotation, random_haar_butterfly(log_n - 1))


def concat_2_by_2_blocks(blocks):
    return torch.cat(
        [
            torch.cat([blocks[0][0], blocks[0][1]], dim=1),
            torch.cat([blocks[1][0], blocks[1][1]], dim=1),
        ],
        dim=0,
    )


def random_diagonal_butterfly(log_n, simple=False):
    if log_n == 0:
        return torch.tensor([[1.0]])
    n = 2**log_n
    theta = torch.tensor(2 * np.pi * np.random.rand(n // 2), dtype=torch.float)
    if simple:
        butterfly = random_diagonal_butterfly(log_n - 1, simple=simple)
        return concat_2_by_2_blocks(
            [
                [
                    torch.cos(theta).unsqueeze(1) * butterfly,
                    torch.sin(theta).unsqueeze(1) * butterfly,
                ],
                [
                    -torch.sin(theta).unsqueeze(1) * butterfly,
                    torch.cos(theta).unsqueeze(1) * butterfly,
                ],
            ]
        )
    butterfly_1 = random_diagonal_butterfly(log_n - 1, simple=simple)
    butterfly_2 = random_diagonal_butterfly(log_n - 1, simple=simple)
    return concat_2_by_2_blocks(
        [
            [
                torch.cos(theta).unsqueeze(1) * butterfly_1,
                torch.sin(theta).unsqueeze(1) * butterfly_2,
            ],
            [
                -torch.sin(theta).unsqueeze(1) * butterfly_1,
                torch.cos(theta).unsqueeze(1) * butterfly_2,
            ],
        ]
    )


if __name__ == "__main__":
    n = 3
    a = random_parker_butterfly(n)
    print((a.t() @ a - torch.eye(2**n)).sum())
    print((a @ a.t() - torch.eye(2**n)).sum())