# --- built in ---
import os
import sys
import time
import math
import logging
import functools

# --- 3rd party ---
import numpy as np
import tensorflow as tf

# --- my module ---

__all__ = [
    'ToyMLP',
    'Energy',
    'Trainer',
]

# --- primitives ---

class ToyMLP(tf.keras.Model):
    def __init__(
        self,
        input_dim=2,
        output_dim=1,
        units=[300, 300],
        silu=True,
        dropout=False
    ):
        """Toy MLP from
        https://github.com/ermongroup/ncsn/blob/master/runners/toy_runner.py#L198
        Args:
            input_dim (int, optional): input dimensions. Defaults to 2.
            output_dim (int, optional): output dimensions. Defaults to 1.
            units (list, optional): hidden units. Defaults to [300, 300].
            silu (bool, optional): use silu as activation function. Set False to use
                soft plus instead. Defaults to True.
            dropout (bool, optional): use dropout layers. Defaults to False.
        """
        super().__init__()
        layers = [tf.keras.layers.Flatten()]
        for unit in units:
            layers.extend([
                tf.keras.layers.Dense(unit),
                tf.keras.layers.Activation('silu' if silu else 'softplus'),
                tf.keras.layers.Dropout(.5) if dropout else tf.keras.layers.Layer()
            ])
        layers.append(tf.keras.layers.Dense(output_dim))

        self.net = tf.keras.Sequential(layers)
        # forwrad dummy tensor
        dummy = tf.keras.Input((input_dim,), dtype=tf.float32)
        self.net(dummy)

    def call(self, x, training=True):
        return self.net(x, training=training)

# --- energy model ---
class Energy(tf.keras.Model):
    def __init__(self, net):
        """A simple energy model

        Args:
            net (tf.keras.Model): An energy function, the output shape of
                the energy function should be (b, 1). The score is
                computed by grad(-E(x))
        """
        super().__init__()
        self.net = net

    def call(self, x, training=True):
        return self.net(x, training=training)

    def score(self, x, sigma=None, training=True):
        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(x)
            logp = -tf.math.reduce_sum(self.net(x, training=training))
        grad = tape.gradient(logp, x)
        return grad

class Trainer():
    def __init__(
        self,
        model,
        learning_rate = 1e-3,
        clipnorm = 100.,
        n_slices = 1,
        loss_type = 'ssm-vr',
        noise_type = 'gaussian',
    ):
        """Energy based model trainer

        Args:
            model (nn.Module): energy-based model
            learning_rate (float, optional): learning rate. Defaults to 1e-4.
            clipnorm (float, optional): gradient clip. Defaults to 100..
            n_slices (int, optional): number of slices for sliced score matching loss.
                Defaults to 1.
            loss_type (str, optional): type of loss. Can be 'ssm-vr', 'ssm', 'deen',
                'dsm'. Defaults to 'ssm-vr'.
            noise_type (str, optional): type of noise. Can be 'radermacher', 'sphere'
                or 'gaussian'. Defaults to 'radermacher'.
        """
        self.model = model
        self.learning_rate = learning_rate
        self.clipnorm = clipnorm
        self.n_slices = n_slices
        self.loss_type = loss_type.lower()
        self.noise_type = noise_type.lower()
        # setup optimizer
        self.optimizer = tf.keras.optimizers.Adam(
            learning_rate=learning_rate,
            clipnorm=clipnorm
        )
        
        self.num_gradsteps = 0
        self.num_epochs = 0
        self.progress = 0
        self.tb_writer = None

    def ssm_loss(self, x, v):
        """SSM loss from
        Sliced Score Matching: A Scalable Approach to Density and Score Estimation

        The loss is computed as
        s = -dE(x)/dx
        loss = vT*(ds/dx)*v + 1/2*(vT*s)^2

        Args:
            x (tf.Tensor): input samples
            v (tf.Tensor): sampled noises

        Returns:
            SSM loss
        """
        x = tf.repeat(x, self.n_slices, axis=0) # (n_slices*b, ...)
        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(x)
            score = self.model.score(x) # (n_slices*b, ...)
            sv    = tf.math.reduce_sum(score * v) # ()
        gsv = tape.gradient(sv, x) # (n_slices*b, ...)
        loss1 = tf.math.reduce_sum(score * v, axis=-1) ** 2 * 0.5 # (n_slices*b,)
        loss2 = tf.math.reduce_sum(v * gsv, axis=-1) # (n_slices*b,)
        loss = tf.math.reduce_mean(loss1 + loss2) # ()
        return loss
    
    def ssm_vr_loss(self, x, v):
        """SSM-VR (variance reduction) loss from
        Sliced Score Matching: A Scalable Approach to Density and Score Estimation

        The loss is computed as
        s = -dE(x)/dx
        loss = vT*(ds/dx)*v + 1/2*||s||^2

        Args:
            x (tf.Tensor): input samples
            v (tf.Tensor): sampled noises

        Returns:
            SSM-VR loss
        """
        x = tf.repeat(x, self.n_slices, axis=0) # (n_slices*b, ...)
        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(x)
            score = self.model.score(x) # (n_slices*b, ...)
            sv    = tf.math.reduce_sum(score * v) # ()
        gsv = tape.gradient(sv, x) # (n_slices*b, ...)
        loss1 = tf.norm(score, axis=-1) ** 2 * 0.5 # (n_slices*b,)
        loss2 = tf.math.reduce_sum(v*gsv, axis=-1) # (n_slices*b,)
        loss = tf.math.reduce_mean(loss1 + loss2) # ()
        return loss
    
    def deen_loss(self, x, v, sigma=0.1):
        """DEEN loss from
        Deep Energy Estimator Networks

        The loss is computed as
        x_ = x + v   # noisy samples
        s = -dE(x_)/dx_
        loss = 1/2*||x - x_ + sigma^2*s||^2

        Args:
            x (tf.Tensor): input samples
            v (tf.Tensor): sampled noises
            sigma (int, optional): noise scale. Defaults to 1.

        Returns:
            DEEN loss
        """
        v = v * sigma
        x_ = x + v
        s = (sigma ** 2) * self.model.score(x_)
        loss = tf.norm(s+v, axis=-1)**2
        loss = 0.5 * tf.math.reduce_mean(loss)
        return loss

    def dsm_loss(self, x, v, sigma=0.1):
        """DSM loss from
        A Connection Between Score Matching
            and Denoising Autoencoders

        The loss is computed as
        x_ = x + v   # noisy samples
        s = -dE(x_)/dx_
        loss = 1/2*||s + (x-x_)/sigma^2||^2

        Args:
            x (tf.Tensor): input samples
            v (tf.Tensor): sampled noises
            sigma (float, optional): noise scale. Defaults to 0.1.
        
        Returns:
            DSM loss
        """
        v = v * sigma
        x_ = x + v
        s = self.model.score(x_)
        loss = tf.norm(s + v/(sigma**2), axis=-1) ** 2
        loss = 0.5 * tf.math.reduce_mean(loss)
        return loss

    def get_random_noise(self, x, n_slices=None):
        """Sampling random noises

        Args:
            x (tf.Tensor): input samples
            n_slices (int, optional): number of slices. Defaults to None.

        Returns:
            tf.Tensor: sampled noises
        """
        if n_slices is None:
            v = tf.random.normal(x.shape)
        else:
            v = tf.random.normal((n_slices, *x.shape))
            v = tf.reshape(v, (-1, *v.shape[2:])) # (n_slices*b, 2)

        if self.noise_type == 'radermacher':
            v = tf.math.sign(v)
        elif self.noise_type == 'sphere':
            v = v/tf.norm(v, axis=-1, keepdims=True) * np.sqrt(v.shape[-1])
        elif self.noise_type == 'gaussian':
            pass
        else:
            raise NotImplementedError(
                f"Noise type '{self.noise_type}' not implemented."
            )
        return v
            
    
    def get_loss(self, x, v=None):
        """Compute loss

        Args:
            x (tf.Tensor): input samples
            v (tf.Tensor, optional): sampled noises. Defaults to None.

        Returns:
            loss
        """
        if self.loss_type == 'ssm-vr':
            v = self.get_random_noise(x, self.n_slices)
            loss = self.ssm_vr_loss(x, v)
        elif self.loss_type == 'ssm':
            v = self.get_random_noise(x, self.n_slices)
            loss = self.ssm_loss(x, v)
        elif self.loss_type == 'deen':
            v = self.get_random_noise(x, None)
            loss = self.deen_loss(x, v)
        elif self.loss_type == 'dsm':
            v = self.get_random_noise(x, None)
            loss = self.dsm_loss(x, v)
        else:
            raise NotImplementedError(
                f"Loss type '{self.loss_type}' not implemented."
            )

        return loss
    
    @tf.function
    def train_step(self, batch, update=True):
        """Train one batch

        Args:
            batch (dict): batch data
            update (bool, optional): whether to update networks. 
                Defaults to True.

        Returns:
            loss
        """
        x = batch['samples']
        # move inputs to device
        x = tf.convert_to_tensor(x, dtype=tf.float32)
        vars = self.model.variables
        with tf.GradientTape() as tape:
            tape.watch(vars)
            # compute losses
            loss = self.get_loss(x)
        # update model
        if update:
            # compute gradients
            grads = tape.gradient(loss, vars)
            self.optimizer.apply_gradients(zip(grads, vars))
        return loss

    def train(self, dataset, batch_size):
        """Train one epoch

        Args:
            dataset (tf.data.Dataset): Tensorflow dataset
            batch_size (int): batch size

        Returns:
            np.ndarray: mean loss
        """        
        all_losses = []
        dataset = dataset.batch(batch_size)
        for batch_data in dataset.as_numpy_iterator():
            sample_batch = {
                'samples': batch_data
            }
            loss = self.train_step(sample_batch)
            self.num_gradsteps += 1
            all_losses.append(loss)
        m_loss = np.mean(all_losses).astype(np.float32)
        return m_loss

    def eval(self, dataset, batch_size):
        """Eval one epoch

        Args:
            dataset (tf.data.Dataset): Tensorflow dataset
            batch_size (int): batch size

        Returns:
            np.ndarray: mean loss
        """        
        all_losses = []
        dataset = dataset.batch(batch_size)
        for batch_data in dataset.as_numpy_iterator():
            sample_batch = {
                'samples': batch_data
            }
            loss = self.train_step(sample_batch, update=False)
            all_losses.append(loss)
        m_loss = np.mean(all_losses).astype(np.float32)
        return m_loss

    def learn(
        self,
        train_dataset,
        eval_dataset = None,
        n_epochs = 5,
        batch_size = 100,
        log_freq = 1,
        eval_freq = 1,
        vis_freq = 1,
        vis_callback = None,
        tb_logdir = None
    ):
        """Train the model

        Args:
            train_dataset (tf.data.Dataset): training dataset
            eval_dataset (tf.data.Dataset, optional): evaluation dataset.
                Defaults to None.
            n_epochs (int, optional): number of epochs to train. Defaults to 5.
            batch_size (int, optional): batch size. Defaults to 100.
            log_freq (int, optional): logging frequency (epoch). Defaults to 1.
            eval_freq (int, optional): evaluation frequency (epoch). Defaults to 1.
            vis_freq (int, optional): visualizing frequency (epoch). Defaults to 1.
            vis_callback (callable, optional): visualization function. Defaults to None.
            tb_logdir (str, optional): path to tensorboard files. Defaults to None.

        Returns:
            self
        """        
        if tb_logdir is not None:
            self.tb_writer = tf.summary.create_file_writer(tb_logdir)

        # initialize
        time_start = time.time()
        time_spent = 0
        total_epochs = n_epochs

        for epoch in range(n_epochs):
            self.num_epochs += 1
            self.progress = float(self.num_epochs) / float(n_epochs)
            # train one epoch
            loss = self.train(train_dataset, batch_size)
            # write tensorboard
            if self.tb_writer is not None:
                with self.tb_writer.as_default():
                    tf.summary.scalar(f'train/loss', loss, step=self.num_epochs)
            
            if (log_freq is not None) and (self.num_epochs % log_freq == 0):
                logging.info(
                    f"[Epoch {self.num_epochs}/{total_epochs}]: loss: {loss}"
                )

            if (eval_dataset is not None) and (self.num_epochs % eval_freq == 0):
                # evaluate
                eval_loss = self.eval(eval_dataset, batch_size)

                if self.tb_writer is not None:
                    with self.tb_writer.as_default():
                        tf.summary.scalar(f'eval/loss', eval_loss, step=self.num_epochs)
                
                logging.info(
                    f"[Eval {self.num_epochs}/{total_epochs}]: loss: {eval_loss}"
                )

            if (vis_callback is not None) and (self.num_epochs % vis_freq == 0):
                logging.debug("Visualizing")
                vis_callback(self)
        return self

# --- built in ---
import os
import sys
import time
import math
import logging
import functools

# --- 3rd party ---
import numpy as np
import tensorflow as tf

# --- my module ---

__all__ = [
    'ToyMLP',
    'Energy',
    'Trainer',
]

# --- primitives ---

class ToyMLP(tf.keras.Model):
    def __init__(
        self,
        input_dim=2,
        output_dim=1,
        units=[300, 300],
        silu=True,
        dropout=False
    ):
        """Toy MLP from
        https://github.com/ermongroup/ncsn/blob/master/runners/toy_runner.py#L198
        Args:
            input_dim (int, optional): input dimensions. Defaults to 2.
            output_dim (int, optional): output dimensions. Defaults to 1.
            units (list, optional): hidden units. Defaults to [300, 300].
            silu (bool, optional): use silu as activation function. Set False to use
                soft plus instead. Defaults to True.
            dropout (bool, optional): use dropout layers. Defaults to False.
        """
        super().__init__()
        layers = [tf.keras.layers.Flatten()]
        for unit in units:
            layers.extend([
                tf.keras.layers.Dense(unit),
                tf.keras.layers.Activation('silu' if silu else 'softplus'),
                tf.keras.layers.Dropout(.5) if dropout else tf.keras.layers.Layer()
            ])
        layers.append(tf.keras.layers.Dense(output_dim))

        self.net = tf.keras.Sequential(layers)
        # forwrad dummy tensor
        dummy = tf.keras.Input((input_dim,), dtype=tf.float32)
        self.net(dummy)

    def call(self, x, training=True):
        return self.net(x, training=training)

