import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import keras
from keras import layers
import tensorflow_datasets as tfds

# 定义种子和固定变量
seed = 42
keras.utils.set_random_seed(seed)
AUTO = tf.data.AUTOTUNE

# 指定数据集参数
dataset_name = "tf_flowers"
batch_size = 64
image_size = (224, 224)

# 从 tfds 加载数据并分割 10% 用作测试集
(train_data, test_data), ds_info = tfds.load(
    dataset_name,
    split=["train[:90%]", "train[90%:]"],
    shuffle_files=True,
    as_supervised=True,
    with_info=True,
)

# 提取类的数量和类名列表
num_classes = ds_info.features["label"].num_classes
class_names = ds_info.features["label"].names

print(f"类的数量: {num_classes}")
print(f"类名: {class_names}")


# 将数据集转换为 NumPy 数组
def dataset_to_array(dataset, image_size, num_classes):
    images, labels = [], []
    for img, lab in dataset.as_numpy_iterator():
        images.append(tf.image.resize(img, image_size).numpy())
        labels.append(tf.one_hot(lab, num_classes))
    return np.array(images), np.array(labels)


img_train, label_train = dataset_to_array(train_data, image_size, num_classes)
img_test, label_test = dataset_to_array(test_data, image_size, num_classes)

num_train_samples = len(img_train)
print(f"训练样本数量: {num_train_samples}")

plt.figure(figsize=(16, 12))
for n in range(30):
    ax = plt.subplot(5, 6, n + 1)
    plt.imshow(img_test[n].astype("uint8"))
    plt.title(np.array(class_names)[label_test[n] == True][0])
    plt.axis("off")

# 定义图像增强模型
image_augmentation = keras.Sequential(
    [
        layers.RandomFlip(mode="horizontal"),
        layers.RandomRotation(factor=0.1),
        layers.RandomZoom(height_factor=(-0.1, -0)),
        layers.RandomContrast(factor=0.1),
    ],
)

# 将增强应用于训练图像并绘制一些示例
img_train = image_augmentation(img_train).numpy()

plt.figure(figsize=(16, 12))
for n in range(30):
    ax = plt.subplot(5, 6, n + 1)
    plt.imshow(img_train[n].astype("uint8"))
    plt.title(np.array(class_names)[label_train[n] == True][0])
    plt.axis("off")

def build_model(num_classes, img_size=image_size[0], top_dropout=0.3):
    """基于预训练的 MobileNetV2 创建一个分类器。

    参数:
        num_classes: 整数，softmax 层中使用的类别数量。
        img_size: 整数，输入图像的正方形大小（默认值为 224）。
        top_dropout: 整数，dropout 层的值（默认值为 0.3）。

    返回:
        未编译的 Keras 模型。
    """

    # 创建 MobileNetV2 的输入和预处理层
    inputs = layers.Input(shape=(img_size, img_size, 3))
    x = layers.Rescaling(scale=1.0 / 127.5, offset=-1)(inputs)
    model = keras.applications.MobileNetV2(
        include_top=False, weights="imagenet", input_tensor=x
    )

    # 冻结预训练权重
    model.trainable = False

    # 重建顶部
    x = layers.GlobalAveragePooling2D(name="avg_pool")(model.output)
    x = layers.Dropout(top_dropout)(x)
    outputs = layers.Dense(num_classes, activation="softmax")(x)
    model = keras.Model(inputs, outputs)

    print("可训练权重:", len(model.trainable_weights))
    print("不可训练权重:", len(model.non_trainable_weights))
    return model


