import torch

from torch_struct import SentCFG


def params_l3():
    """
    seq = x y z, t0, t1 & n0, n1, n2
    """
    terms = [[2, 1], [1, 2], [1, 1]]
    # term4 = [[1, 1], [2, 1], [1, 2]]
    roots = [1, 1, 1]
    rule1 = [
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 5],
        [1, 1, 1, 2, 1],
    ]
    rule2 = [
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 6],
        [1, 1, 1, 2, 1],
    ]
    rule3 = [
        [1, 1, 1, 1, 1],
        [1, 1, 1, 4, 5],
        [1, 1, 1, 8, 9],
        [1, 1, 1, 1, 4],
        [1, 1, 1, 2, 1],
    ]
    rules = [rule1, rule2, rule3]

    terms = (
        torch.tensor(terms, dtype=torch.float64, requires_grad=True)
        .unsqueeze(0)
        .float()
    )
    roots = (
        torch.tensor(roots, dtype=torch.float64, requires_grad=True)
        .unsqueeze(0)
        .float()
    )
    rules = (
        torch.tensor(rules, dtype=torch.float64, requires_grad=True)
        .unsqueeze(0)
        .float()
    )

    length = torch.tensor([3]).long()

    # print('term:\n', terms, terms.shape)
    # print('root:\n', roots, roots.shape)
    # print('rule:\n', rules, rules.shape)
    return ((terms, rules, roots), length)


def extract_parse(span, length):
    tree = [(i, str(i)) for i in range(length)]
    tree = dict(tree)
    spans = []
    cover = (span > 0).float().nonzero()
    for i in range(cover.shape[0]):
        w, r, A = cover[i].tolist()
        w = w + 1
        r = r + w
        l = r - w
        spans.append((l, r, A))
        span = "({} {})".format(tree[l], tree[r])
        tree[r] = tree[l] = span
    return spans, tree[0]


def extract_topk(matrix, lengths):
    batch, K, N = matrix.shape[:3]
    spans = []
    trees = []
    for b in range(batch):
        for k in range(K):
            this_span = matrix[b][k]
            span, tree = extract_parse(this_span, lengths[b])
            trees.append(tree)
            spans.append(span)
            # print(span)
            # print(tree)
        # break
    return spans, trees


def extract_parses(matrix, lengths):
    batch, K, N = matrix.shape[:3]
    spans = []
    trees = []
    for b in range(batch):
        span, tree = extract_parse(matrix[b], lengths[b])
        trees.append(tree)
        spans.append(span)
        # print(span, tree)
        # break
    return spans, trees


def test_l3_kbest():
    params, lengths = params_l3()
    dist = SentCFG(params, lengths=lengths)

    _, _, _, spans = dist.argmax
    spans, trees = extract_parses(spans, lengths)
    best_trees = "((0 1) 2)"
    best_spans = [(0, 1, 2), (0, 2, 2)]
    assert spans[0] == best_spans
    assert trees[0] == best_trees

    _, _, _, spans = dist.topk(4)
    size = (1, 0) + tuple(range(2, spans.dim()))
    spans = spans.permute(size)
    spans, trees = extract_topk(spans, lengths)
    best_trees = "((0 1) 2)"
    best_spans = [
        [(0, 1, 2), (0, 2, 2)],
        [(0, 1, 2), (0, 2, 2)],
        [(0, 1, 1), (0, 2, 2)],
        [(0, 1, 1), (0, 2, 2)],
    ]
    for i, (span, tree) in enumerate(zip(spans, trees)):
        assert span == best_spans[i]
        assert tree == best_trees


if __name__ == "__main__":
    test_l3_kbest()