# Path: models/illumination_integration_net/illumination_integration_model.py
# class IlluminationNetwork(K.Model):
#     def __init__(self, args, img_height: int = 128, **kwargs):
#         super(IlluminationNetwork, self).__init__(**kwargs)
# 
#         self.cnn_encoder = CnnEncoder(
#             args.latent_dim, img_height=img_height, trainable=True
#         )
#         path = os.path.join("data/illumination/encoder.npy")
#         self.cnn_encoder.set_weights(np.load(path, allow_pickle=True))
# 
#         self.cnn_decoder = CnnDecoder(
#             args.latent_dim, 3, img_height=img_height, trainable=False
#         )
#         path = os.path.join("data/illumination/decoder.npy")
#         self.cnn_decoder.set_weights(np.load(path, allow_pickle=True))
# 
#         self.illumination_network = FiLMIlluminationNetwork(args, **kwargs)
# 
#         # Setup the state restoration
#         states = [
#             StateRestorationItem("illumination_network", self.illumination_network),
#             StateRestorationItem("cnn_encoder", self.cnn_encoder),
#         ]
#         self.state_restoration = StateRestoration(args, states)
# 
#         # Setup losses
#         self.global_batch_size = args.batch_size * get_num_gpus()
# 
#         self.male = multi_gpu_wrapper(
#             MeanAbsoluteLogarithmicError,
#             self.global_batch_size,
#         )
# 
#     def save(self, step):
#         # Save weights for step
#         self.state_restoration.save(step)
# 
#     def restore(self, step: Optional[int] = None) -> int:
#         # Restore weights from step or if None the latest one
#         return self.state_restoration.restore(step)
# 
#     @tf.function
#     def train_step(
#         self,
#         environment_map: tf.Tensor,
#         directions_random: tf.Tensor,
#         roughness_random: tf.Tensor,
#         targets_random: tf.Tensor,
#         optimizer: tf.keras.optimizers.Optimizer,
#         train_cnn: bool = False,
#     ):
#         tf.debugging.assert_shapes(
#             [
#                 (environment_map, ("B", "H", "W", 3)),
#                 (directions_random, ("B", "R", 3)),
#                 (roughness_random, ("B", "R", 1)),
#                 (targets_random, ("B", "R", 3)),
#             ]
#         )
# 
#         with tf.GradientTape() as tape:
#             z = self.cnn_encoder(environment_map)
#             if not train_cnn:
#                 z = tf.stop_gradient(z)
# 
#             reconstruction_full = self.illumination_network.eval_env_map(
#                 z, 0, img_height=environment_map.shape[1]
#             )
# 
#             reconstruction_random = self.illumination_network.call_multi_samples(
#                 directions_random, roughness_random, tf.stop_gradient(z)
#             )
# 
#             loss_full = self.male(environment_map, reconstruction_full) / (
#                 environment_map.shape[1] * environment_map.shape[2]
#             )
#             loss_random = (
#                 self.male(targets_random, reconstruction_random)
#                 / directions_random.shape[1]
#             )
# 
#             loss = loss_full + loss_random
# 
#         grad_vars = (
#             self.illumination_network.trainable_variables
#             + self.cnn_encoder.trainable_variables
#         )
# 
#         gradients = tape.gradient(loss, grad_vars)
# 
#         gradients, _ = tf.clip_by_global_norm(gradients, 10.0)
# 
#         optimizer.apply_gradients(zip(gradients, grad_vars))
# 
#         losses = {
#             "loss": loss,
#             "random_loss": loss_random,
#             "full_loss": loss_full,
#         }
# 
#         return losses
# 
#     @classmethod
#     def add_args(self, parser):
#         parser.add_argument("--latent_dim", type=int, default=128)
# 
#         parser.add_argument("--net_width", type=int, default=128)
#         parser.add_argument("--net_depth", type=int, default=4)
# 
#         parser.add_argument("--cond_frequency", type=int, default=2)
#         parser.add_argument("--cond_width", type=int, default=32)
# 
#         parser.add_argument("--mapping_width", type=int, default=128)
#         parser.add_argument("--mapping_depth", type=int, default=5)
# 
#         parser.add_argument("--single_offset", action="store_true")
# 
#         return parser
# 
# Path: nn_utils/tensorboard_visualization.py
# def hdr_to_tb(name, data):
#     tf.summary.image(
#         name,
#         tf.clip_by_value(  # Just for safety
#             tf.math.pow(
#                 data / (tf.ones_like(data) + data),
#                 1.0 / 2.2,
#             ),
#             0,
#             1,
#         ),  # Reinhard tone mapping
#     )


import tensorflow as tf
import dataflow.illumination_integration as dataflow
import utils.training_setup_utils as train_utils
from models.illumination_integration_net import IlluminationNetwork
from nn_utils.tensorboard_visualization import hdr_to_tb

def main(args):
    # Setup directories, logging etc.
    with train_utils.SetupDirectory(
        args,
        copy_files=True,
        main_script=__file__,
        copy_data="data/illumination",
    ):
        strategy = (
            tf.distribute.get_strategy()
            if train_utils.get_num_gpus() <= 1
            else tf.distribute.MirroredStrategy()
        )

        global_batch_size = args.batch_size * train_utils.get_num_gpus()

        print("Start reading the dataset ...")
        train_df, val_df, _ = dataflow.get_train_val_data(
            args.datadir,
            args.steps_per_epoch,
            args.val_holdout,
            global_batch_size,
            0,
            args.num_random_roughness_samples,
        )
        print("... Dataset read!")

        # Optimizer and models
        with strategy.scope():

