"""
The most atomic way to train and run inference for a GPT in pure, dependency-free Python.
This file is the complete algorithm.
Everything else is just efficiency.

@karpathy
"""

import os       # os.path.exists
import math     # math.log, math.exp
import random   # random.seed, random.choices, random.gauss, random.shuffle
random.seed(42) # Let there be order among chaos

# Let there be a Dataset `docs`: list[str] of documents (e.g. a list of names)
if not os.path.exists('input.txt'):
    import urllib.request
    names_url = 'https://raw.githubusercontent.com/karpathy/makemore/988aa59/names.txt'
    urllib.request.urlretrieve(names_url, 'input.txt')
docs = [line.strip() for line in open('input.txt') if line.strip()]
random.shuffle(docs)
print(f"num docs: {len(docs)}")

# Let there be a Tokenizer to translate strings to sequences of integers ("tokens") and back
uchars = sorted(set(''.join(docs))) # unique characters in the dataset become token ids 0..n-1
BOS = len(uchars) # token id for a special Beginning of Sequence (BOS) token
vocab_size = len(uchars) + 1 # total number of unique tokens, +1 is for BOS
print(f"vocab size: {vocab_size}")

# Let there be Autograd to recursively apply the chain rule through a computation graph
class Value:
    __slots__ = ('data', 'grad', '_children', '_local_grads') # Python optimization for memory usage

    def __init__(self, data, children=(), local_grads=()):
        self.data = data                # scalar value of this node calculated during forward pass
        self.grad = 0                   # derivative of the loss w.r.t. this node, calculated in backward pass
        self._children = children       # children of this node in the computation graph
        self._local_grads = local_grads # local derivative of this node w.r.t. its children

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data + other.data, (self, other), (1, 1))

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data * other.data, (self, other), (other.data, self.data))

    def __pow__(self, other): return Value(self.data**other, (self,), (other * self.data**(other-1),))
    def log(self): return Value(math.log(self.data), (self,), (1/self.data,))
    def exp(self): return Value(math.exp(self.data), (self,), (math.exp(self.data),))
    def relu(self): return Value(max(0, self.data), (self,), (float(self.data > 0),))
    def __neg__(self): return self * -1
    def __radd__(self, other): return self + other
    def __sub__(self, other): return self + (-other)
    def __rsub__(self, other): return other + (-self)
    def __rmul__(self, other): return self * other
    def __truediv__(self, other): return self * other**-1
    def __rtruediv__(self, other): return other * self**-1

    def backward(self):
        topo = []
        visited = set()
        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._children:
                    build_topo(child)
                topo.append(v)
        build_topo(self)
        self.grad = 1
        for v in reversed(topo):
            for child, local_grad in zip(v._children, v._local_grads):
                child.grad += local_grad * v.grad

# ===== [added] Settings and functions for showing intermediate values =====
import sys          # [added] to switch the display output to a file
TRACE_STEP = 0      # [added] which training step (0-based) to display
TRACE_EVERY = 100   # [added] if not 0, also display every this many steps from TRACE_STEP (100 → steps 1, 101, 201, …)
                    #         set to 0 to display only TRACE_STEP once (same as before the extension)
TRACE_FILE = 'microgpt_trace_out.txt'  # [added] file to write the display to. Set to None to show it on the console
                    #         each write goes to a new file with a serial number (microgpt_trace_out_1.txt, _2.txt, …)
TRACE_COUNT = 0     # [added] number of times the display has been written (serial number in the file name)
TRACE = None        # [added] becomes a dict for recording while displaying
PAUSE = True        # [added] True: advance one stage at a time with Enter (does not stop when writing to TRACE_FILE)
CONSOLE = sys.stdout  # [added] the console (the original output)
TRACE_FP = None     # [added] the file currently being written
WEIGHT_EVERY = 100  # [added] write out all weight matrices before training and every this many steps (and at the last step). 0 = do not write
WEIGHT_FILE = 'microgpt_weights_out.txt'  # [added] file to write the weights to
                    #         each write goes to a new file with a serial number (microgpt_weights_out_1.txt, _2.txt, …)
WEIGHT_COUNT = 0    # [added] number of times the weights have been written (serial number in the file name)

def numbered(fname, k):
    """Insert a serial number before the file extension ('a.txt', 3 → 'a_3.txt')"""
    root, ext = os.path.splitext(fname)
    return f"{root}_{k}{ext}"

def is_trace_step(step):
    """Whether to display the computation of this step (0-based)"""
    if step == TRACE_STEP:
        return True
    return TRACE_EVERY > 0 and step > TRACE_STEP and (step - TRACE_STEP) % TRACE_EVERY == 0

