with(LinearAlgebra):

# return the polynomial p_a(n) := number of r x n hardinian arrays, valid for
# sufficiently large n. the outputs with a = 2 should match the "k = r"
# formulas at https://oeis.org/A253223.
hardinPoly := proc(n, a, r)
    local f, P, Q, expr, deg:

    f := combinedGF(x, a, r):
    P := numer(f):
    Q := denom(f):

    # this removes the initial (generating function) polynomial part.
    f := rem(P, Q, x) / Q:
    expr := convert(f, FormalPowerSeries):
    expr := expand(op(1, expr) / x^n):

    deg := degree(expr, x):
    expr := expand(expr / x^deg):
    expand(subs(n = n - deg, expr)):
end:

# return the generating function for a(n) := number of r x (r + n) hardinian
# arrays with parameter a.
combinedGF := proc(x, a, r)
    local symbols, rectCols, M, Mgf, squareCounts, finalStates, gfs, k, col, count, i, gf, j:
    # Two parts:
    #   - number of r x r "start" arrays;
    #   - number of ways to continue to a valid r x (r + n) array.

    symbols := {seq(0..a)}:
    rectCols := rectStates(symbols, [a, r]):

    # This matrix solves part 2. Specifically, M^n(i, j) gives the number of
    # ways to create a valid r x (n + r) array where the rth column is
    # rectCols[i] and the final column is rectCols[j].
    M := trimmedTM([a, r], rectStates, symbols, rectValidTrans, rectIsFinal):

    # Mgf(i, j) is the generating function of a(n) := M^n(i, j).
    Mgf := (IdentityMatrix(Dimension(M)[1]) - x * M)^(-1):

    # This solves part 1. Specifically, squareCounts[2][i] gives the number of
    # r x r "starter arrays" which have final column squareCounts[1][i].
    squareCounts := squareCounter(a, r):

    # States which, if we stopped at, would produce a valid hardinian array.
    finalStates := select(i -> rectIsFinal(rectCols[i]), [seq(1..nops(rectCols))]):

    gfs := []:
    for k from 1 to nops(squareCounts[1]) do
        col := squareCounts[1][k]:
        count := squareCounts[2][k]:

        # Which row of M corresponds to this column?
        i := ListTools[Search](col, rectCols):

        if i = 0 then
            print(`ahhhh`):
            return FAIL:
        fi:

        gf := count * add(Mgf[i][j], j in finalStates):
        gfs := [op(gfs), gf]:
    od:

    return simplify(x^r * add(gfs)):
end:

# return a pair [L, R] such that R[k] is the number of partially-filled
# Hardinian arrays of size r x r and parameter a with last column L[k].
squareCounter := proc(a, r)
    local symbols, M, states, retStates, retCounts, counter, k:
    if r = 1 then
        return [[[a]], [1]]:
    fi:

    symbols := {seq(0..a)}:
    M := TM([a, r], squareStates, symbols, squareValidTrans, squareIsFinal):

    states := squareStates(symbols, [a, r]):

    retStates := []:
    retCounts := []:
    counter := (M^r)[1]:
    for k from 1 to nops(states) do
        if counter[k] > 0 then
            retStates := [op(retStates), states[k][..-2]]:
            retCounts := [op(retCounts), counter[k]]:
        fi:
    od:

    return [retStates, retCounts]:
end:

squareStates := proc(symbols, params)
    local a, r, cols, states, col, j, good, i:
    a := params[1]:
    r := params[2]:
    cols := gen_columns(r,symbols):
    states := [[op(cols[1]), 0]]:
    for col in cols[2..] do
        # append the index of the column to the state
        for j from 1 to r do
            # check that this column is valid according to the going-down rule.
            good := true:
            for i from 1 to r - 1 do
                good := KDCheck(col[i], i, j, col[i + 1], i + 1, j, a):
                if not good then
                    break
                fi:
            od:

            if good then
                states := [op(states), [op(col), j]]:
            fi:
        od:
    od:

    states:
end:

KDCheck := proc(x, i1, j1, y, i2, j2, a)
    local KD1, KD2:
    KD1 := max(i1, j1) - 1:
    KD2 := max(i2, j2) - 1:

    # x = v - KD + a, and v >= 0.
    if x + KD1 - a < 0 or y + KD2 - a < 0 then
        return false:
    fi:

    if not x + KD1 - (y + KD2) in {0, -1} then
        return false:
    fi:

    return true:
end:

#is the state s an accept state?
squareIsFinal := proc(s, params)
return(evalb(s[-1] = nops(s) - 1)):
end:

# THE COLUMN OF -1 IS FOR THE PURPOSES OF INITIALIZING
gen_columns := proc(r, symbols)
local L,T,i,j,c:
L := [[(-1)$r]]:
T:=combinat:-cartprod([seq([seq(i,i=1..nops(symbols))],j=1..r)]):
while not T[finished] do 
c:=T[nextvalue]():
L := [op(L),[seq(symbols[c[i]],i=1..r)]]:
od:

L:
end:

squareValidTrans := proc(s1,s2, params)
    local i, r, a:

    r := nops(s1) - 1:
    a := params[1]:

    if s2[-1] - s1[-1] <> 1 then
        return false:
    fi:

    if s1 = [-1 $ r, 0] then
        return true:
    fi:

    for i from 1 to r do 
        if not KDCheck(s1[i], i, s1[-1], s2[i], i, s2[-1], a) then
            return false:
        elif i < r and not KDCheck(s1[i], i, s1[-1], s2[i + 1], i + 1, s2[-1], a) then
            return false:
        fi:
    od:
    return(true):
end:

rectIsFinal := proc(s)
    evalb(s[-1] = 0):
end:

rectStates := proc(symbols, params)
    local a, r, cols, states, col, good, i:
    a := params[1]:
    r := params[2]:
    cols := gen_columns(r,symbols):
    states := []:

    # cols[1] is the initial vector [-1, -1, ..., -1]. we don't want that here.
    for col in cols[2..] do
        # check that the column satisfies the going-down rule.
        good := true:
        for i from 1 to r - 1 do
            if not col[i + 1] - col[i] in {0, 1} then
                good := false:
                break:
            fi:
        od:
        if good then
            states := [op(states), col]:
        fi:
    od:
    states:
end:

rectValidTrans := proc(s1,s2, params)
    local i, r, a:

    r := nops(s1):
    a := params[1]:

    if s1 = [-1 $ r] then
        return true:
    fi:

    for i from 1 to r do 
        if not s1[i] - s2[i] in {0, 1} then
            return false:
        elif i < r and not s1[i] - s2[i + 1] in {0, 1} then
            return false:
        fi:
    od:

    true:
end:

# Transition Matrix
#   params: list of parameters passed to user defined functions
#   gen_states(sym, params): return list of states
#   sym: list of symbols
#   valid_trans(s1, s2, params): check whether s1 -> s2 is a valid transition
#   is_final(s, params): determine whether a state is final
TM := proc(params, gen_states, sym, valid_trans, is_final):
    local S,i,j,M:

    S := gen_states(sym, params):
    M:=Matrix(nops(S)+1,nops(S)+1):
    for i from 1 to nops(S) do
    for j from 1 to nops(S) do

        if valid_trans(S[i],S[j], params) then
            M[i,j] := 1:
        fi:
        
    od:
        if is_final(S[i], params) then 
            M[i,-1] := 1:
        fi:	
    od: 

    M:
end:

# TM but with the last row / column removed, so that final states are not tracked.
trimmedTM := proc(params, gen_states, sym, valid_trans, is_final)
    local M, d:
    M := TM(params, gen_states, sym, valid_trans, is_final):
    d := Dimension(M)[1]:

    DeleteColumn(DeleteRow(M, d), d):
end:

# Computes the first n terms of the sequence giving the number of valid configurations
# r is the number of rows
comp_seq := proc(r,n)
local M,N,L,i:
L := []:
M := TM(r):
N := M:
for i from 1 to n do
N := N.M:
L := [op(L),N[1][-1]]:
od:
L:
end: