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

import os
os.environ["PATH"] += os.pathsep + 'C:/Program Files (x86)/Graphviz2.38/bin/'

# 将一组模型合并为一个平均预测的模型方法
def get_model():
    inputs = tf.keras.Input(shape=(128,))
    outputs = layers.Dense(1)(inputs)
    return keras.Model(inputs, outputs)

model1 = get_model()
model2 = get_model()
model3 = get_model()

inputs = keras.Input(shape=(128,))
y1 = model1(inputs)
y2 = model2(inputs)
y3 = model3(inputs)
outputs = layers.average([y1, y2, y3])
ensemble_model = keras.Model(inputs=inputs, outputs=outputs)

# ResNet模型
# 使用CIFAR10 建立一个Demo ResNet
inputs = keras.Input(shape=(32, 32, 3), name="img")
x = layers.Conv2D(32, 3, activation="relu")(inputs)
x = layers.Conv2D(64, 3, activation="relu")(x)
block_1_output = layers.MaxPooling2D(3)(x)

x = layers.Conv2D(64, 3, activation="relu", padding="same")(block_1_output)
x = layers.Conv2D(64, 3, activation="relu", padding="same")(x)
block_2_output = layers.add([x, block_1_output])

x = layers.Conv2D(64, 3, activation="relu", padding="same")(block_2_output)
x = layers.Conv2D(64, 3, activation="relu", padding="same")(x)
block_3_output = layers.add([x, block_2_output])

x = layers.Conv2D(64, 3, activation="relu")(block_3_output)
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dense(256, activation="relu")(x)
x = layers.Dropout(0.5)(x)
outputs = layers.Dense(10)(x)


model = keras.Model(inputs, outputs, name="resnet-model")
model.summary()

# 绘制模型的可视化结构图
keras.utils.plot_model(model, "resnet-cifar.png", show_shapes=True)
# 训练模型
(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()

x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0
y_train = keras.utils.to_categorical(y_train, 10)
y_test = keras.utils.to_categorical(y_test, 10)

model.compile(
    optimizer=keras.optimizers.RMSprop(1e-3),
    loss=keras.losses.CategoricalCrossentropy(from_logits=True),
    metrics=["acc"],

)

# model.fit(x_train[:1000], y_train[:1000], batch_size=64, epochs=1, validation_split=0.2)

# 数据预处理
tfds.disable_progress_bar()

# 数据集格式划分(训练集(80%),验证（10%），测试(10%))
(raw_train, raw_validation, raw_test), metadata = tfds.load(
    'cifar10',
    split=['train[:80%]', 'train[80%:90%]', 'train[90%:]'],
    with_info=True,
    as_supervised=True,
)


# # 生成tf.data.Datasets对象包含(image, label)信息
# print(raw_train)
# print(raw_validation)
# print(raw_test)

# # 显示训练集中的两个图像和标签
# get_label_name = metadata.features['label'].int2str
#
# for image, label, in raw_train.take(2):
#     plt.figure()
#     plt.imshow(image)
#     plt.title(get_label_name(label))
#     plt.show()

# 使用tf.image 模块格式化图像
# 将图像调整为固定输入尺寸，然后将输入通道调整为一定范围[-1,1]
IMG_SIZE = 160  # 将所有图像调整为160*160

def format_example(image, label):
    image = tf.cast(image, tf.float32)
    image = (image / 127.5) - 1
    image = tf.image.resize(image, (IMG_SIZE, IMG_SIZE))
    return image, label

# 使用map方法将函数应用到数据集的每个模块
train = raw_train.map(format_example)
validation = raw_validation.map(format_example)
test = raw_test.map(format_example)

# 随即批处理数据
BATCH_SIZE = 32
SHUFFLE_BUFFER_SIZE = 1000

train_batches = train.shuffle(SHUFFLE_BUFFER_SIZE).batch(BATCH_SIZE)
validation_batches = validation.batch(BATCH_SIZE)
test_batches = test.batch(BATCH_SIZE)

# 检查一批数据
for image_batch, label_batch, in train_batches.take(1):
    pass
image_batch.shape
print("TensorShape", image_batch.shape)

# 使用MobileNet v2模型， 从预训练的卷积开始创建模型
# 加载预训练模型
IMG_SHAPE = (IMG_SIZE, IMG_SIZE, 3)

# 从预先训练好的模型ResNet50创建基本模型
base_model = model(input_shape=IMG_SHAPE,
                                               include_top=False,
                                               weights='resnet-model')

# 查看它对示例图像批次有什么作用
feature_batch = base_model(image_batch)
print(feature_batch.shape)

# 特征提取
# 将整个模型的可训练标记设置为False,将冻结所有层
base_model.trainable = False

# 查看基本模型结构体系
base_model = keras.Model(inputs, outputs, name="resnet-model")
base_model.summary()

# 添加分类层
# 为了从特征模块生成预测，需要对5*5的空间进行平均操作
global_average_layer = tf.keras.layers.GlobalAveragePooling2D()
feature_batch_average = global_average_layer(feature_batch)
print(feature_batch_average.shape)

# 使用tf.keras.layers.Dense,这里不需要激活层
# 因此预测将被视为logit或者原始预测值，正数表示1类，负数表示0类
prediction_layer = tf.keras.layers.Dense(1)
prediction_batch = prediction_layer(feature_batch)
print(prediction_batch.shape)

# 使用tf.keras.Sequential 模型堆叠特征提取器提取两层

model = tf.keras.Sequential(
    [
        base_model,
        global_average_layer,
        prediction_layer
    ]
)

# 编译模型 (训练模型之前必须进行编译)
base_learning_rate = 0.0001
model.compile(optimizer=tf.keras.optimizers.RMSprop(lr=base_learning_rate),
              loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
              metrics=['accuracy'])

model.summary()

# MobileNet中有2.5k参数被冻结，Dense层中有1.2k参数可训练模型
# 分为权重和偏差两个对象
len(model.trainable_variables)
print((len(model.trainable_variables)))

# 训练模型
initial_epochs = 10
validation_steps = 20

loss0, accuracy0 = model.evaluate(validation_batches, steps= validation_steps)

# 打印
print("initial loss".format(loss0))
print("initial accuracy".format(accuracy0))

history = model.fit(train_batches,
                    epochs=initial_epochs,
                    validation_data=validation_batches)
# 查看MobileNet V2 作为固定功能提取器是训练和验证专区恶行/损失的学习曲线
acc = history.history['accuracy']
val_acc = history.history['val_accuracy']

loss = history.history['loss']
val_loss = history.history['val_loss']

plt.figure(figsize=(8,8))
plt.subplot(2, 1, 1)
plt.plot(acc, label='Training Accuracy')
plt.plot(val_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.ylabel('accuracy')
plt.ylim([min(plt.ylim()), 1])
plt.title('Training and Valkidation Accuracy')

plt.subplot(2, 1, 2)
plt.plot(loss, label='training loss')
plt.plot(val_loss, label='Validation_loss')
plt.legend(loc='upper right')
plt.ylabel('Cross Entroy')
plt.title([0, 1.0])
plt.title('Training and Validation loss')
plt.xlabel('epoch')
plt.show()