def trace_begin():
    """Open a new file with a serial number and switch the display output to it (does nothing if TRACE_FILE is None)"""
    global TRACE_FP, TRACE_COUNT
    if TRACE_FILE is not None:
        TRACE_COUNT += 1
        TRACE_FP = open(numbered(TRACE_FILE, TRACE_COUNT), 'w', encoding='utf-8')
        sys.stdout = TRACE_FP

def trace_end(step):
    """Close the file and switch the display output back to the console"""
    global TRACE_FP
    if TRACE_FP is not None:
        TRACE_FP.close()
        TRACE_FP = None
        sys.stdout = CONSOLE
        print(f"Wrote the computation of training step {step+1} to {numbered(TRACE_FILE, TRACE_COUNT)}.")

def pause():
    global PAUSE
    if PAUSE and TRACE_FILE is None:
        a = input("  ---- [Enter] next  /  [a+Enter] show the rest without stopping ---- ")
        if a.strip().lower() == 'a':
            PAUSE = False
    print()

def val(x):
    return x.data if isinstance(x, Value) else x

def show(title, M, rows, cols, fmt="{:+.3f}", w=7, mask=None):
    """Display matrix M (a list of lists) in 2D with row and column labels"""
    print(f"■ {title}   ({len(M)} rows x {len(M[0])} cols)")
    lw = max(len(r) for r in rows) + 1
    print(" " * lw + "".join(f"{c:>{w}}" for c in cols))
    for i, (r, row) in enumerate(zip(rows, M)):
        cells = []
        for j, x in enumerate(row):
            if mask is not None and mask(i, j):
                cells.append(f"{'·':>{w}}")
            else:
                cells.append(f"{fmt.format(val(x)):>{w}}")
        print(f"{r:<{lw}}" + "".join(cells))
    print()

def T(M):
    return [list(c) for c in zip(*M)]

# Initialize the parameters, to store the knowledge of the model
n_layer = 1     # depth of the transformer neural network (number of layers)
n_embd = 16     # width of the network (embedding dimension)
block_size = 16 # maximum context length of the attention window (note: the longest name is 15 characters)
n_head = 4      # number of attention heads
head_dim = n_embd // n_head # derived dimension of each head
matrix = lambda nout, nin, std=0.08: [[Value(random.gauss(0, std)) for _ in range(nin)] for _ in range(nout)]
state_dict = {'wte': matrix(vocab_size, n_embd), 'wpe': matrix(block_size, n_embd), 'lm_head': matrix(vocab_size, n_embd)}
for i in range(n_layer):
    state_dict[f'layer{i}.attn_wq'] = matrix(n_embd, n_embd)
    state_dict[f'layer{i}.attn_wk'] = matrix(n_embd, n_embd)
    state_dict[f'layer{i}.attn_wv'] = matrix(n_embd, n_embd)
    state_dict[f'layer{i}.attn_wo'] = matrix(n_embd, n_embd)
    state_dict[f'layer{i}.mlp_fc1'] = matrix(4 * n_embd, n_embd)
    state_dict[f'layer{i}.mlp_fc2'] = matrix(n_embd, 4 * n_embd)
params = [p for mat in state_dict.values() for row in mat for p in row] # flatten params into a single list[Value]
print(f"num params: {len(params)}")

# Define the model architecture: a function mapping tokens and parameters to logits over what comes next
# Follow GPT-2, blessed among the GPTs, with minor differences: layernorm -> rmsnorm, no biases, GeLU -> ReLU
def linear(x, w):
    return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]

def softmax(logits):
    max_val = max(val.data for val in logits)
    exps = [(val - max_val).exp() for val in logits]
    total = sum(exps)
    return [e / total for e in exps]

def rmsnorm(x):
    ms = sum(xi * xi for xi in x) / len(x)
    scale = (ms + 1e-5) ** -0.5
    return [xi * scale for xi in x]