# --- energy model ---
class Energy(tf.keras.Model):
    def __init__(self, net):
        """A simple energy model

        Args:
            net (tf.keras.Model): An energy function, the output shape of
                the energy function should be (b, 1). The score is
                computed by grad(-E(x))
        """
        super().__init__()
        self.net = net

    def call(self, x, training=True):
        return self.net(x, training=training)

    def score(self, x, sigma=None, training=True):
        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(x)
            logp = -tf.math.reduce_sum(self.net(x, training=training))
        grad = tape.gradient(logp, x)
        return grad

class Trainer():
    def __init__(
        self,
        model,
        learning_rate = 1e-3,
        clipnorm = 100.,
        n_slices = 1,
        loss_type = 'ssm-vr',
        noise_type = 'gaussian',
    ):
        """Energy based model trainer

        Args:
            model (nn.Module): energy-based model
            learning_rate (float, optional): learning rate. Defaults to 1e-4.
            clipnorm (float, optional): gradient clip. Defaults to 100..
            n_slices (int, optional): number of slices for sliced score matching loss.
                Defaults to 1.
            loss_type (str, optional): type of loss. Can be 'ssm-vr', 'ssm', 'deen',
                'dsm'. Defaults to 'ssm-vr'.
            noise_type (str, optional): type of noise. Can be 'radermacher', 'sphere'
                or 'gaussian'. Defaults to 'radermacher'.
        """
        self.model = model
        self.learning_rate = learning_rate
        self.clipnorm = clipnorm
        self.n_slices = n_slices
        self.loss_type = loss_type.lower()
        self.noise_type = noise_type.lower()
        # setup optimizer
        self.optimizer = tf.keras.optimizers.Adam(
            learning_rate=learning_rate,
            clipnorm=clipnorm
        )
        
        self.num_gradsteps = 0
        self.num_epochs = 0
        self.progress = 0
        self.tb_writer = None

    def ssm_loss(self, x, v):
        """SSM loss from
        Sliced Score Matching: A Scalable Approach to Density and Score Estimation

        The loss is computed as
        s = -dE(x)/dx
        loss = vT*(ds/dx)*v + 1/2*(vT*s)^2

        Args:
            x (tf.Tensor): input samples
            v (tf.Tensor): sampled noises

        Returns:
            SSM loss
        """
        x = tf.repeat(x, self.n_slices, axis=0) # (n_slices*b, ...)
        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(x)
            score = self.model.score(x) # (n_slices*b, ...)
            sv    = tf.math.reduce_sum(score * v) # ()
        gsv = tape.gradient(sv, x) # (n_slices*b, ...)
        loss1 = tf.math.reduce_sum(score * v, axis=-1) ** 2 * 0.5 # (n_slices*b,)
        loss2 = tf.math.reduce_sum(v * gsv, axis=-1) # (n_slices*b,)
        loss = tf.math.reduce_mean(loss1 + loss2) # ()
        return loss
    
    def ssm_vr_loss(self, x, v):
        """SSM-VR (variance reduction) loss from
        Sliced Score Matching: A Scalable Approach to Density and Score Estimation

        The loss is computed as
        s = -dE(x)/dx
        loss = vT*(ds/dx)*v + 1/2*||s||^2

        Args:
            x (tf.Tensor): input samples
            v (tf.Tensor): sampled noises

        Returns:
            SSM-VR loss
        """
        x = tf.repeat(x, self.n_slices, axis=0) # (n_slices*b, ...)
        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(x)
            score = self.model.score(x) # (n_slices*b, ...)
            sv    = tf.math.reduce_sum(score * v) # ()
        gsv = tape.gradient(sv, x) # (n_slices*b, ...)
        loss1 = tf.norm(score, axis=-1) ** 2 * 0.5 # (n_slices*b,)
        loss2 = tf.math.reduce_sum(v*gsv, axis=-1) # (n_slices*b,)
        loss = tf.math.reduce_mean(loss1 + loss2) # ()
        return loss
    
    def deen_loss(self, x, v, sigma=0.1):
        """DEEN loss from
        Deep Energy Estimator Networks

        The loss is computed as
        x_ = x + v   # noisy samples
        s = -dE(x_)/dx_
        loss = 1/2*||x - x_ + sigma^2*s||^2

        Args:
            x (tf.Tensor): input samples
            v (tf.Tensor): sampled noises
            sigma (int, optional): noise scale. Defaults to 1.

        Returns:
            DEEN loss
        """
        v = v * sigma
        x_ = x + v
        s = (sigma ** 2) * self.model.score(x_)
        loss = tf.norm(s+v, axis=-1)**2
        loss = 0.5 * tf.math.reduce_mean(loss)
        return loss

    def dsm_loss(self, x, v, sigma=0.1):
        """DSM loss from
        A Connection Between Score Matching
            and Denoising Autoencoders

        The loss is computed as
        x_ = x + v   # noisy samples
        s = -dE(x_)/dx_
        loss = 1/2*||s + (x-x_)/sigma^2||^2

        Args:
            x (tf.Tensor): input samples
            v (tf.Tensor): sampled noises
            sigma (float, optional): noise scale. Defaults to 0.1.
        
        Returns:
            DSM loss
        """
        v = v * sigma
        x_ = x + v
        s = self.model.score(x_)
        loss = tf.norm(s + v/(sigma**2), axis=-1) ** 2
        loss = 0.5 * tf.math.reduce_mean(loss)
        return loss

    def get_random_noise(self, x, n_slices=None):
        """Sampling random noises

        Args:
            x (tf.Tensor): input samples
            n_slices (int, optional): number of slices. Defaults to None.

        Returns:
            tf.Tensor: sampled noises
        """
        if n_slices is None:
            v = tf.random.normal(x.shape)
        else:
            v = tf.random.normal((n_slices, *x.shape))
            v = tf.reshape(v, (-1, *v.shape[2:])) # (n_slices*b, 2)

        if self.noise_type == 'radermacher':
            v = tf.math.sign(v)
        elif self.noise_type == 'sphere':
            v = v/tf.norm(v, axis=-1, keepdims=True) * np.sqrt(v.shape[-1])
        elif self.noise_type == 'gaussian':
            pass
        else:
            raise NotImplementedError(
                f"Noise type '{self.noise_type}' not implemented."
            )
        return v
            
    
    def get_loss(self, x, v=None):
        """Compute loss

        Args:
            x (tf.Tensor): input samples
            v (tf.Tensor, optional): sampled noises. Defaults to None.

        Returns:
            loss
        """
        if self.loss_type == 'ssm-vr':
            v = self.get_random_noise(x, self.n_slices)
            loss = self.ssm_vr_loss(x, v)
        elif self.loss_type == 'ssm':
            v = self.get_random_noise(x, self.n_slices)
            loss = self.ssm_loss(x, v)
        elif self.loss_type == 'deen':
            v = self.get_random_noise(x, None)
            loss = self.deen_loss(x, v)
        elif self.loss_type == 'dsm':
            v = self.get_random_noise(x, None)
            loss = self.dsm_loss(x, v)
        else:
            raise NotImplementedError(
                f"Loss type '{self.loss_type}' not implemented."
            )

        return loss
    
    @tf.function
    def train_step(self, batch, update=True):
        """Train one batch

        Args:
            batch (dict): batch data
            update (bool, optional): whether to update networks. 
                Defaults to True.

        Returns:
            loss
        """
        x = batch['samples']
        # move inputs to device
        x = tf.convert_to_tensor(x, dtype=tf.float32)
        vars = self.model.variables
        with tf.GradientTape() as tape:
            tape.watch(vars)
            # compute losses
            loss = self.get_loss(x)
        # update model
        if update:
            # compute gradients
            grads = tape.gradient(loss, vars)
            self.optimizer.apply_gradients(zip(grads, vars))
        return loss

    def train(self, dataset, batch_size):
        """Train one epoch

        Args:
            dataset (tf.data.Dataset): Tensorflow dataset
            batch_size (int): batch size

        Returns:
            np.ndarray: mean loss
        """        
        all_losses = []
        dataset = dataset.batch(batch_size)
        for batch_data in dataset.as_numpy_iterator():
            sample_batch = {
                'samples': batch_data
            }
            loss = self.train_step(sample_batch)
            self.num_gradsteps += 1
            all_losses.append(loss)
        m_loss = np.mean(all_losses).astype(np.float32)
        return m_loss

    def eval(self, dataset, batch_size):
        """Eval one epoch

        Args:
            dataset (tf.data.Dataset): Tensorflow dataset
            batch_size (int): batch size

        Returns:
            np.ndarray: mean loss
        """        
        all_losses = []
        dataset = dataset.batch(batch_size)
        for batch_data in dataset.as_numpy_iterator():
            sample_batch = {
                'samples': batch_data
            }
            loss = self.train_step(sample_batch, update=False)
            all_losses.append(loss)
        m_loss = np.mean(all_losses).astype(np.float32)
        return m_loss

    def learn(
        self,
        train_dataset,
        eval_dataset = None,
        n_epochs = 5,
        batch_size = 100,
        log_freq = 1,
        eval_freq = 1,
        vis_freq = 1,
        vis_callback = None,
        tb_logdir = None
    ):
        """Train the model

        Args:
            train_dataset (tf.data.Dataset): training dataset
            eval_dataset (tf.data.Dataset, optional): evaluation dataset.
                Defaults to None.
            n_epochs (int, optional): number of epochs to train. Defaults to 5.
            batch_size (int, optional): batch size. Defaults to 100.
            log_freq (int, optional): logging frequency (epoch). Defaults to 1.
            eval_freq (int, optional): evaluation frequency (epoch). Defaults to 1.
            vis_freq (int, optional): visualizing frequency (epoch). Defaults to 1.
            vis_callback (callable, optional): visualization function. Defaults to None.
            tb_logdir (str, optional): path to tensorboard files. Defaults to None.

        Returns:
            self
        """        
        if tb_logdir is not None:
            self.tb_writer = tf.summary.create_file_writer(tb_logdir)

        # initialize
        time_start = time.time()
        time_spent = 0
        total_epochs = n_epochs

        for epoch in range(n_epochs):
            self.num_epochs += 1
            self.progress = float(self.num_epochs) / float(n_epochs)
            # train one epoch
            loss = self.train(train_dataset, batch_size)
            # write tensorboard
            if self.tb_writer is not None:
                with self.tb_writer.as_default():
                    tf.summary.scalar(f'train/loss', loss, step=self.num_epochs)
            
            if (log_freq is not None) and (self.num_epochs % log_freq == 0):
                logging.info(
                    f"[Epoch {self.num_epochs}/{total_epochs}]: loss: {loss}"
                )

            if (eval_dataset is not None) and (self.num_epochs % eval_freq == 0):
                # evaluate
                eval_loss = self.eval(eval_dataset, batch_size)

                if self.tb_writer is not None:
                    with self.tb_writer.as_default():
                        tf.summary.scalar(f'eval/loss', eval_loss, step=self.num_epochs)
                
                logging.info(
                    f"[Eval {self.num_epochs}/{total_epochs}]: loss: {eval_loss}"
                )

            if (vis_callback is not None) and (self.num_epochs % vis_freq == 0):
                logging.debug("Visualizing")
                vis_callback(self)
        return self

# --- built in --- # --- 3rd party --- # --- my module --- # --- primitives --- Toy MLP from         https://github.com/ermongroup/ncsn/blob/master/runners/toy_runner.py#L198         Args:             input_dim (int, optional): input dimensions. Defaults to 2.             output_dim (int, optional): output dimensions. Defaults to 1.             units (list, optional): hidden units. Defaults to [300, 300].             silu (bool, optional): use silu as activation function. Set False to use                 soft plus instead. Defaults to True.             dropout (bool, optional): use dropout layers. Defaults to False. # forwrad dummy tensor # --- energy model --- A simple energy model          Args:             net (tf.keras.Model): An energy function, the output shape of                 the energy function should be (b, 1). The score is                 computed by grad(-E(x)) Energy based model trainer          Args:             model (nn.Module): energy-based model             learning_rate (float, optional): learning rate. Defaults to 1e-4.             clipnorm (float, optional): gradient clip. Defaults to 100..             n_slices (int, optional): number of slices for sliced score matching loss.                 Defaults to 1.             loss_type (str, optional): type of loss. Can be 'ssm-vr', 'ssm', 'deen',                 'dsm'. Defaults to 'ssm-vr'.             noise_type (str, optional): type of noise. Can be 'radermacher', 'sphere'                 or 'gaussian'. Defaults to 'radermacher'. # setup optimizer SSM loss from         Sliced Score Matching: A Scalable Approach to Density and Score Estimation          The loss is computed as         s = -dE(x)/dx         loss = vT*(ds/dx)*v + 1/2*(vT*s)^2          Args:             x (tf.Tensor): input samples             v (tf.Tensor): sampled noises          Returns:             SSM loss # (n_slices*b, ...) # (n_slices*b, ...) # () # (n_slices*b, ...) # (n_slices*b,) # (n_slices*b,) # () SSM-VR (variance reduction) loss from         Sliced Score Matching: A Scalable Approach to Density and Score Estimation          The loss is computed as         s = -dE(x)/dx         loss = vT*(ds/dx)*v + 1/2*||s||^2          Args:             x (tf.Tensor): input samples             v (tf.Tensor): sampled noises          Returns:             SSM-VR loss # (n_slices*b, ...) # (n_slices*b, ...) # () # (n_slices*b, ...) # (n_slices*b,) # (n_slices*b,) # () DEEN loss from         Deep Energy Estimator Networks          The loss is computed as         x_ = x + v   # noisy samples         s = -dE(x_)/dx_         loss = 1/2*||x - x_ + sigma^2*s||^2          Args:             x (tf.Tensor): input samples             v (tf.Tensor): sampled noises             sigma (int, optional): noise scale. Defaults to 1.          Returns:             DEEN loss DSM loss from         A Connection Between Score Matching             and Denoising Autoencoders          The loss is computed as         x_ = x + v   # noisy samples         s = -dE(x_)/dx_         loss = 1/2*||s + (x-x_)/sigma^2||^2          Args:             x (tf.Tensor): input samples             v (tf.Tensor): sampled noises             sigma (float, optional): noise scale. Defaults to 0.1.                  Returns:             DSM loss Sampling random noises          Args:             x (tf.Tensor): input samples             n_slices (int, optional): number of slices. Defaults to None.          Returns:             tf.Tensor: sampled noises # (n_slices*b, 2) Compute loss          Args:             x (tf.Tensor): input samples             v (tf.Tensor, optional): sampled noises. Defaults to None.          Returns:             loss Train one batch          Args:             batch (dict): batch data             update (bool, optional): whether to update networks.                  Defaults to True.          Returns:             loss # move inputs to device # compute losses # update model # compute gradients Train one epoch          Args:             dataset (tf.data.Dataset): Tensorflow dataset             batch_size (int): batch size          Returns:             np.ndarray: mean loss Eval one epoch          Args:             dataset (tf.data.Dataset): Tensorflow dataset             batch_size (int): batch size          Returns:             np.ndarray: mean loss Train the model          Args:             train_dataset (tf.data.Dataset): training dataset             eval_dataset (tf.data.Dataset, optional): evaluation dataset.                 Defaults to None.             n_epochs (int, optional): number of epochs to train. Defaults to 5.             batch_size (int, optional): batch size. Defaults to 100.             log_freq (int, optional): logging frequency (epoch). Defaults to 1.             eval_freq (int, optional): evaluation frequency (epoch). Defaults to 1.             vis_freq (int, optional): visualizing frequency (epoch). Defaults to 1.             vis_callback (callable, optional): visualization function. Defaults to None.             tb_logdir (str, optional): path to tensorboard files. Defaults to None.          Returns:             self # initialize # train one epoch # write tensorboard # evaluate