#This file contains Python code for the Algorithms in the Appendix of the Masters thesis of Tonya Reeves.

from collections import deque    # for queue management in A.3 and A.4

################################################################################################
# A.2 Verifying Nonrepetitive edge-colorings of k-ary trees
# The root is vertex 1, its children are 2,3,..,k+1. At the next level we find k+2, ... and so on
# The coloring is given as a vector whose i-th entry is the color of the edge from vertex i to its parent
# For the algorithms to work correctly the first entry should be 0 (as the root has no edge going up)
#     and the color 0 should be avoided otherwise
################################################################################################

def parent(z,k):   #Algorithm 1
#returns the parent of vertex z in k-ary tree. The root 1 has parent 0 (no parent)
    return (z+k-2)//k

def color(z,coloring): #Algorithm 2
    return coloring[z-1] 
#color(z,coloring) gives the color of the edge ABOVE vertex z in coloring.
#The shift is due to lists starting at 0 in Python




def verify(u,v,k,coloring): #Algorithm 4
#Input: ancestor u of a vertex v (i.e. u<v) in coloring of k-ary tree.
#Returns True if there is a repetition (of some length 2r) with u=v_1 and v=v_{r+1}  
    x=u
    y=v
    rep=True
    while y>u and rep==True: #If y=u we have a repetition
        rep=False
        if color(y,coloring)==color(x,coloring):
            x=parent(x,k)
            rep=True
        else:
            for i in range(2,k+2): #i will go from 2 to k+1
                z=k*(x-1)+i        #the children of x
                if color(z,coloring)==color(y,coloring):
                    x=z
                    rep=True
        y=parent(y,k)
    return rep #True for a repetition


def RepFree(k,coloring): #Algorithm 3
#Returns False if a PROPER coloring of k-ary tree on n vertices has a repetition
    repf=True
    n=len(coloring)
    for v in range(k+2,n+1): #v will go from k+2 to n
        w=parent(v,k)
        u=parent(w,k)
        while u>0 and repf==True:
            if color(w,coloring)!=color(v,coloring): #color(w) does not equal color(v)
                if verify(u,v,k,coloring)==True:
                    repf=False
                    print (u,v) #u and v are the vertices repetition occurs at
            w=parent(w,k)
            u=parent(u,k)
    return repf      #True means repitition-free

def Proper(k,coloring):  #Not in thesis
#Returns True if the coloring is proper
    n=len(coloring)    # Total number of vertices
    w=parent(n,k)      # Last vertex that has an edge going down from it
    for v in range(1,w):  #Go through all vertices (except last) that have edges going down
        D={color(x,coloring) for x in range(k*(v-1)+2,k*(v-1)+k+2)}.union({color(v,coloring)})
        if len(D)<=k:
            return [False, v]
    D={color(x,coloring) for x in range(k*(w-1)+2,n+1)}.union({color(w,coloring)}) # Check last vertex
    if len(D)<n-k*(w-1):
        return [False, w]
    return True

#####################################################################################
# A.3 Verifying if a sequence is k-special
# S is a sequence (1st entry should not be 0 as above, but should be 1st entry of sequence) 
####################################################################################

def kSpecial(k,S):  #Algorithm 5
# Returns True if S is k-special
    for u in range(len(S)):
        for v in range(u+1, len(S)):
            R=uvSpecial(k,u,v,S)
            if R!=False:
                #return ['Repetition starting at ',u,v, 'because of', R ] #
                return False
    return True    

def uvSpecial(k,u,v,S): #Algorithm 6
    # Returns True if there is a repetition s_1...s_2r in S with points  
    # u=s_1, v=s_r+1 (or u=s_2r and v=s_r) such that u<...<v
    # from v to y always move to left,
    # from u to x last move was in direction dir
    # long=True if reached x on step of size k to the left,
    #        or if started to right from u on a step of size k 
    if S[u]!= S[v]:  #verify starting point of repetition
        return False
    if v-u <= 2*k-1:
        #return 'closeness' #There is repetition of length 2 or 4
        return True
    Q=deque([[u,v,'left',False]])
    while len(Q)>0:
        x=Q[0][0]
        y=Q[0][1]
        dir=Q[0][2]
        long=Q[0][3]
        if y<=u+k and [y,dir,long]!=[u+k,'right',True]:
            return True
        if dir=='right':
            Right(k,Q,S,x,y,long,u)
        else:
            Left(k,Q,S,x,y,long,u)
        Q.popleft()
    return False


def Right(k,Q,S,x,y,long,u): #Algorithm 7
# Procedure if x was reached on step to the right 
    for b in range(1,k+1):
        for a in range(1,k+1):
            if x+a<=len(S)-1 and y-b>u: # Line added to avoid exploring improper edges
                if S[x+a]==S[y-b]: 
                    Q.append([x+a,y-b,'right',long])