def gpt(token_id, pos_id, keys, values):
    tok_emb = state_dict['wte'][token_id] # token embedding
    pos_emb = state_dict['wpe'][pos_id] # position embedding
    x = [t + p for t, p in zip(tok_emb, pos_emb)] # joint token and position embedding
    x = rmsnorm(x) # note: not redundant due to backward pass via the residual connection
    if TRACE is not None:  # [added]
        TRACE['tok'].append(tok_emb); TRACE['pos'].append(pos_emb); TRACE['x0'].append(x)

    for li in range(n_layer):
        # 1) Multi-head Attention block
        x_residual = x
        x = rmsnorm(x)
        q = linear(x, state_dict[f'layer{li}.attn_wq'])
        k = linear(x, state_dict[f'layer{li}.attn_wk'])
        v = linear(x, state_dict[f'layer{li}.attn_wv'])
        if TRACE is not None:  # [added]
            TRACE['xn1'].append(x); TRACE['q'].append(q); TRACE['k'].append(k); TRACE['v'].append(v)
        keys[li].append(k)
        values[li].append(v)
        x_attn = []
        for h in range(n_head):
            hs = h * head_dim
            q_h = q[hs:hs+head_dim]
            k_h = [ki[hs:hs+head_dim] for ki in keys[li]]
            v_h = [vi[hs:hs+head_dim] for vi in values[li]]
            attn_logits = [sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5 for t in range(len(k_h))]
            attn_weights = softmax(attn_logits)
            if TRACE is not None:  # [added]
                TRACE['score'][h].append(attn_logits); TRACE['weight'][h].append(attn_weights)
            head_out = [sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h))) for j in range(head_dim)]
            x_attn.extend(head_out)
        x = linear(x_attn, state_dict[f'layer{li}.attn_wo'])
        if TRACE is not None:  # [added]
            TRACE['wo'].append(x)
        x = [a + b for a, b in zip(x, x_residual)]
        # 2) MLP block
        if TRACE is not None:  # [added]
            TRACE['attn_cat'].append(x_attn); TRACE['x1'].append(x)
        x_residual = x
        x = rmsnorm(x)
        x = linear(x, state_dict[f'layer{li}.mlp_fc1'])
        x = [xi.relu() for xi in x]
        if TRACE is not None:  # [added]
            TRACE['hidden'].append(x)
        x = linear(x, state_dict[f'layer{li}.mlp_fc2'])
        if TRACE is not None:  # [added]
            TRACE['fc2'].append(x)
        x = [a + b for a, b in zip(x, x_residual)]

    logits = linear(x, state_dict['lm_head'])
    if TRACE is not None:  # [added]
        TRACE['x2'].append(x)
    return logits

# ===== [added] The display routines =====
def lab(t):
    return 'BOS' if t == BOS else uchars[t]

def dump_weights(k):
    """Write all weight matrices, as they are after the update of training step k, to a new file with a serial number (k=0 is before training)"""
    global WEIGHT_COUNT
    W = lambda name: [[p.data for p in row] for row in state_dict[name]]
    D = [f"d{j}" for j in range(n_embd)]
    C = [f"c{j}" for j in range(n_embd)]
    H = [f"h{j}" for j in range(4 * n_embd)]
    V = [lab(i) for i in range(vocab_size)]
    WEIGHT_COUNT += 1
    fname = numbered(WEIGHT_FILE, WEIGHT_COUNT)
    fp = open(fname, 'w', encoding='utf-8')
    sys.stdout = fp
    print("=" * 70)
    if k == 0:
        print("Weights before training (initial values)")
    else:
        print(f"Weights after the update of training step {k}")
    print("=" * 70 + "\n")
    show("wte  token embedding (rows = character)", W('wte'), V, D)
    show("wpe  position embedding (rows = position)", W('wpe'), [f"p{i}" for i in range(block_size)], D)
    show("Wq  attn_wq (rows = component of q, cols = component of input)", W('layer0.attn_wq'), [f"q{j}" for j in range(n_embd)], D)
    show("Wk  attn_wk (rows = component of k, cols = component of input)", W('layer0.attn_wk'), [f"k{j}" for j in range(n_embd)], D)
    show("Wv  attn_wv (rows = component of v, cols = component of input)", W('layer0.attn_wv'), [f"v{j}" for j in range(n_embd)], D)
    show("Wo  attn_wo (rows = component of output, cols = component of the concatenated 4 heads)", W('layer0.attn_wo'), D, C)
    show("W1  mlp_fc1 (rows = component of hidden layer, cols = component of input)", W('layer0.mlp_fc1'), H, D)
    show("W2  mlp_fc2 (64 columns, so transposed: rows = component of hidden layer, cols = component of output)", T(W('layer0.mlp_fc2')), H, D)
    show("lm_head  output matrix (rows = character)", W('lm_head'), V, D)
    fp.close()
    sys.stdout = CONSOLE
    print(f"Wrote the weights after {k} training steps to {fname}.")

