import math, random

text = open("data.txt", encoding="utf-8").read()[:20000]
N = 6

chars = sorted(set(text))
idx = {c: i for i, c in enumerate(chars)}
V = len(chars)

# N grids, one per position in the context
W = [[[random.gauss(0, 0.1) for _ in range(V)] for _ in range(V)] for _ in range(N)]
print("parameters:", N * V * V)

# each example: (list of N letter ids, next letter id)
examples = [([idx[c] for c in text[i:i+N]], idx[text[i+N]])
            for i in range(len(text) - N)]

def softmax(row):
    m = max(row)
    exps = [math.exp(x - m) for x in row]
    s = sum(exps)
    return [e / s for e in exps]

def predict(ctx):
    # add up the votes from each position's grid
    scores = [0.0] * V
    for pos, letter in enumerate(ctx):
        for j in range(V):
            scores[j] += W[pos][letter][j]
    return softmax(scores)

lr = 0.02
for step in range(30):
    total_loss = 0
    for ctx, nxt in examples:
        probs = predict(ctx)
        total_loss += -math.log(probs[nxt])
        probs[nxt] -= 1
        for pos, letter in enumerate(ctx):
            for j in range(V):
                W[pos][letter][j] -= lr * probs[j]
    if step % 1 == 0:
        print("step", step, "loss", round(total_loss / len(examples), 3))

# Generate
ctx = [idx[c] for c in text[:N]]
out = text[:N]
for _ in range(200):
    nxt = random.choices(range(V), weights=predict(ctx))[0]
    out += chars[nxt]
    ctx = ctx[1:] + [nxt]
print(out)