"""On Solving Simple Curved Nonograms
by Maarten Löffler, Günter Rote, Soeren Terziadis, and Alexandra Weinberger"""
# indices start at 0 (following Python conventions) and not at 1 (as in the paper)

from collections import defaultdict

def decomposition_tree(f, show=False):
    """construct the binary decomposition tree.
    f is a the "face pattern", a string (or list of numbers or other objects)
    with one entry for each segment of the curve.
    Equal letters indicate segments belonging to the same face.

    Tree has three types of nodes, see Figure 5:
    * "Singleton", a leaf.
    * "Separate", with two children (white node).
    * "Progressive", with two children (black node): progressive sleuth + bracket
    "Separate" and "Progressive" store a triplet (i,j,k) of indices:
    f[i..k] (f[i:k+1] in Python notation) is the current range,
    and it is split after f[j].  """ 
    assert all(c!=d for c,d in zip(f,f[1:])) # Consecutive letters must be distinct.
    where = defaultdict(list)
    for i,c in enumerate(f):
        where[c].append(i) # store in which positions each letter appears

    # The tree is built by the recursive procedure "decompose". This is
    # not described in the paper.
    def decompose(i,j):
        """Either f[i]=f[j]=c and i is the first occurrence of c, or i=0,
        or the symbol f[i-1] occurs again after f[i..j] but not in f[i..j]."""
        c = f[j] # pick pieces from the end
        first = where[c][0]
        if first<i: # illegal pattern
            d = f[i-1]
            raise ValueError(c,first,j, d,i-1,f.index(d,i))
        if i==j:
            return ("Singleton",i)
        if first>i:
            return("Separate", (i,first-1,j),
                   decompose(i,first-1),
                   decompose(first,j)) # complete group
        # else: # first==i
        where[c].pop()
        prev=where[c][-1]
        return("Progressive", (i,prev,j),
               # progressive sleuth
                   decompose(i,prev),
               # plus bracket
                   ("Separate", (prev+1,j-1,j),
                      decompose(prev+1,j-1),
                      ("Singleton",j)))
    try:
        Tree =  decompose(0,len(f)-1)
    except ValueError as err:
        c,i1,i2, d,j1,j2 = err.args
        print(f"Error: Sequence {f} has interleaving faces.")
        print(f"f[{i1}]=f[{i2}]={c}, f[{j1}]=f[{j2}]={d}")
        assert f[i1]==f[i2]==c and f[j1]==f[j2]==d
        assert c!=d
        assert i1<j1<i2<j2
        return
    if show:
        print("decomposition tree for",repr(f)+":")
        print_tree(Tree,f)
    return Tree

def print_tree(t,f,level=0):
    typ = t[0]
    if typ == "Singleton":
        print("  "*level,*t, f[t[1]])
    else:
        ranges = t[1]
        print("  "*level,typ,ranges,f[ranges[0]:ranges[-1]+1])
        for sub in t[2:]:
            print_tree(sub,f,level+1)

def settle(progress_descriptor, clue_sequence, face_pattern, show_tree=False):
    """progress_descriptor is a string of '0', '1', and '?'.   
    It is assumed that it contains the added artificial '0' at the
    beginning and at the end.

    clue_sequence is a sequence of positive integers (clues)

    face_pattern gives the sequence of incident faces along the curve segments.    
    face_pattern has the same length as progress_descriptor.
    """
    f = face_pattern
    clue_bitstring = [0]
    for n in clue_sequence:
        clue_bitstring += [1]*n+[0]
    a = ''.join(str(c) for c in clue_bitstring) # convert to 0-1 string
    k = len(a)
    l = len(face_pattern)
    assert l==len(progress_descriptor)
    psi = progress_descriptor
    print(f"{a=} {k=}\n{f=} {l=}\nprogress_descriptor {psi=}")
    T = decomposition_tree(face_pattern, show=show_tree)

    Match = defaultdict(bool) # defaults to False
    # For convenience, we use dictionaries instead of arrays.

    def solve_bottom_up(t):
        if t[0]=="Singleton":
            j = t[1]
            for i in range(k):
                Match[i,i, j,j] = psi[j]=="?" or psi[j]==a[i]
        else:
            typ,split, sub1,sub2 = t
            solve_bottom_up(sub1) # visit subproblems first
            solve_bottom_up(sub2)
            
            lenient = typ=="Separate" # (strict for typ=="Progressive")
            j1,j2,j3 = split
            for i1 in range(k):
                for i3 in range(i1,k): # recursion (2) of the paper
                    Match[i1,i3, j1,j3] = (
                        any(
                            Match[i1,i2,   j1,j2] and
                            Match[i2+1,i3, j2+1,j3] and
                            (lenient or a[i2]==a[i3])
                               for i2 in range(i1,i3))
                        or
                        any(a[i2]=="0" and
                            Match[i1,i2, j1,j2] and
                            Match[i2,i3, j2+1,j3] and
                            (lenient or a[i2]==a[i3])
                               for i2 in range(i1,i3+1))
                        )
    def solve_top_down(t):
        if t[0]=="Singleton":
            j = t[1]
            if psi[j]=="?":
                for i in range(k):
                    if Extensible[i,i,j,j]:
                        Possible_letters[j,a[i]] = True
        else:
            typ,split, sub1,sub2 = t            
            lenient = typ=="Separate" # need not check extra equality
            j1,j2,j3 = split
            for i1 in range(k):
                for i3 in range(i1,k):
                    if Extensible[i1,i3, j1,j3]:
                        for i2 in range(i1,i3):
                            if (Match[i1,i2, j1,j2] and
                              Match[i2+1,i3, j2+1,j3] and
                              (lenient or a[i2]==a[i3])):                               
                                Extensible[i1,i2,   j1,j2] = True
                                Extensible[i2+1,i3, j2+1,j3] = True
                        for i2 in range(i1,i3+1):
                            if (a[i2]=="0" and
                              Match[i1,i2, j1,j2] and
                              Match[i2,i3, j2+1,j3] and
                              (lenient or a[i2]==a[i3])):
                                Extensible[i1,i2, j1,j2] = True
                                Extensible[i2,i3, j2+1,j3] = True
            solve_top_down(sub1)
            solve_top_down(sub2)

    solve_bottom_up(T)
    result = Match[0,k-1, 0,l-1]
    assert result           

    Extensible = defaultdict(bool)
    Extensible[0,k-1, 0,l-1] = True
    Possible_letters = defaultdict(bool)
    
    solve_top_down(T)
    
    Settlestring = list(psi)
    for i,ps in enumerate(psi):
        if ps=="?":
            if not Possible_letters[i,"0"]:
                Settlestring[i]="1"
            if not Possible_letters[i,"1"]:
                Settlestring[i]="0"
    return "".join(Settlestring)

####### A FEW TEST CASES #######

print("Example from the paper, Fig.4/5")
decomposition_tree("abcdefghebib", show=True)

print("\nBad example: not nested, should raise an error.")
decomposition_tree("ababcdefghebib", show=True)

print("\nExample from the paper, Fig.3 (simple nonogram)")
Psi = "0???1???10??1??0?0"
se = settle( Psi, (5,1,2), list(range(len(Psi)))) # f is trivial (basic nonogram)
print(f"settled:                '{se}'\n")
assert se=="0???11??10??1?0000" # as claimed in Fig.3

print("Example from the paper, Fig.6 (advanced nonogram)")
print("settled:", settle( Psi, (5,1,2), "abcdedfghigjdklmbn", show_tree = True))