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

print(vocab)

char_to_id = {}

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

print(char_to_id)

def encode(text):
    token_ids = []

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

    return token_ids

print(encode("hello"))

def decode(token_ids):
    characters = []

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

    return "".join(characters)

print(decode([3, 2, 4, 4, 5]))

original = "hello world"
reconstructed = decode(encode(original))

assert reconstructed == original
print("Round trip passed!")

inputs = []
targets = []

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

    inputs.extend(tokens[:-1])
    targets.extend(tokens[1:])

for current_id, next_id in zip(inputs, targets):
    current_character = vocab[current_id]
    next_character = vocab[next_id]

    print(repr(current_character), "->", repr(next_character))

counts = {}

for current_id, next_id in zip(inputs, targets):
    if current_id not in counts:
        counts[current_id] = {}

    following = counts[current_id]

    if next_id not in following:
        following[next_id] = 0

    following[next_id] += 1

l_id = char_to_id["l"]

for next_id, count in counts[l_id].items():
    print(repr(vocab[next_id]), count)

probabilities = {}

for current_id, following in counts.items():
    total = sum(following.values())
    probabilities[current_id] = {}

    for next_id, count in following.items():
        probabilities[current_id][next_id] = count / total

for next_id, probability in probabilities[l_id].items():
    print(repr(vocab[next_id]), probability)

def generate(start_character, max_new_tokens):
    current_id = char_to_id[start_character]
    genereated_ids = [current_id]

    for _ in range(max_new_tokens):
        if current_id not in probabilities:
            break

        following = probabilities[current_id]

        next_id = random.choices(
            population=list(following.keys()),
            weights=list(following.values()),
            k=1,
        )[0]

        if next_id == end_id:
            break

        genereated_ids.append(next_id)
        current_id = next_id

    return decode(genereated_ids)
print(generate("h", 30))

total_loss = 0

for current_id, next_id in zip(inputs, targets):
    probability = probabilities[current_id][next_id]
    loss = -math.log(probability)
    total_loss += loss

    print(
        repr(vocab[current_id]),
        "->",
        repr(vocab[next_id]),
        "loss:",
        round(loss, 3)
    )

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

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]

print(softmax([0, 0, 0]))
print(softmax([2, 0, 0]))

weights = []

for current_id in range(len(vocab)):
    row = []

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

    weights.append(row)

total_loss = 0

for current_id, next_id in zip(inputs, targets):
    scores = weights[current_id]
    predicted_probabilities = softmax(scores)

    probability = predicted_probabilities[next_id]
    total_loss += -math.log(probability)

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

learning_rate = 0.1

current_id = inputs[0]  # "<START>"
next_id = targets[0]  # "h"

predicted = softmax(weights[current_id])
print("Before:", -math.log(predicted[next_id]))

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

    if candidate_id == next_id:
        target = 1.0

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

updated = softmax(weights[current_id])
print("After: ", -math.log(updated[next_id]))

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

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

            if candidate_id == next_id:
                target = 1.0

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

    # Measure loss after this epoch's updates.
    total_loss = 0

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

    average_loss = total_loss / len(targets)

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

def generate_learned(max_new_tokens):
    current_id = start_id
    generated_ids = []

    for _ in range(max_new_tokens):
        predicted = softmax(weights[current_id])

        # START is an input marker, never generated text.
        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)
        current_id = next_id

    return decode(generated_ids)

print("Counting model:", generate("h", 30))
for _ in range(5):
    print("Learned model: ", repr(generate_learned(30)))

d_id = char_to_id["d"]
predicted = softmax(weights[d_id])

print("After 'd':")

for token_id, probability in enumerate(predicted):
    print(
        repr(vocab[token_id]),
        f"{probability:.1%}",
    )