# Path: awesome_gans/datasets.py
# class Pix2PixDataSet:
#     def __init__(
#         self,
#         height=64,
#         width=64,
#         channel=3,
#         use_split=False,
#         split_rate=0.15,
#         random_state=42,
#         n_threads=8,
#         ds_path=None,
#         ds_name=None,
#         use_save=False,
#         save_type='to_h5',
#         save_file_name=None,
#     ):
# 
#         """
#         # General Settings
#         :param height: image height, default 64
#         :param width: image width, default 64
#         :param channel: image channel, default 3 (RGB)
# 
#         # Pre-Processing Option
#         :param use_split: using DataSet split, default False
#         :param split_rate: image split rate (into train & test), default 0.2
#         :param random_state: random seed for shuffling, default 42
#         :param n_threads: the number of threads for multi-threading, default 8
# 
#         # DataSet Option
#         :param ds_path: DataSet's Path, default None
#         :param ds_name: DataSet's Name, default None
#         :param use_save: saving into another file format
#         :param save_type: file format to save
#         :param save_file_name: file name to save
#         """
# 
#         self.height = height
#         self.width = width
#         self.channel = channel
#         self.image_shape = (self.height, self.width, self.channel)
# 
#         self.use_split = use_split
#         self.split_rate = split_rate
#         self.random_state = random_state
#         self.n_threads = n_threads  # change this value to the fitted value for ur system
# 
#         """
#         Expected ds_path : pix2pix/...
#         Expected ds_name : apple2orange
#         """
#         self.ds_path = ds_path
#         self.ds_name = ds_name
#         # single grid : testA, testB, (trainA, trainB)
#         # double grid : train, val, (test, sample)
#         self.ds_single_grid = [
#             'apple2orange',
#             'horse2zebra',
#             'monet2photo',
#             'summer2winter_yosemite',
#             'vangogh2photo',
#             'ae_photos',
#             'cezanne2photo',
#             'ukivoe2photo',
#             'iphone2dslr_flower',
#         ]
#         self.ds_double_grid = ['cityscapes', 'edges2handbags', 'edges2shoes', 'facades', 'maps']
# 
#         # Single Grid DatSet - the number of images
#         self.n_sg_images_a = 400
#         self.n_sg_images_b = 6287
#         # Double Grid DatSet - the number of images
#         self.n_dg_images_a = 0
#         self.n_dg_images_b = 0
# 
#         self.use_save = use_save
#         self.save_type = save_type
#         self.save_file_name = save_file_name
# 
#         try:
#             if self.use_save:
#                 assert self.save_file_name
#         except AssertionError:
#             raise AssertionError("[-] save-file/folder-name is required!")
# 
#         if self.ds_name in self.ds_single_grid:
#             self.images_a = DataSetLoader(
#                 path=self.ds_path + "/" + self.ds_name + "/trainA/",
#                 size=self.image_shape,
#                 use_save=self.use_save,
#                 name=self.save_type,
#                 save_file_name=self.save_file_name,
#                 use_image_scaling=True,
#                 image_scale='0,1',
#             ).raw_data  # numpy arrays
# 
#             self.images_b = DataSetLoader(
#                 path=self.ds_path + "/" + self.ds_name + "/trainB/",
#                 size=self.image_shape,
#                 use_save=self.use_save,
#                 name=self.save_type,
#                 save_file_name=self.save_file_name,
#                 use_image_scaling=True,
#                 image_scale='0,1',
#             ).raw_data  # numpy arrays
#             self.n_images_a = self.n_sg_images_a
#             self.n_images_b = self.n_sg_images_b
#         elif self.ds_name in self.ds_double_grid:
#             # To-Do
#             # Implement this!
#             self.n_images_a = self.n_dg_images_a
#             self.n_images_b = self.n_dg_images_b
#         else:
#             raise NotImplementedError("[-] Not Implemented yet")


import sys
import time
import tensorflow as tf
import awesome_gans.discogan.discogan_model as discogan
import awesome_gans.image_utils as iu
from awesome_gans.datasets import Pix2PixDataSet as DataSets
        # load model & graph & weight
        global_step = 0
        ckpt = tf.train.get_checkpoint_state('./model/')
        if ckpt and ckpt.model_checkpoint_path:
            # Restores from checkpoint
            model.saver.restore(s, ckpt.model_checkpoint_path)

            global_step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
            print("[+] global step : %s" % global_step, " successfully loaded")
        else:
            print('[-] No checkpoint file found')

        # initializing variables
        tf.global_variables_initializer().run()

        d_overpowered = False  # G loss > D loss * 2
        for epoch in range(paras['epoch']):
            for step in range(1000):
                offset_a = (step * paras['batch_size']) % (dataset.images_a.shape[0] - paras['batch_size'])
                offset_b = (step * paras['batch_size']) % (dataset.images_b.shape[0] - paras['batch_size'])

                # batch data set
                batch_a = dataset.images_a[offset_a : (offset_a + paras['batch_size']), :]
                batch_b = dataset.images_b[offset_b : (offset_b + paras['batch_size']), :]

                # update D network
                if not d_overpowered:
                    s.run(model.d_op, feed_dict={model.A: batch_a})

                # update G network

