class CyclicEmbeddingLoss(tf.keras.losses.Loss):
    def call(self, initial_embedding, cyclic_embedding):
        return tf.reduce_mean(math_utils.l2Norm(initial_embedding - cyclic_embedding)) class DistanceLoss(tf.keras.losses.Loss):
    def euclidean_distance_all_pairs(self, x):
        x0 = tf.expand_dims(x, 0)
        x1 = tf.expand_dims(x, 1)
        return tf.sqrt(tf.reduce_sum(tf.square(x0 - x1), axis=-1) + 1e-7)

    def call(self, y_true, y_pred):
        # Calculate the maximum distance possible in euclidean space 0 to 1
        y_true_dims = y_true.shape[-1]
        y_pred_dims = y_pred.shape[-1]

        max_y_true_dist = tf.sqrt(tf.cast(y_true_dims, tf.float32))
        max_y_pred_dist = tf.sqrt(tf.cast(y_pred_dims, tf.float32))

        y_true_distances = self.euclidean_distance_all_pairs(y_true)
        y_pred_distances = self.euclidean_distance_all_pairs(y_pred * 0.5 + 0.5)

        # Normalize the distances
        norm_y_true = y_true_distances / max_y_true_dist
        norm_y_pred = y_pred_distances / max_y_pred_dist

        distances = norm_y_true - norm_y_pred
        mask = tf.ones_like(y_true_distances) - tf.linalg.diag(
            tf.ones_like(y_true[:, 0])
        )
        return tf.reduce_mean(tf.nn.relu(distances * mask)) class LsganDiscriminatorLoss(tf.keras.losses.Loss):
    def call(self, real_logits, fake_logits):
        pos = tf.reduce_mean(tf.square(real_logits - 1))  # Move to 1
        neg = tf.reduce_mean(tf.square(fake_logits))  # Move to 0
        return 0.5 * pos + 0.5 * neg def lsganGeneratorLoss(fake_logits):
    return 0.5 * tf.reduce_mean(tf.square(fake_logits - 1)) def smoothnessLoss(x):
    return tf.reduce_mean(tf.square(x)) class BrdfDecoder(tf.keras.Model):
    def __init__(
        self,
        args,
        output_dimensions: int = 7,
        fourier_frequency_embedding: int = 10,
        activation="relu",
        final_activation="sigmoid",
        kernel_initializer="he_uniform",
        bias_initializer=tf.keras.initializers.Zeros,
        **kwargs,
    ):
        super(BrdfDecoder, self).__init__(**kwargs)

        self.latent_dim = args.latent_dim

        embedder = FourierEmbedding(
            fourier_frequency_embedding, input_dim=args.latent_dim
        )
        print(args.latent_dim, embedder.get_output_dimensionality(), args.net_w)

        net = [tf.keras.layers.InputLayer(input_shape=(args.latent_dim,))]
        if fourier_frequency_embedding >= 1:
            net.append(embedder)
            net.append(
                tf.keras.layers.Lambda(
                    lambda x: tf.ensure_shape(
                        x, [None, embedder.get_output_dimensionality()]
                    )
                )
            )

        for _ in range(args.net_d):
            net.append(
                tf.keras.layers.Dense(
                    args.net_w,
                    activation=activation,
                    kernel_initializer=kernel_initializer,
                    bias_initializer=bias_initializer,
                )
            )

        print(args.net_w)

        net.append(
            tf.keras.layers.Dense(
                output_dimensions,
                activation=final_activation,
                kernel_initializer=kernel_initializer,
                bias_initializer=bias_initializer,
            )
        )

        print(output_dimensions)

        self.net = tf.keras.Sequential(net)
        print("Done")

    @tf.function
    def call(self, x):
        x = self.net(x)

        return x

    @tf.function
    def random_sample(self, samples, mean=0.0, stddev=0.3):
        x = tf.clip_by_value(
            tf.random.normal(
                shape=(samples, self.latent_dim),
                mean=mean,
                stddev=stddev,
                dtype=tf.float32,
            ),
            -1,
            1,
        )
        return self.call(x) class BrdfDiscriminator(tf.keras.Model):
    def __init__(
        self,
        args,
        input_dimensions: int = 7,
        activation="relu",
        final_activation="linear",
        kernel_initializer="he_uniform",
        bias_initializer=tf.keras.initializers.Zeros,
        **kwargs,
    ):
        super(BrdfDiscriminator, self).__init__(**kwargs)

        net = [tf.keras.Input((input_dimensions,))]
        for _ in range(args.disc_d):
            net.append(
                tf.keras.layers.Dense(
                    args.disc_w,
                    activation=activation,
                    kernel_initializer=kernel_initializer,
                    bias_initializer=bias_initializer,
                )
            )

        net.append(
            tf.keras.layers.Dense(
                1,
                activation=final_activation,
                kernel_initializer=kernel_initializer,
                bias_initializer=bias_initializer,
            )
        )

        self.net = tf.keras.Sequential(net)

    @tf.function
    def call(self, x):
        x = self.net(x)

        return x class BrdfEncoder(tf.keras.Model):
    def __init__(
        self,
        args,
        input_dimensions: int = 7,
        activation="relu",
        final_activation="linear",
        kernel_initializer="he_uniform",
        bias_initializer=tf.keras.initializers.Zeros,
        **kwargs,
    ):
        super(BrdfEncoder, self).__init__(**kwargs)

        net = [tf.keras.Input((input_dimensions,))]
        for _ in range(args.net_d):
            net.append(
                tf.keras.layers.Dense(
                    args.net_w,
                    activation=activation,
                    kernel_initializer=kernel_initializer,
                    bias_initializer=bias_initializer,
                )
            )
        net.append(
            tf.keras.layers.Dense(
                args.latent_dim,
                activation=final_activation,
                kernel_initializer=kernel_initializer,
                bias_initializer=bias_initializer,
            )
        )

        self.net = tf.keras.Sequential(net)

    @tf.function
    def call(self, x):
        x = self.net(x)

        return x class Interpolation(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(Interpolation, self).__init__(trainable=False, **kwargs)

    @tf.function
    def call(self, x1, x2, alpha):
        return x1 * tf.cast(1 - alpha, x1.dtype) + x2 * tf.cast(alpha, x2.dtype)

    @tf.function
    def multiple_alphas(self, x1, x2, alphas):
        alphas = tf.reshape(alphas, (-1, *[1 for _ in x1.shape]))
        return x1[None, ...] * tf.cast(1 - alphas, x1.dtype) + x2[None, ...] * tf.cast(
            alphas, x2.dtype
        ) class StateRestoration:
    def __init__(self, args, to_watch: List[StateRestorationItem]):
        self.args = args
        self.to_watch = to_watch

    def save(self, i):
        for sri in self.to_watch:
            if sri.is_variable:
                StateRestoration.save_variable(self.args, sri.obj, sri.name, i)
            else:
                StateRestoration.save_weights(self.args, sri.obj, sri.name, i)

    def restore(self, step: Optional[int] = None) -> int:
        if step is None:
            name_to_step = lambda name: int(
                name.replace(self.to_watch[0].name + "_", "").replace(".npy", ""),
            )
            ckpts = sorted(
                [
                    f
                    for f in os.listdir(
                        os.path.join(self.args.basedir, self.args.expname)
                    )
                    if self.to_watch[0].name in f
                ],
                key=name_to_step,
            )
            print("Found ckpts", ckpts)
            if len(ckpts) > 0:
                step = name_to_step(ckpts[-1])
            else:
                return 0

        for sri in self.to_watch:
            print("Restoring", sri.obj, sri.name, step)
            if sri.is_variable:
                StateRestoration.restore_variable(self.args, sri.obj, sri.name, step)
            else:
                StateRestoration.restore_weights(self.args, sri.obj, sri.name, step)

        return step

    @classmethod
    def build_save_path(cls, args, prefix, i):
        return os.path.join(
            args.basedir, args.expname, "{}_{:06d}.npy".format(prefix, i)
        )

    @classmethod
    def save_weights(cls, args, net, prefix, i):
        path = StateRestoration.build_save_path(args, prefix, i)
        np.save(path, net.get_weights())
        print("saved weights", prefix, "at", path)

    @classmethod
    def save_variable(cls, args, variable, prefix, i):
        path = StateRestoration.build_save_path(args, prefix, i)
        np.save(path, variable)
        print("saved variable", prefix, "at", path)

    @classmethod
    def restore_weights(cls, args, net, prefix, i):
        path = StateRestoration.build_save_path(args, prefix, i)
        net.set_weights(np.load(path, allow_pickle=True))
        print("reloaded weights", prefix, "from", path)

    @classmethod
    def restore_variable(cls, args, variable, prefix, i):
        path = StateRestoration.build_save_path(args, prefix, i)
        data = np.load(path, allow_pickle=True)
        variable.assign(data)
        print("reloaded variable", prefix, "from", path) class StateRestorationItem:
    name: str
    obj: object
    is_variable: bool = False
from typing import Optional
from losses.brdf_ae import (
    CyclicEmbeddingLoss,
    DistanceLoss,
    LsganDiscriminatorLoss,
    lsganGeneratorLoss,
    smoothnessLoss,
)
from models.brdf_ae_net.models import BrdfDecoder, BrdfDiscriminator, BrdfEncoder
from nn_utils.utils_layers import Interpolation
from utils.training_setup_utils import StateRestoration, StateRestorationItem
import tensorflow as tf




class BrdfInterpolatingAutoEncoder(tf.keras.Model):
    def __init__(
        self,
        args,
        **kwargs,
    ):
        super(BrdfInterpolatingAutoEncoder, self).__init__(**kwargs)

        self.no_roughness = args.no_roughness
        input_dimensions = 6 if args.no_roughness else 7

        self.encoder = BrdfEncoder(args, input_dimensions=input_dimensions)
        self.decoder = BrdfDecoder(args, output_dimensions=input_dimensions)
        self.discriminator = BrdfDiscriminator(args, input_dimensions=input_dimensions)

        states = [
            StateRestorationItem("encoder", self.encoder),
            StateRestorationItem("decoder", self.decoder),
            StateRestorationItem("discriminator", self.discriminator),
        ]

        self.state_restoration = StateRestoration(args, states)