def compile_and_train(
    model,
    training_data,
    training_labels,
    metrics=[keras.metrics.AUC(name="auc"), "acc"],
    optimizer=keras.optimizers.Adam(),
    patience=5,
    epochs=5,
):
    """编译并训练模型。

    参数:
        model: 未编译的 Keras 模型。
        training_data: NumPy 数组，训练数据。
        training_labels: NumPy 数组，训练标签。
        metrics: Keras/TF 指标，至少需要 'auc' 指标（默认是
                `[keras.metrics.AUC(name='auc'), 'acc']`）。
        optimizer: Keras/TF 优化器（默认是 `keras.optimizers.Adam()`）。
        patience: 整数，EarlyStopping 的耐心轮数（默认值为 5）。
        epochs: 整数，训练的轮数（默认值为 5）。

    返回:
        训练的 Keras 模型的训练历史。
    """

    stopper = keras.callbacks.EarlyStopping(
        monitor="val_auc",
        mode="max",
        min_delta=0,
        patience=patience,
        verbose=1,
        restore_best_weights=True,
    )

    model.compile(loss="categorical_crossentropy", optimizer=optimizer, metrics=metrics)

    history = model.fit(
        x=training_data,
        y=training_labels,
        batch_size=batch_size,
        epochs=epochs,
        validation_split=0.1,
        callbacks=[stopper],
    )
    return history


def unfreeze(model, block_name, verbose=0):
    """解冻 Keras 模型层。

    参数:
        model: Keras 模型。
        block_name: 字符串，层名称，例如 block_name = 'block4'。
                    检查提供的字符串是否在层名称中。
        verbose: 整数，0 表示静默，1 打印出层的可训练状态。

    返回:
        Keras 模型，所有在指定的 block_name 之后（包括该层）的层可训练，排除 BatchNormalization 层。
    """

    # 从 block_name 开始解冻
    set_trainable = False

    for layer in model.layers:
        if block_name in layer.name:
            set_trainable = True
        if set_trainable and not isinstance(layer, layers.BatchNormalization):
            layer.trainable = True
            if verbose == 1:
                print(layer.name, "可训练")
        else:
            if verbose == 1:
                print(layer.name, "不可训练")
    print("可训练权重:", len(model.trainable_weights))
    print("不可训练权重:", len(model.non_trainable_weights))
    return model

def train_model(training_data, training_labels):
    """训练模型如下：

    - 仅训练顶部层10个周期。
    - 解冻更深层。
    - 再训练20个周期。

    参数：
        training_data: NumPy数组，训练数据。
        training_labels: NumPy数组，训练标签。

    返回：
        模型精度。
    """

    model = build_model(num_classes)

    # 编译并训练顶部层
    history = compile_and_train(
        model,
        training_data,
        training_labels,
        metrics=[keras.metrics.AUC(name="auc"), "acc"],
        optimizer=keras.optimizers.Adam(),
        patience=3,
        epochs=10,
    )

    # 从第10块开始解冻模型
    model = unfreeze(model, "block_10")

    # 用较低的学习率编译并训练20个周期
    fine_tune_epochs = 20
    total_epochs = history.epoch[-1] + fine_tune_epochs

    history_fine = compile_and_train(
        model,
        training_data,
        training_labels,
        metrics=[keras.metrics.AUC(name="auc"), "acc"],
        optimizer=keras.optimizers.Adam(learning_rate=1e-4),
        patience=5,
        epochs=total_epochs,
    )

    # 计算模型在测试集上的精度
    _, _, acc = model.evaluate(img_test, label_test)
    return np.round(acc, 4)

def train_iteratively(sample_splits=[0.05, 0.1, 0.25, 0.5], iter_per_split=5):
    """针对多个样本分割迭代训练模型。

    参数：
        sample_splits: 列表/Numpy数组，包含训练集中用于训练的比例。
        iter_per_split: 整数，每个样本分割重复训练模型的次数。

    返回：
        所有分割和迭代的训练准确性及每个分割用于训练的样本数量。
    """
    # 训练所有样本模型并计算准确性
    train_acc = []
    sample_sizes = []

    for fraction in sample_splits:
        print(f"分割比例: {fraction}")
        # 对于每个样本大小重复训练3次
        sample_accuracy = []
        num_samples = int(num_train_samples * fraction)
        for i in range(iter_per_split):
            print(f"第 {i+1} 次运行，共 {iter_per_split} 次:")
            # 创建部分子集
            rand_idx = np.random.randint(num_train_samples, size=num_samples)
            train_img_subset = img_train[rand_idx, :]
            train_label_subset = label_train[rand_idx, :]
            # 训练模型并计算准确性
            accuracy = train_model(train_img_subset, train_label_subset)
            print(f"准确性: {accuracy}")
            sample_accuracy.append(accuracy)
        train_acc.append(sample_accuracy)
        sample_sizes.append(num_samples)
    return train_acc, sample_sizes


