import tensorflow as tf
from tensorflow.keras import layers, models, Input
from tensorflow.keras.datasets import mnist

# MNISTデータセットの読み込み
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# データの前処理
train_images = train_images.reshape((60000, 28 * 28)).astype('float32') / 255
test_images = test_images.reshape((10000, 28 * 28)).astype('float32') / 255

epochs = 10
batch_size = 64

print('\n教師モデルのトレーニング')
# 教師モデルの定義とトレーニング
teacher_model = models.Sequential([
    Input(shape=(784,)),
    layers.Dense(256, activation='relu'),
    layers.Dense(128, activation='relu'),
    layers.Dense(64, activation='relu'),
    layers.Dense(10, activation='softmax')
])
teacher_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
teacher_model.fit(train_images, train_labels, epochs=epochs, batch_size=batch_size, validation_split=0.2)

print('\n生徒モデルのトレーニング')

# 生徒モデルの定義
student_model = models.Sequential([
    Input(shape=(784,)),
    layers.Dense(32, activation='relu'),
    layers.Dense(16, activation='relu'),
    layers.Dense(10, activation='softmax')
])

# 蒸留損失関数の定義
def distillation_loss(y_true, y_pred, teacher_logits, student_logits, temperature=3, alpha=0.5):

    y_true = tf.one_hot(tf.cast(y_true, tf.int32), depth=10)
    teacher_probs = tf.nn.softmax(teacher_logits / temperature)
    student_probs = tf.nn.softmax(student_logits / temperature)

    hard_loss = tf.keras.losses.categorical_crossentropy(y_true, y_pred)
    soft_loss = tf.keras.losses.categorical_crossentropy(teacher_probs, student_probs)

    loss = alpha * hard_loss + (1 - alpha) * soft_loss # , hard_loss, soft_loss

    return loss

# 生徒モデルのコンパイル
student_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# カスタムトレーニングループ
# バッチごとに教師モデルの出力を取得し、生徒モデルをトレーニングする
num_batches = len(train_images) // batch_size
temperature = 3

for epoch in range(epochs):
    print(f'Epoch {epoch+1}/{epochs}')
    # 各エポックでバッチごとにデータを取得。
    for i in range(num_batches):
        batch_start = i * batch_size
        batch_end = (i + 1) * batch_size
        batch_images = train_images[batch_start:batch_end]
        batch_labels = train_labels[batch_start:batch_end]
        
        with tf.GradientTape() as tape:
            # 教師モデルの出力（ロジット）を取得。
            teacher_logits = teacher_model(batch_images, training=False)
            # 生徒モデルの出力を取得。
            student_logits = student_model(batch_images, training=True)
            y_pred = tf.nn.softmax(student_logits)
            # カスタム損失関数を使用して損失を計算。
            loss = distillation_loss(batch_labels, y_pred, teacher_logits, student_logits, temperature)
        
        # 勾配を計算し、生徒モデルのパラメータを更新。
        gradients = tape.gradient(loss, student_model.trainable_variables)
        student_model.optimizer.apply_gradients(zip(gradients, student_model.trainable_variables))
    
    # エポックごとの評価
    test_loss, test_acc = student_model.evaluate(test_images, test_labels, verbose=0)
    print(f'Test accuracy: {test_acc}')

print('\n')
test_loss, test_acc = teacher_model.evaluate(test_images, test_labels)
print(f'Teacher Model Test accuracy: {test_acc}\n')

# 最終評価
test_loss, test_acc = student_model.evaluate(test_images, test_labels)
print(f'Student Model Test accuracy: {test_acc}')