with(LinearAlgebra):
with(ListTools):

Help:=proc()
	if args = NULL then
		print(`Find123PlusGenFunction(k,r,x)`):
	elif args = 'Find123PlusGenFunction' then
		print(`Find123PlusGenFunction(k,r,x) finds the generating function enumerating the words on the alphabet [n]^r which avoid both 123 and 1k(k-1)...2.`):
		print(`For example: Find123PlusGenFunction(6,2,x);`): 
	fi:
end:


get_L_allowed:=proc(k,r) local i,output,l,j:
	output:=[[]]:
	for i from 1 to k-1 do
		for l in output do
			if nops(l) = i-1 then
				for j from 1 to r do
					output:=[op(output), [op(l),j]]
				od:
			fi:
		od:
	od:
	output:
end:

R:=proc(L) local i,output:
	output:=[]:
	for i from 1 to nops(L) do
		if L[i] <> 0 then
			output:=[op(output), L[i]]:
		fi:
	od:
	output:
end:


#Count123PlusAvoiders(k,n,r) counts the number of words on the alphabet {1^r,2^r,...,a^r} that avoid both 123 and 1k(k-1)..32
Count123PlusAvoiders:=proc(k,n,r)
Count123PlusAvoiders1(k,r,n,0,[0$(k-2)]):
end:

#Count123PlusAvoiders1(k,r,a,b,L) counts the number of words on the alphabet {1^r,2^r,...a^r,(a+1)^b,(a+1+i)^L[i]} that avoid both 123 and 1k(k-1)..32 Note that L should have length k-2
Count123PlusAvoiders1:=proc(k,r,a,b,L) local i, lower_lim,extras,max_nonzero,output,Lp:
option remember:
#print(k,r,a,b,L):
#if nops(L) <> k-2 then
# print(`L should have k-2 elements`):
# print(L):
# return(FAIL):
#fi:
if (a=0 and b=0) then
 return(1):
fi:
if (a < 0 or b < 0) then
 return(0):
fi:
for i from 1 to nops(L) do
 if L[i] < 0 then
  return(0):
 fi:
od:
max_nonzero:=0:
for i from 1 to nops(L) do
 if L[i] > 0 then
  max_nonzero:=i:
 fi:
od:
extras:=[op(Nonzero([b,op(L)]))]:
output:=add(Count123PlusAvoiders1(k,r,i-1,r-1,[r$(a-i),op(extras),0$(i-a+k-2-nops(extras))]),i=a-k+nops(extras)+2..a):
output:=output+Count123PlusAvoiders1(k,r,a,b-1,L):
if max_nonzero >0 then
 Lp:=L:
 Lp[max_nonzero]:=Lp[max_nonzero]-1:
 output:=output+Count123PlusAvoiders1(k,r,a,b,Lp):
fi:
output:
end:
				
Find123PlusGenFunction:=proc(k,r,x) local L_allowed,n,M,row_ind,inits,new_init,j,f,b,L:
	L_allowed := get_L_allowed(k,r):
	n:=nops(L_allowed):
	M:=Array([[0$(n*r)]$(n*r)]):
	inits:=Array([[0$(n*r)]$(n*r)]):
	for b from 0 to r-1 do
		for L in L_allowed do
			M[args2ind(k,r,b,L)]:=Fun2Vec(k,r,b,L):
			new_init:=[]:
			for j from 0 to r*n-1 do
				if nops(L)+b > j or ((nops(L) + b - j) mod r) <> 0 then
					new_init:=[op(new_init), 0]:
				else
					new_init:=[op(new_init), Count123PlusAvoiders1(k,r,(j-nops(L)-b)/r,b,L)]:
				fi:
			od:
			inits[args2ind(k,r,b,L)]:=Array(new_init):
		od:
	od:
	f:=Mat2Gen(Matrix(M),Matrix(inits),x):
	return(normal(subs(x=x^(1/r), f[1]))):
end:

Fun2Vec:=proc(k,r,b,L) local Fp, M,last_pow,t,j:
M:=Array([0$nops(get_L_allowed(k,r))*r]):
t := nops(L):

if b >= 1 then
	if t =0 then
		for j from 0 to k-2-t-1 do
			M[args2ind(k,r,r-1,[r$j,b,op(L)])]:= M[args2ind(k,r,r-1,[r$j,b,op(L)])]+1:
		od:
		M[args2ind(k,r,b-1,L)] := M[args2ind(k,r,b-1,L)]+1:
	else
		for j from 0 to k-2-t-1 do
			M[args2ind(k,r,r-1,[r$j,b,op(L)])]:= M[args2ind(k,r,r-1,[r$j,b,op(L)])]+1:
		od:
		M[args2ind(k,r,b,R([op(1..-2,L),L[-1]-1]))] :=M[args2ind(k,r,b,R([op(1..-2,L),L[-1]-1]))] +1:
		M[args2ind(k,r,b-1,L)] := M[args2ind(k,r,b-1,L)]+1:
	fi:
else
	if t = 0 then
		for j from 0 to (k-2-t) do
			M[args2ind(k,r,r-1,[r$j,op(L)])]:=M[args2ind(k,r,r-1,[r$j,op(L)])] +1:
		od:
	else
		for j from 0 to (k-2-t) do
			M[args2ind(k,r,r-1,[r$j,op(L)])]:=M[args2ind(k,r,r-1, [r$j,op(L)])] +1:
		od:
		M[args2ind(k,r,0,R([op(1..-2,L),L[-1]-1]))] := M[args2ind(k,r,0,R([op(1..-2,L),L[-1]-1]))]+1:
	fi:
fi:
return(M):
end:

args2ind:=proc(k,r,b,L) local L_allowed:
	L_allowed := get_L_allowed(k,r):
	return(b*nops(L_allowed)+Search(L,L_allowed)):
end:

Mat2Gen:=proc(M,inits,x) local n,xI,deno,rat,output,i,j,ini,vars,eqns,RHS,q,ans,k:
n:=RowDimension(M):
xI:=Matrix([seq([0$(j-1),x,0$(n-j)],j=1..n)]):
deno := Determinant(xI-M):
deno := normal(subs(x=1/x,deno)*x^n):
rat := -(deno-1):
output :=[]:
for k from 1 to n do
	ini := inits[k]:
	vars := {seq(q[i], i=0..n-1)}:
	eqns := {}:
	RHS := add(q[i]*x^i,i=0..n-1)*add(rat^i,i=0..n-1):
	for j from 0 to n-1 do
		eqns := eqns union {ini[j+1] = coeff(RHS,x,j)}:
	od:
	ans := solve(eqns,vars):
	output:=[op(output), add(subs(ans,q[i])*x^i,i=0..n-1)/deno]:
od:
output:
end:

Nonzero:=proc(L) local i,output:
output:=[]:
for i from 1 to nops(L) do
 if L[i] <> 0 then
  output:=[op(output), L[i]]:
 fi:
od:
output:
end: