if __name__ == '__main__':
    (X_train, _), (_, _) = mnist.load_data()
    #取出預設的資料集圖像
    # Rescale -1 to 1
    X_train = (X_train.astype(np.float32) - 127.5) / 127.5
    X_train = np.expand_dims(X_train, axis=3)

    gan = GAN()
    gan.train(X_train)

def __init__(self, width=28, height=28, channels=1):
        #size:28 * 28 灰階

        self.width = width
        self.height = height
        self.channels = channels

        self.shape = (self.width, self.height, self.channels)

        self.optimizer = Adam(lr=0.0002, beta_1=0.5, decay=8e-8)

        self.G = self.__generator()
        self.G.compile(loss='binary_crossentropy', optimizer=self.optimizer)

        self.D = self.__discriminator()
        self.D.compile(loss='binary_crossentropy', optimizer=self.optimizer, metrics=['accuracy'])

        self.stacked_generator_discriminator = self.__stacked_generator_discriminator()

        self.stacked_generator_discriminator.compile(loss='binary_crossentropy', optimizer=self.optimizer)

def __generator(self):
        """ Declare generator """

        model = Sequential()
        model.add(Dense(256, input_shape=(100,)))
        model.add(LeakyReLU(alpha=0.2))
        model.add(BatchNormalization(momentum=0.8))
        model.add(Dense(512))
        model.add(LeakyReLU(alpha=0.2))
        model.add(BatchNormalization(momentum=0.8))
        model.add(Dense(1024))
        model.add(LeakyReLU(alpha=0.2))
        model.add(BatchNormalization(momentum=0.8))
        model.add(Dense(self.width  * self.height * self.channels, activation='tanh'))
        model.add(Reshape((self.width, self.height, self.channels)))

        return model

def __discriminator(self):
        """ Declare discriminator """

        model = Sequential() 
        model.add(Flatten(input_shape=self.shape))
        model.add(Dense((self.width * self.height * self.channels), input_shape=self.shape))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dense(np.int64((self.width * self.height * self.channels)/2)))
        model.add(LeakyReLU(alpha=0.2))
        model.add(Dense(1, activation='sigmoid'))
        model.summary()

        return model

def __stacked_generator_discriminator(self):

        self.D.trainable = False ##固定識別器參數

        model = Sequential()
        
        #將generator和discriminator合併
        model.add(self.G) 
        model.add(self.D)
        return model

 def train(self, X_train, epochs=20000, batch = 32, save_interval = 100):

        for cnt in range(epochs): #訓練20000次

            ## train discriminator
            random_index = np.random.randint(0, len(X_train) - np.int64(batch/2)) #隨機產生index
            legit_images = X_train[random_index : random_index + np.int64(batch/2)].reshape(np.int64(batch/2), self.width, self.height, self.channels)   #從raw data隨機抽取真實圖片資料
            #隨機高斯產生分部的數值 均值為0，標準差為1，產生16個陣列，裡面包含100個數值
            gen_noise = np.random.normal(0, 1, (np.int64(batch/2), 100)) 
            syntetic_images = self.G.predict(gen_noise) #利用生成generator假圖片，給discriminator訓練用的假圖片(self.width, self.height, self.channels)

            x_combined_batch = np.concatenate((legit_images, syntetic_images)) #合併真假圖片
        #真圖片分類給1，假圖片給0，合併
            y_combined_batch = np.concatenate((np.ones((np.int64(batch/2), 1)), np.zeros((np.int64(batch/2), 1)))) 

            d_loss = self.D.train_on_batch(x_combined_batch, y_combined_batch) #訓練discriminator分辨真假圖片


            # train generator

            noise = np.random.normal(0, 1, (batch, 100)) #隨機生成高斯分部數值，32個陣列，裡面包含100個數值
            y_mislabled = np.ones((batch, 1)) #產生100個為1的陣列，將上面數值所產生當作真圖片

            g_loss = self.stacked_generator_discriminator.train_on_batch(noise, y_mislabled) #fix discriminator 訓練generator
        #d_loss代表Discriminator cross entropy loss value，值越小代表分類器分的越好
        #g_loss代表在fix Discriminator情況下， Discriminator分辨generator產生圖片的cross entropy loss value，值越小代表generator越好
        print ('epoch: %d, [Discriminator :: d_loss: %f], [ Generator :: loss: %f]' % (cnt, d_loss[0], g_loss))
            if cnt % save_interval == 0: #每100次生成手寫圖片
                self.plot_images(save2file=True, step=cnt)

def plot_images(self, save2file=False, samples=16, step=0):
        ''' Plot and generated images '''
        if not os.path.exists("./images"):
            os.makedirs("./images")
        filename = "./images/mnist_%d.png" % step
        noise = np.random.normal(0, 1, (samples, 100))

        images = self.G.predict(noise)

        plt.figure(figsize=(10, 10))

        for i in range(images.shape[0]):
            plt.subplot(4, 4, i+1)
            image = images[i, :, :, :]
            image = np.reshape(image, [self.height, self.width])
            plt.imshow(image, cmap='gray')
            plt.axis('off')
        plt.tight_layout()

        if save2file:
            plt.savefig(filename)
            plt.close('all')
        else:
            plt.show()