# Path: nerfactor/networks/seq.py
# class Network(BaseNetwork):
#     """Assuming simple sequential flow.
#     """
#     def build(self, input_shape):
#         seq = tf.keras.Sequential(self.layers)
#         seq.build(input_shape)
#         for layer in self.layers:
#             assert layer.built, "Some layers not built"
# 
#     def __call__(self, tensor):
#         x = tensor
#         for layer in self.layers:
#             y = layer(x)
#             x = y
#         return y


import tensorflow as tf
from util import logging as logutil
from .seq import Network as BaseNetwork




logger = logutil.Logger(loggee="networks/mlp")


class Network(BaseNetwork):
    def __init__(self, widths, act=None, skip_at=None):
        super().__init__()
        depth = len(widths)
        if act is None:
            act = [None] * depth
        assert len(act) == depth, \
            "If not `None`, `act` must have the save length as `widths`"
        for w, a in zip(widths, act):
            if isinstance(a, str):
                a = tf.keras.layers.Activation(a)
            layer = tf.keras.layers.Dense(w, activation=a)
            self.layers.append(layer)
        self.skip_at = skip_at

    def __call__(self, x):
        # Shortcircuit if simply sequential
        if self.skip_at is None:
            return super().__call__(x)
        # Need to concatenate input at some levels
        x_ = x + 0 # make a copy
        for i, layer in enumerate(self.layers):
            y = layer(x_)

# Path: nerfactor/networks/seq.py
# class Network(BaseNetwork):
#     """Assuming simple sequential flow.
#     """
#     def build(self, input_shape):
#         seq = tf.keras.Sequential(self.layers)
#         seq.build(input_shape)
#         for layer in self.layers:
#             assert layer.built, "Some layers not built"
# 
#     def __call__(self, tensor):
#         x = tensor
#         for layer in self.layers:
#             y = layer(x)
#             x = y
#         return y


# Path: nerfactor/networks/mlp.py
import tensorflow as tf
from util import logging as logutil
from .seq import Network as BaseNetwork





logger = logutil.Logger(loggee="networks/mlp")


class Network(BaseNetwork):
    def __init__(self, widths, act=None, skip_at=None):
        super().__init__()
        depth = len(widths)
        if act is None:
            act = [None] * depth
        assert len(act) == depth, \
            "If not `None`, `act` must have the save length as `widths`"
        for w, a in zip(widths, act):
            if isinstance(a, str):
                a = tf.keras.layers.Activation(a)
            layer = tf.keras.layers.Dense(w, activation=a)
            self.layers.append(layer)
        self.skip_at = skip_at

    def __call__(self, x):
        # Shortcircuit if simply sequential
        if self.skip_at is None:
            return super().__call__(x)
        # Need to concatenate input at some levels
        x_ = x + 0 # make a copy
        for i, layer in enumerate(self.layers):
            y = layer(x_)

            if i in self.skip_at: