# Path: model/pipeline.py
# class Pipeline(PipelineNetDef):
# 
#     def __init__(self, *args, **kwargs):
#         super().__init__(*args, **kwargs)
# 
#     # training ------------------------------------------
# 
#     def resume(self, sess):
#         self.train.groups["ae"].trainer.run(sess)
# 
#     # posterior parameters -----------------------------------------------------------------------
# 
#     def posterior_param(self, sess, output_fn, is_large=False, save_img_data=True):
#         if save_img_data:
#             output_list = self.posterior.outputs
#         else:
#             output_list = self.posterior.outputs
#             output_list.pop('data')
#         print(output_list)
#         r = OneEpochRunner(
#             self.posterior.data_module,
#             output_fn = output_fn,
#             num_samples=self.opt.test_limit,
#             output_list=output_list,
#             disp_time_interval=self.opt.disp_time_interval,
#             is_large=is_large)
#         return r.run(sess)
# 
#     """
#     def dump_posterior_param(self, pp, output_fn):
#         dir_path = os.path.dirname(output_fn)
#         if not os.path.exists(dir_path):
#             os.makedirs(dir_path)
#         # pickle.dump(pp, open(output_fn + ".p", "wb"))
#         if "vis" in pp["decoded"]:
#             pp["decoded"]["vis"] = self.output_scaled_color_image(pp["decoded"]["vis"])
#         if "data" in pp:
#             pp["data"] = self.output_scaled_color_image(pp["data"])
#         savemat(output_fn + ".mat", pp)
#     """
# 
#     # all test  -----------------------------------------------------------------------------
#     def test(self, sess, output_dir, is_snapshot=False, is_large=False, save_img_data=True):
# 
#         def nprint(*args, **kwargs):
#             if not is_snapshot:
#                 print(*args, **kwargs)
# 
#         # not necessary for snapshot
#         nprint('========== Posterior parameters')
#         self.posterior_param(sess, os.path.join(output_dir, "posterior_param"), is_large, save_img_data)
#         nprint('-- Done')


import tensorflow as tf
import os
import sys
from copy import copy
from model.pipeline import Pipeline
from tensorflow.python import debug as tf_debug
        "train_color_jittering": True,
        "train_random_mirroring": False,
        "train_batch_size": 8,
        "train_shuffle_capacity": 1000,
        "learning_rate": learning_rate,
        "max_epochs": 2000,  
        "weight_decay": 1e-6,
        "test_steps": 5000,  
        "test_limit": 200,
        "recon_weight": recon_weight,
    }
    opt["encoder_options"] = {
        "keypoint_num": num_keypoints,
        "patch_feature_dim": patch_feature_dim,
        "ae_recon_type": opt["recon_name"],  
        "keypoint_concentration_loss_weight": 100.,    
        "keypoint_axis_balancing_loss_weight": 200.,   
        "keypoint_separation_loss_weight": keypoint_separation_loss_weight,
        "keypoint_separation_bandwidth": keypoint_separation_bandwidth,
        "keypoint_transform_loss_weight": kp_transform_loss,
        "keypoint_decoding_heatmap_levels": decoding_levels,
        "keypoint_decoding_heatmap_level_base": 0.5**(1/2),
        "image_channels": 3,

    }
    opt["decoder_options"] = copy(opt["encoder_options"])


    # -------------------------------------
    model_dir = os.path.join("results/celeba_30")

# Path: model/pipeline.py
# class Pipeline(PipelineNetDef):
# 
#     def __init__(self, *args, **kwargs):
#         super().__init__(*args, **kwargs)
# 
#     # training ------------------------------------------
# 
#     def resume(self, sess):
#         self.train.groups["ae"].trainer.run(sess)
# 
#     # posterior parameters -----------------------------------------------------------------------
# 
#     def posterior_param(self, sess, output_fn, is_large=False, save_img_data=True):
#         if save_img_data:
#             output_list = self.posterior.outputs
#         else:
#             output_list = self.posterior.outputs
#             output_list.pop('data')
#         print(output_list)
#         r = OneEpochRunner(
#             self.posterior.data_module,
#             output_fn = output_fn,
#             num_samples=self.opt.test_limit,
#             output_list=output_list,
#             disp_time_interval=self.opt.disp_time_interval,
#             is_large=is_large)
#         return r.run(sess)
# 
#     """
#     def dump_posterior_param(self, pp, output_fn):
#         dir_path = os.path.dirname(output_fn)
#         if not os.path.exists(dir_path):
#             os.makedirs(dir_path)
#         # pickle.dump(pp, open(output_fn + ".p", "wb"))
#         if "vis" in pp["decoded"]:
#             pp["decoded"]["vis"] = self.output_scaled_color_image(pp["decoded"]["vis"])
#         if "data" in pp:
#             pp["data"] = self.output_scaled_color_image(pp["data"])
#         savemat(output_fn + ".mat", pp)
#     """
# 
#     # all test  -----------------------------------------------------------------------------
#     def test(self, sess, output_dir, is_snapshot=False, is_large=False, save_img_data=True):
# 
#         def nprint(*args, **kwargs):
#             if not is_snapshot:
#                 print(*args, **kwargs)
# 
#         # not necessary for snapshot
#         nprint('========== Posterior parameters')
#         self.posterior_param(sess, os.path.join(output_dir, "posterior_param"), is_large, save_img_data)
#         nprint('-- Done')


# Path: exp-ae-celeba-mafl-30.py
import tensorflow as tf
import os
import sys
from copy import copy
from model.pipeline import Pipeline
from tensorflow.python import debug as tf_debug

        "train_color_jittering": True,
        "train_random_mirroring": False,
        "train_batch_size": 8,
        "train_shuffle_capacity": 1000,
        "learning_rate": learning_rate,
        "max_epochs": 2000,  
        "weight_decay": 1e-6,
        "test_steps": 5000,  
        "test_limit": 200,
        "recon_weight": recon_weight,
    }
    opt["encoder_options"] = {
        "keypoint_num": num_keypoints,
        "patch_feature_dim": patch_feature_dim,
        "ae_recon_type": opt["recon_name"],  
        "keypoint_concentration_loss_weight": 100.,    
        "keypoint_axis_balancing_loss_weight": 200.,   
        "keypoint_separation_loss_weight": keypoint_separation_loss_weight,
        "keypoint_separation_bandwidth": keypoint_separation_bandwidth,
        "keypoint_transform_loss_weight": kp_transform_loss,
        "keypoint_decoding_heatmap_levels": decoding_levels,
        "keypoint_decoding_heatmap_level_base": 0.5**(1/2),
        "image_channels": 3,

    }
    opt["decoder_options"] = copy(opt["encoder_options"])


    # -------------------------------------
    model_dir = os.path.join("results/celeba_30")

    vp = Pipeline(None, opt, model_dir=model_dir)