# 运行上述函数生成以下输出
train_acc = [
    [0.8202, 0.7466, 0.8011, 0.8447, 0.8229],
    [0.861, 0.8774, 0.8501, 0.8937, 0.891],
    [0.891, 0.9237, 0.8856, 0.9101, 0.891],
    [0.8937, 0.9373, 0.9128, 0.8719, 0.9128],
]

sample_sizes = [165, 330, 825, 1651]

def fit_and_predict(train_acc, sample_sizes, pred_sample_size):
    """拟合学习曲线以模型训练准确度结果。

    参数：
        train_acc: 列表/Numpy 数组，所有模型
                    训练拆分和迭代的训练准确度。
        sample_sizes: 列表/Numpy 数组，每个拆分用于训练的样本数量。
        pred_sample_size: 整数，基于
                        拟合学习曲线预测模型准确度的样本大小。
    """
    x = sample_sizes
    mean_acc = tf.convert_to_tensor([np.mean(i) for i in train_acc])
    error = [np.std(i) for i in train_acc]

    # 定义均方误差成本和指数曲线拟合函数
    mse = keras.losses.MeanSquaredError()

    def exp_func(x, a, b):
        return a * x**b

    # 定义变量、学习率和用于拟合的训练轮数
    a = tf.Variable(0.0)
    b = tf.Variable(0.0)
    learning_rate = 0.01
    training_epochs = 5000

    # 将指数函数拟合到数据
    for epoch in range(training_epochs):
        with tf.GradientTape() as tape:
            y_pred = exp_func(x, a, b)
            cost_function = mse(y_pred, mean_acc)
        # 获取梯度并计算调整后的权重
        gradients = tape.gradient(cost_function, [a, b])
        a.assign_sub(gradients[0] * learning_rate)
        b.assign_sub(gradients[1] * learning_rate)
    print(f"曲线拟合权重: a = {a.numpy()} 和 b = {b.numpy()}.")

    # 现在我们可以估算 pred_sample_size 的准确度
    max_acc = exp_func(pred_sample_size, a, b).numpy()

    # 打印预测的 x 值并附加到图值
    print(f"预计在 {pred_sample_size} 个样本中模型准确度为 {max_acc}。")
    x_cont = np.linspace(x[0], pred_sample_size, 100)

    # 构建图形
    fig, ax = plt.subplots(figsize=(12, 6))
    ax.errorbar(x, mean_acc, yerr=error, fmt="o", label="均值准确度和标准差。")
    ax.plot(x_cont, exp_func(x_cont, a, b), "r-", label="拟合的指数曲线。")
    ax.set_ylabel("模型分类准确度。", fontsize=12)
    ax.set_xlabel("训练样本大小。", fontsize=12)
    ax.set_xticks(np.append(x, pred_sample_size))
    ax.set_yticks(np.append(mean_acc, max_acc))
    ax.set_xticklabels(list(np.append(x, pred_sample_size)), rotation=90, fontsize=10)
    ax.yaxis.set_tick_params(labelsize=10)
    ax.set_title("学习曲线：模型准确度与样本大小。", fontsize=14)
    ax.legend(loc=(0.75, 0.75), fontsize=10)
    ax.xaxis.grid(True)
    ax.yaxis.grid(True)
    plt.tight_layout()
    plt.show()

    # 计算曲线拟合的平均绝对误差（MAE），以查看拟合数据的效果。
    # 误差越小，拟合效果越好。
    mae = keras.losses.MeanAbsoluteError()
    print(f"曲线拟合的 mae 为 {mae(mean_acc, exp_func(x, a, b)).numpy()}。")


# 我们使用整个训练集来预测模型准确度
fit_and_predict(train_acc, sample_sizes, pred_sample_size=num_train_samples)

# 现在用完整数据集训练模型以获得实际准确度
accuracy = train_model(img_train, label_train)
print(f"在{num_train_samples}张图像上达到的模型准确度为{accuracy}!")