# Path: models/illumination_integration_net/illumination_integration_model.py
# class IlluminationNetwork(K.Model):
#     def __init__(self, args, img_height: int = 128, **kwargs):
#         super(IlluminationNetwork, self).__init__(**kwargs)
# 
#         self.cnn_encoder = CnnEncoder(
#             args.latent_dim, img_height=img_height, trainable=True
#         )
#         path = os.path.join("data/illumination/encoder.npy")
#         self.cnn_encoder.set_weights(np.load(path, allow_pickle=True))
# 
#         self.cnn_decoder = CnnDecoder(
#             args.latent_dim, 3, img_height=img_height, trainable=False
#         )
#         path = os.path.join("data/illumination/decoder.npy")
#         self.cnn_decoder.set_weights(np.load(path, allow_pickle=True))
# 
#         self.illumination_network = FiLMIlluminationNetwork(args, **kwargs)
# 
#         # Setup the state restoration
#         states = [
#             StateRestorationItem("illumination_network", self.illumination_network),
#             StateRestorationItem("cnn_encoder", self.cnn_encoder),
#         ]
#         self.state_restoration = StateRestoration(args, states)
# 
#         # Setup losses
#         self.global_batch_size = args.batch_size * get_num_gpus()
# 
#         self.male = multi_gpu_wrapper(
#             MeanAbsoluteLogarithmicError,
#             self.global_batch_size,
#         )
# 
#     def save(self, step):
#         # Save weights for step
#         self.state_restoration.save(step)
# 
#     def restore(self, step: Optional[int] = None) -> int:
#         # Restore weights from step or if None the latest one
#         return self.state_restoration.restore(step)
# 
#     @tf.function
#     def train_step(
#         self,
#         environment_map: tf.Tensor,
#         directions_random: tf.Tensor,
#         roughness_random: tf.Tensor,
#         targets_random: tf.Tensor,
#         optimizer: tf.keras.optimizers.Optimizer,
#         train_cnn: bool = False,
#     ):
#         tf.debugging.assert_shapes(
#             [
#                 (environment_map, ("B", "H", "W", 3)),
#                 (directions_random, ("B", "R", 3)),
#                 (roughness_random, ("B", "R", 1)),
#                 (targets_random, ("B", "R", 3)),
#             ]
#         )
# 
#         with tf.GradientTape() as tape:
#             z = self.cnn_encoder(environment_map)
#             if not train_cnn:
#                 z = tf.stop_gradient(z)
# 
#             reconstruction_full = self.illumination_network.eval_env_map(
#                 z, 0, img_height=environment_map.shape[1]
#             )
# 
#             reconstruction_random = self.illumination_network.call_multi_samples(
#                 directions_random, roughness_random, tf.stop_gradient(z)
#             )
# 
#             loss_full = self.male(environment_map, reconstruction_full) / (
#                 environment_map.shape[1] * environment_map.shape[2]
#             )
#             loss_random = (
#                 self.male(targets_random, reconstruction_random)
#                 / directions_random.shape[1]
#             )
# 
#             loss = loss_full + loss_random
# 
#         grad_vars = (
#             self.illumination_network.trainable_variables
#             + self.cnn_encoder.trainable_variables
#         )
# 
#         gradients = tape.gradient(loss, grad_vars)
# 
#         gradients, _ = tf.clip_by_global_norm(gradients, 10.0)
# 
#         optimizer.apply_gradients(zip(gradients, grad_vars))
# 
#         losses = {
#             "loss": loss,
#             "random_loss": loss_random,
#             "full_loss": loss_full,
#         }
# 
#         return losses
# 
#     @classmethod
#     def add_args(self, parser):
#         parser.add_argument("--latent_dim", type=int, default=128)
# 
#         parser.add_argument("--net_width", type=int, default=128)
#         parser.add_argument("--net_depth", type=int, default=4)
# 
#         parser.add_argument("--cond_frequency", type=int, default=2)
#         parser.add_argument("--cond_width", type=int, default=32)
# 
#         parser.add_argument("--mapping_width", type=int, default=128)
#         parser.add_argument("--mapping_depth", type=int, default=5)
# 
#         parser.add_argument("--single_offset", action="store_true")
# 
#         return parser
# 
# Path: nn_utils/tensorboard_visualization.py
# def hdr_to_tb(name, data):
#     tf.summary.image(
#         name,
#         tf.clip_by_value(  # Just for safety
#             tf.math.pow(
#                 data / (tf.ones_like(data) + data),
#                 1.0 / 2.2,
#             ),
#             0,
#             1,
#         ),  # Reinhard tone mapping
#     )


# Path: train_illumination_net.py
import tensorflow as tf
import dataflow.illumination_integration as dataflow
import utils.training_setup_utils as train_utils
from models.illumination_integration_net import IlluminationNetwork
from nn_utils.tensorboard_visualization import hdr_to_tb


def main(args):
    # Setup directories, logging etc.
    with train_utils.SetupDirectory(
        args,
        copy_files=True,
        main_script=__file__,
        copy_data="data/illumination",
    ):
        strategy = (
            tf.distribute.get_strategy()
            if train_utils.get_num_gpus() <= 1
            else tf.distribute.MirroredStrategy()
        )

        global_batch_size = args.batch_size * train_utils.get_num_gpus()

        print("Start reading the dataset ...")
        train_df, val_df, _ = dataflow.get_train_val_data(
            args.datadir,
            args.steps_per_epoch,
            args.val_holdout,
            global_batch_size,
            0,
            args.num_random_roughness_samples,
        )
        print("... Dataset read!")

        # Optimizer and models
        with strategy.scope():

            illumination_model = IlluminationNetwork(args)