# Path: awesome_gans/datasets.py
# class Pix2PixDataSet:
#     def __init__(
#         self,
#         height=64,
#         width=64,
#         channel=3,
#         use_split=False,
#         split_rate=0.15,
#         random_state=42,
#         n_threads=8,
#         ds_path=None,
#         ds_name=None,
#         use_save=False,
#         save_type='to_h5',
#         save_file_name=None,
#     ):
# 
#         """
#         # General Settings
#         :param height: image height, default 64
#         :param width: image width, default 64
#         :param channel: image channel, default 3 (RGB)
# 
#         # Pre-Processing Option
#         :param use_split: using DataSet split, default False
#         :param split_rate: image split rate (into train & test), default 0.2
#         :param random_state: random seed for shuffling, default 42
#         :param n_threads: the number of threads for multi-threading, default 8
# 
#         # DataSet Option
#         :param ds_path: DataSet's Path, default None
#         :param ds_name: DataSet's Name, default None
#         :param use_save: saving into another file format
#         :param save_type: file format to save
#         :param save_file_name: file name to save
#         """
# 
#         self.height = height
#         self.width = width
#         self.channel = channel
#         self.image_shape = (self.height, self.width, self.channel)
# 
#         self.use_split = use_split
#         self.split_rate = split_rate
#         self.random_state = random_state
#         self.n_threads = n_threads  # change this value to the fitted value for ur system
# 
#         """
#         Expected ds_path : pix2pix/...
#         Expected ds_name : apple2orange
#         """
#         self.ds_path = ds_path
#         self.ds_name = ds_name
#         # single grid : testA, testB, (trainA, trainB)
#         # double grid : train, val, (test, sample)
#         self.ds_single_grid = [
#             'apple2orange',
#             'horse2zebra',
#             'monet2photo',
#             'summer2winter_yosemite',
#             'vangogh2photo',
#             'ae_photos',
#             'cezanne2photo',
#             'ukivoe2photo',
#             'iphone2dslr_flower',
#         ]
#         self.ds_double_grid = ['cityscapes', 'edges2handbags', 'edges2shoes', 'facades', 'maps']
# 
#         # Single Grid DatSet - the number of images
#         self.n_sg_images_a = 400
#         self.n_sg_images_b = 6287
#         # Double Grid DatSet - the number of images
#         self.n_dg_images_a = 0
#         self.n_dg_images_b = 0
# 
#         self.use_save = use_save
#         self.save_type = save_type
#         self.save_file_name = save_file_name
# 
#         try:
#             if self.use_save:
#                 assert self.save_file_name
#         except AssertionError:
#             raise AssertionError("[-] save-file/folder-name is required!")
# 
#         if self.ds_name in self.ds_single_grid:
#             self.images_a = DataSetLoader(
#                 path=self.ds_path + "/" + self.ds_name + "/trainA/",
#                 size=self.image_shape,
#                 use_save=self.use_save,
#                 name=self.save_type,
#                 save_file_name=self.save_file_name,
#                 use_image_scaling=True,
#                 image_scale='0,1',
#             ).raw_data  # numpy arrays
# 
#             self.images_b = DataSetLoader(
#                 path=self.ds_path + "/" + self.ds_name + "/trainB/",
#                 size=self.image_shape,
#                 use_save=self.use_save,
#                 name=self.save_type,
#                 save_file_name=self.save_file_name,
#                 use_image_scaling=True,
#                 image_scale='0,1',
#             ).raw_data  # numpy arrays
#             self.n_images_a = self.n_sg_images_a
#             self.n_images_b = self.n_sg_images_b
#         elif self.ds_name in self.ds_double_grid:
#             # To-Do
#             # Implement this!
#             self.n_images_a = self.n_dg_images_a
#             self.n_images_b = self.n_dg_images_b
#         else:
#             raise NotImplementedError("[-] Not Implemented yet")


# Path: awesome_gans/discogan/discogan_train.py
import sys
import time
import tensorflow as tf
import awesome_gans.discogan.discogan_model as discogan
import awesome_gans.image_utils as iu
from awesome_gans.datasets import Pix2PixDataSet as DataSets

        # load model & graph & weight
        global_step = 0
        ckpt = tf.train.get_checkpoint_state('./model/')
        if ckpt and ckpt.model_checkpoint_path:
            # Restores from checkpoint
            model.saver.restore(s, ckpt.model_checkpoint_path)

            global_step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
            print("[+] global step : %s" % global_step, " successfully loaded")
        else:
            print('[-] No checkpoint file found')

        # initializing variables
        tf.global_variables_initializer().run()

        d_overpowered = False  # G loss > D loss * 2
        for epoch in range(paras['epoch']):
            for step in range(1000):
                offset_a = (step * paras['batch_size']) % (dataset.images_a.shape[0] - paras['batch_size'])
                offset_b = (step * paras['batch_size']) % (dataset.images_b.shape[0] - paras['batch_size'])

                # batch data set
                batch_a = dataset.images_a[offset_a : (offset_a + paras['batch_size']), :]
                batch_b = dataset.images_b[offset_b : (offset_b + paras['batch_size']), :]

                # update D network
                if not d_overpowered:
                    s.run(model.d_op, feed_dict={model.A: batch_a})

                # update G network

                s.run(model.g_op, feed_dict={model.B: batch_b})