def Left(k,Q,S,x,y,long,u): #Algorithm 8
# Procedure if x was reached on step to the left
    for b in range(1,k+1):
        for a in range(1,k):
            if x>=a and S[x-a]==S[y-b]:             # Added to avoid exploring improper edges
                Q.append([x-a,y-b,'left',False])
            if S[x+a]==S[y-b]  :
                Q.append([x+a,y-b,'right',False])
        if x>=k and S[x-k]==S[y-b]:                 # Added to avoid exploring improper edges
            Q.append([x-k,y-b,'left',True])
        if S[x+k]==S[y-b]and long==False and k!=1:  # Added to avoid exploring improper edges
            Q.append([x+k,y-b,'right',x==u])


##############################################################################
# A.4 Exhaustively searching for long k-special sequences
##############################################################################

def ExhaustiveSpecial(k,m): # Same as Algorithm 9, except line added to remove symmetries
    Q=[[x+1 for x in range(2*k)]]  #Start search from List [1...2k] to reduce symmetry
    maxlength=0            # (-2k)+ length of current longest maximal k-special sequence 
    vector=[0]             # n-th entry of vector is current number of maximal k-special sequences of length 2k+n
    while len(Q)>0:
        A=Q.pop()          # LIFO stack for DFS since this may reduce memory usage over BFS
        Q1=[]
        bad={A[-x-1] for x in range(2*k-1)} # Colors used already on previous 2k-1 entries
        s=min(max(A)+1,m)    ### remove symmetry by only using smallest symbol not used so far
        D={x+1 for x in range(s)}-bad       # Possible next entries
        for j in D:
            B=A+[j]
            Truthiness=True      # True if are there no k-bad sequences starting at end so far 
            for u in range(len(B)-1):              #Instead of verifying whole sequence, just check last entry
                Truthiness=Truthiness and uvSpecial(k,u,len(B)-1,B)==False
            if Truthiness:
                Q1.append(B)
        if len(Q1)==0:
            j=len(A)-2*k
            if j>maxlength:      #if we found a longer k-special sequence we have to extend vector
                for i in range(maxlength,j):
                    vector.append(0)
                maxlength=j
                print(j+2*k)     # progress report that we found a longer sequence, can comment out
            vector[j]=vector[j]+1
            # print(A) #print out newly found maximal sequence, best commented out if there are many
        else:
            Q.extend(Q1)
    return [maxlength+2*k,vector]    

##############################################################################################
#   Verifying ALL possible m-colorings of T_{k,h} for nonrepetitiveness
#      Not in thesis appendix, but adapted from A.4 algorithm
##############################################################################################


import itertools
def findsubsets(S,m):
    return set(itertools.combinations(S, m))

def ExhaustiveGeneral(k,h,m):      # Finding Nonrepetitive m-colorings for $T_{k,h}$
    print ('Are', k,'-ary trees of height', h, m,'-colorable?') 
    target=(k**(h+1)-1)/(k-1)      # Number of vertices of k-ary tree of height h
    Q=[[x for x in range(k+1)]]    # Start search from List [0,1,...,k] to reduce symmetry
    while len(Q)>0:
#        print(len(Q))              # tracking number of partial colorings that can still be checked
        A=Q.pop()                  # LIFO stack for DFS since this may reduce memory usage over BFS
        z=parent(len(A)+1,k)       # Find the vertex below which the next k colored edges will go
        bad={color(z,A)}           # Start set of colors not usable below z by making it proper        
        D={x+1 for x in range(0,m)}-bad
        for j in D:
            B=A+[j]
            Possibility=True        # True if are there no repetitions starting at end so far 
            u=parent(z,k)
            while u>0:              #Instead of verifying whole sequence, just check last entry
                Possibility=Possibility and  verifyModified(u,len(B),k,B)==False
                u=parent(u,k)
            if Possibility==False:
                bad.add(j)
        D={x+1 for x in range(m)}-bad       # Possible next colors
        if len(D)>=k:                       # We have enough colors to extend A
#            print(A, 'can be extended by',D)
            j=len(A)
            if j+k >= target:               #we found a long enough coloring
               return( 'Yes! Coloring of length', j, A,'can be extended from', D)
            else:
               Q1=findsubsets(D,k)
               for Col in Q1:
                   B=A+list(Col)
                   Q.append(B)
#        else:
#            print(A, 'only has available colors', D)
    return ('No!')


def verifyModified(u,v,k,coloring):
#Input: ancestor u of a vertex v (i.e. u<v) in coloring of k-ary tree.
#Returns True if there is a repetition (of some length 2r) with u=v_1 and v=v_{r+1}
#Modified in order to check only paths that don't use vertices >v.
    x=u
    y=v
    rep=True
    while y>u and rep==True: #If y=u we have a repetition
        rep=False
        if color(y,coloring)==color(x,coloring):
            x=parent(x,k)
            rep=True
        else:
            for i in range(2,k+2): #i will go from 2 to k+1
                z=k*(x-1)+i        #the children of x
                if z<v and color(z,coloring)==color(y,coloring):  # make sure z stays above v
                    x=z
                    rep=True
        y=parent(y,k)
    return rep #True for a repetition