import numpy as np
from tensorflow import keras
from tensorflow.keras import layers

def to_numpy(examples):
    examples["pixel_values"] = [np.array(image) for image in examples["image"]]
    return examples

def preprocess():
    test_dataset = load_dataset("active-learning/test_mnist")
    train_dataset = load_dataset("active-learning/labeled_samples")
    train_dataset = train_dataset.map(to_numpy, batched=True)
    test_dataset = test_dataset.map(to_numpy, batched=True)
    
    x_train = train_dataset["train"]["pixel_values"]
    y_train = train_dataset["train"]["label"]
    
    x_test = test_dataset["test"]["pixel_values"]
    y_test = test_dataset["test"]["label"]

    x_train = np.expand_dims(x_train, -1)
    x_test = np.expand_dims(x_test, -1)

    num_classes = 10
    input_shape = (28, 28, 1)
    
    y_train = keras.utils.to_categorical(y_train, num_classes)
    y_test = keras.utils.to_categorical(y_test, num_classes)

    return x_train, y_train, x_test, y_test

def training():
    x_train, y_train, x_test, y_test = preprocess()

    model = keras.Sequential(
        [
            keras.Input(shape=input_shape),
            layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
            layers.MaxPooling2D(pool_size=(2, 2)),
            layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
            layers.MaxPooling2D(pool_size=(2, 2)),
            layers.Flatten(),
            layers.Dropout(0.5),
            layers.Dense(num_classes, activation="softmax"),
        ]
    )
    
    model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
    model.fit(x_train, y_train, batch_size=128, epochs=15, validation_split=0.1)

    score = model.evaluate(x_test, y_test, verbose=0)
    print("Test loss:", score[0])
    print("Test accuracy:", score[1])