def show_forward(step, doc, tokens, n, loss):
    P = [f"{i}:{lab(tokens[i])}" for i in range(n)]           # row label = position:input character
    D = [f"d{j}" for j in range(n_embd)]
    print("\n" + "=" * 70)
    print(f"Showing the computation of training step {step+1} (learning rate etc. unchanged)")
    print("=" * 70 + "\n")
    print(f"■ 0. Input: the name used in this step = '{doc}'")
    print("  token ids :", tokens)
    print("  input → target (next character):")
    for i in range(n):
        print(f"    pos {i}: {lab(tokens[i]):>3} → {lab(tokens[i+1])}")
    print("  Each row is one position (position:character input at that position).\n")
    pause()
    show("1a. Token embedding: rows of wte (one vector per character)", TRACE['tok'], P, D); pause()
    show("1b. Position embedding: rows of wpe (one vector per position)", TRACE['pos'], P, D); pause()
    show("1c. RMSNorm of the sum = input x to the block", TRACE['x0'], P, D); pause()
    show("2a. Query q = Wq·RMSNorm(x)", TRACE['q'], P, D); pause()
    show("2b. Key   k = Wk·RMSNorm(x)", TRACE['k'], P, D); pause()
    show("2c. Value v = Wv·RMSNorm(x)", TRACE['v'], P, D)
    print(f"  * The 16 dimensions are split into groups of 4: d0–d3 go to head 0, d4–d7 to head 1, and so on.\n"); pause()
    for h in range(n_head):
        S = [row + [0.0] * (n - len(row)) for row in TRACE['score'][h]]
        W = [row + [0.0] * (n - len(row)) for row in TRACE['weight'][h]]
        fut = lambda i, j: j > i
        print(f"--- Head {h} (dimensions d{h*head_dim}–d{h*head_dim+head_dim-1}) ---")
        show(f"3a. Attention scores q·k/√{head_dim} (rows = query position, cols = key position)", S, P, P, mask=fut)
        show("3b. Attention weights softmax(scores) (each row sums to 1)", W, P, P, fmt="{:.3f}", mask=fut)
        print("  * '·' marks future positions not yet computed (their keys do not exist yet).\n")
        pause()
    show("4a. Concatenated outputs of the 4 heads (values mixed by the weights)", TRACE['attn_cat'], P, D); pause()
    show("4b. After multiplying by Wo", TRACE['wo'], P, D); pause()
    show("4c. Residual connection: 4b + block input x (1c) = output of the attention block", TRACE['x1'], P, D); pause()
    H = [f"h{j}" for j in range(4 * n_embd)]
    show("5a. MLP hidden layer ReLU(W1·RMSNorm(x)) (64 dims, so transposed: rows = component, cols = position)",
         T(TRACE['hidden']), H, P, fmt="{:.3f}", mask=lambda i, j: val(TRACE['hidden'][j][i]) == 0)
    print("  * '·' marks components set to 0 by ReLU.\n"); pause()
    show("5b. After multiplying by W2", TRACE['fc2'], P, D); pause()
    show("5c. Residual connection: 5b + 4c = output of the MLP block (output of the last layer)", TRACE['x2'], P, D); pause()
    V = [lab(i) for i in range(vocab_size)]
    show("6a. Logits (scores for the next character) (transposed: rows = character, cols = position)", T(TRACE['logits']), V, P); pause()
    show("6b. Probabilities softmax(logits) (transposed: rows = character, cols = position, each column sums to 1)", T(TRACE['probs']), V, P, fmt="{:.3f}")
    print("■ 6c. Probability of the correct character and loss −ln p")
    for i in range(n):
        tgt = tokens[i + 1]
        print(f"    pos {i}: target {lab(tgt):>3}  p = {val(TRACE['probs'][i][tgt]):.4f}   −ln p = {val(TRACE['loss'][i]):.4f}")
    print(f"  Loss (average) L = {loss.data:.4f}    (for reference, random guessing gives ln {vocab_size} = {math.log(vocab_size):.4f})\n")
    pause()

def show_backward(loss):
    D = [f"d{j}" for j in range(n_embd)]
    print("■ 7. Backpropagation loss.backward() is done. Each parameter's .grad now holds ∂L/∂w.\n")
    G = [[p.grad for p in row] for row in state_dict['layer0.attn_wq']]
    show("7a. Gradient of Wq, ∂L/∂Wq (rows = component of q, cols = component of input)", G, [f"q{j}" for j in range(n_embd)], D,
         fmt="{:+.4f}", w=8)
    pause()

