#!/usr/bin/env python3

from collections import deque
from math import inf
from random import randrange
import json

# A proposed "set of simple selectors" s_{1..k} for a configuration is consistent iff
# for each selector s_i,
#  * s_i is a boolean fixed column; and
#  * all polynomial constraints involving s_i are of the form s_i . t = 0 where t does
#    not involve any of s_{1..k}.
#
# Given a set of simple selectors, we can potentially find a partition of that set
# that combines some selectors. For example, if s_1 and s_2 are used only on disjoint
# sets of rows, then we can combine them into s', where
#
#   s' = 0, if s_1 = s_2 = 0
#      = 1, if s_1 = 1
#      = 2, if s_2 = 1
#
# and then replace
#
#   s_1 . t_1 = 0 by s' . (2-s') . t_1 = 0
#   s_2 . t_2 = 0 by s' . (1-s') . t_2 = 0
#
# More generally, we interpolate whether each original selector in a given combination
# is zero/non-zero from its combined selector (we can neglect dividing by the constant
# to normalize to 1).

"""Iterate through combinations with one entry chosen from each disjoint set."""
def iter_combinations(width, X, degrees, degree_bound):
    i = 0
    seen = [False]*width

    while i < width:
        combination = deque()
        d = 0

        for j in range(width):
            def is_excluded():
                if seen[j]:
                    return True

                for k in combination:
                    assert k < j
                    if X[j][k] == 1:
                        return True

                return False

            if is_excluded():
                continue

            new_d = max(d, degrees[j])
            if new_d + len(combination) + 1 > degree_bound:
                continue
            d = new_d

            i += 1
            combination.append(j)

            # Don't repeat entries across combinations.
            seen[j] = True

        yield list(combination)


def check(result, selectors, degrees, degree_bound):
    for combination in result:
        d = max([degrees[s] for s in combination]) + len(combination)
        assert d <= degree_bound, "%r has degree %r, which violates degree bound %r" % (combination, d, degree_bound)

    n = len(selectors[0])
    for row in range(n):
        for combination in result:
            # At most one selector in the combination must be set on this row.
            count = sum([selectors[s][row] for s in combination])
            assert count <= 1, "row %r has %r selectors from %r set" % (row, count, combination)


def compute_combinations(selectors, degrees, degree_bound):
    width = len(selectors)
    assert width > 0
    n = len(selectors[0])

    print("  selectors:")
    for col in selectors:
        print("  " + str(col))

    print("\n  degrees:")
    print("  " + str(degrees))
    print("  degree_bound =", degree_bound)
    print("  ")

    assert len(degrees) == width

    # Compute a conflict matrix X, in which X[i][j] is set if selectors
    # i and j are excluded from being combined.
    # Since X is symmetric and diagonal entries are zero, we only need to
    # store the lower triangular subset of entries.
    X = [[0]*i for i in range(width)]

    prev = []
    for row in range(n):
        # Each pair of selectors set in this row are conflicted.
        selected = [h for h in range(width) if selectors[h][row] == 1]

        for (i, j) in enumerate(selected):
            for k in selected[:i]:
                #print("  %r <-> %r" % (j, k))
                X[j][k] = 1

    for j in range(width):
        print("  " + str(X[j]))
    print()

    return list(iter_combinations(width, X, degrees, degree_bound))

"""
Is the list xs a subset of (including equal to) the list ys?
Both lists are sorted and have no duplicates.
"""
def is_subset(xs, ys):
    len_ys = len(ys)
    if len(xs) > len_ys:
        return False

    i = 0
    for x in xs:
        while ys[i] < x:
            i += 1
            if i == len_ys:
                return False
        if ys[i] > x:
            return False

    return True


def transpose(M):
    cols = len(M)
    rows = len(M[0])
    print("  cols =", cols)
    print("  rows =", rows)
    return [[M[col][row] for col in range(cols)] for row in range(rows)]


"""
selectors = transpose([
    [0, 0, 0, 0, 1],
    [1, 0, 0, 0, 1],
    [0, 1, 1, 0, 0],
    [0, 1, 0, 0, 0],
    [0, 0, 1, 1, 0],
])
width = 5
degree_bound = 11
degrees = (
    [3, 6, 5, 8, 2]
)
"""

with open("action-circuit-selectors.json") as f:
    content = json.loads(f.read())

selectors = content["selectors"]
width = len(selectors)
degree_bound = 9
#degrees = [randrange(4, 8) for i in range(width)]
degrees = [7, 4, 7, 4, 6, 6, 6, 7, 7, 6, 7, 4, 4, 6, 5, 6, 7, 4, 5, 5, 4, 6, 4, 4, 4, 4, 5, 4, 4, 5, 6, 6, 4, 7, 4, 7, 5, 6, 4, 7]

degs_sels = list(zip(degrees, selectors))
degs_sels.sort(key=lambda ds: ds[0], reverse=True)
degrees   = [ds[0] for ds in degs_sels]
selectors = [ds[1] for ds in degs_sels]

result = compute_combinations(selectors, degrees, degree_bound)
check(result, selectors, degrees, degree_bound)
print("  " + str(result))