# Part three: predict the next token from two previous tokens.

import math
import random

training_texts = [
    "hello world",
    "hello",
    "world",
]

vocab = sorted(set("".join(training_texts)))

vocab.append("<START>")
start_id = len(vocab) - 1

vocab.append("<END>")
end_id = len(vocab) - 1

char_to_id = {}

for token_id, character in enumerate(vocab):
    char_to_id[character] = token_id

def encode(text):
    token_ids = []

    for character in text:
        token_ids.append(char_to_id[character])

    return token_ids

def decode(token_ids):
    characters = []

    for token_id in token_ids:
        characters.append(vocab[token_id])

    return "".join(characters)

# Next: build training examples with two tokens of context.
inputs = []
targets = []

for example in training_texts:
    tokens = [start_id, start_id] + encode(example) + [end_id]

    for position in range(2, len(tokens)):
        context = (tokens[position - 2], tokens[position - 1])

        inputs.append(context)
        targets.append(tokens[position])

    for context, target in zip(inputs, targets):
        first_id, second_id = context

        print(
            repr(vocab[first_id]),
            repr(vocab[second_id]),
            "->",
            repr(vocab[target]),
        )

weights = {}

for first_id in range(len(vocab)):
    for second_id in range(len(vocab)):
        context = (first_id, second_id)
        row = []

        for next_id in range(len(vocab)):
            row.append(0.0)

        weights[context] = row

print("Context rows:", len(weights))
print("Total weights:", len(weights) * len(vocab))

he_context = (char_to_id["h"], char_to_id["e"])
print("Scores after 'he':", weights[he_context])

def softmax(scores):
    largest = max(scores)
    exponentials = []

    for score in scores:
        exponentials.append(math.exp(score - largest))

    total = sum(exponentials)

    return [value / total for value in exponentials]

total_loss = 0

for context, target_id in zip(inputs, targets):
    predicted = softmax(weights[context])
    probability = predicted[target_id]

    total_loss += -math.log(probability)

average_loss = total_loss / len(targets)
print("Untrained model loss:", round(average_loss, 3))

learning_rate = 0.1

for epoch in range(100):
    for context, target_id in zip(inputs, targets):
        predicted = softmax(weights[context])

        for candidate_id in range(len(vocab)):
            target = 0.0

            if candidate_id == target_id:
                target = 1.0

            gradient = predicted[candidate_id] - target
            weights[context][candidate_id] -= learning_rate * gradient

    total_loss = 0

    for context, target_id in zip(inputs, targets):
        predicted = softmax(weights[context])
        total_loss += -math.log(predicted[target_id])

    average_loss = total_loss / len(targets)

    if (epoch + 1) % 10 == 0:
        print(
            "Epoch:", epoch + 1,
            "Loss:", round(average_loss, 3),
        )

def generate(max_new_tokens):
    context = (start_id, start_id)
    generated_ids = []

    for _ in range(max_new_tokens):
        predicted = softmax(weights[context])
        predicted[start_id] = 0.0

        next_id = random.choices(
            population=range(len(vocab)),
            weights=predicted,
            k=1,
        )[0]

        if next_id == end_id:
            break

        generated_ids.append(next_id)

        # Keep the newest two tokens.
        context = (context[1], next_id)

    return decode(generated_ids)

for _ in range(5):
    print(repr(generate(30)))