def show_update(before, lr_t, step):
    D = [f"d{j}" for j in range(n_embd)]
    R = [f"q{j}" for j in range(n_embd)]
    after = [[p.data for p in row] for row in state_dict['layer0.attn_wq']]
    diff = [[a - b for a, b in zip(ra, rb)] for ra, rb in zip(after, before['layer0.attn_wq'])]
    print(f"■ 8. Update by Adam (learning rate for this step: {lr_t:.5f})\n")
    show("8a. Wq before the update", before['layer0.attn_wq'], R, D); pause()
    show("8b. Amount of update (after − before)", diff, R, D, fmt="{:+.4f}", w=8)
    if step == 0:
        print("  * In the first step, Adam's bias-corrected ratio m̂/√v̂ equals the sign of the gradient,")
        print("    so every update is ±(learning rate) (only the direction matters, not the size of the gradient).\n")
    pause()
    show("8c. Wq after the update", after, R, D)
    if TRACE_EVERY == 0:
        print("This completes one training step. From here on, training continues as in the original program, without display.\n")
    else:
        print(f"This completes the display of training step {step+1}.\n")
    pause()

# Let there be Adam, the blessed optimizer and its buffers
learning_rate, beta1, beta2, eps_adam = 0.01, 0.85, 0.99, 1e-8
m = [0.0] * len(params) # first moment buffer
v = [0.0] * len(params) # second moment buffer
if WEIGHT_EVERY > 0:  # [added] write out the weights before training
    dump_weights(0)

# Repeat in sequence
num_steps = 1000 # number of training steps
for step in range(num_steps):

    # Take single document, tokenize it, surround it with BOS special token on both sides
    doc = docs[step % len(docs)]
    tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS]
    n = min(block_size, len(tokens) - 1)
    if is_trace_step(step):  # [added] start recording
        TRACE = {key: [] for key in ['tok','pos','x0','xn1','q','k','v','attn_cat','wo','x1','hidden','fc2','x2','logits','probs','loss']}
        TRACE['score'] = [[] for _ in range(n_head)]; TRACE['weight'] = [[] for _ in range(n_head)]

    # Forward the token sequence through the model, building up the computation graph all the way to the loss
    keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
    losses = []
    for pos_id in range(n):
        token_id, target_id = tokens[pos_id], tokens[pos_id + 1]
        logits = gpt(token_id, pos_id, keys, values)
        probs = softmax(logits)
        loss_t = -probs[target_id].log()
        if TRACE is not None:  # [added]
            TRACE['logits'].append(logits); TRACE['probs'].append(probs); TRACE['loss'].append(loss_t)
        losses.append(loss_t)
    loss = (1 / n) * sum(losses) # final average loss over the document sequence. May yours be low.
    if TRACE is not None:  # [added] show intermediate values of the forward pass
        trace_begin()
        show_forward(step, doc, tokens, n, loss)

    # Backward the loss, calculating the gradients with respect to all model parameters
    loss.backward()
    if TRACE is not None:  # [added] show gradients and save the values before the update
        show_backward(loss)
        before = {name: [[p.data for p in row] for row in state_dict[name]] for name in ('layer0.attn_wq',)}

    # Adam optimizer update: update the model parameters based on the corresponding gradients
    lr_t = learning_rate * (1 - step / num_steps) # linear learning rate decay
    for i, p in enumerate(params):
        m[i] = beta1 * m[i] + (1 - beta1) * p.grad
        v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2
        m_hat = m[i] / (1 - beta1 ** (step + 1))
        v_hat = v[i] / (1 - beta2 ** (step + 1))
        p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)
        p.grad = 0
    if TRACE is not None:  # [added] show the result of the update and stop recording
        show_update(before, lr_t, step)
        trace_end(step)
        TRACE = None
    if WEIGHT_EVERY > 0 and ((step + 1) % WEIGHT_EVERY == 0 or step + 1 == num_steps):  # [added] write out the weights
        dump_weights(step + 1)

    print(f"step {step+1:4d} / {num_steps:4d} | loss {loss.data:.4f}", end='\r')

# Inference: may the model babble back to us
temperature = 0.5 # in (0, 1], control the "creativity" of generated text, low to high
print("\n--- inference (new, hallucinated names) ---")
for sample_idx in range(20):
    keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
    token_id = BOS
    sample = []
    for pos_id in range(block_size):
        logits = gpt(token_id, pos_id, keys, values)
        probs = softmax([l / temperature for l in logits])
        token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0]
        if token_id == BOS:
            break
        sample.append(uchars[token_id])
    print(f"sample {sample_idx+1:2d}: {''.join(sample)}")