# Path: nlt/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: nlt/networks/elements.py
# def conv(kernel_size, n_ch_out, stride=1):
#     return tf.keras.layers.Conv2D(
#         n_ch_out,
#         kernel_size,
#         strides=stride,
#         padding='same')
# 
# def norm(type_):
#     if type_ == 'batch':
#         norm_layer = tf.keras.layers.BatchNormalization(
#             momentum=0.99, epsilon=0.001)
#     elif type_ == 'layer':
#         norm_layer = tf.keras.layers.LayerNormalization(
#             epsilon=0.001, center=True, scale=True)
#     elif type_ == 'instance':
#         norm_layer = instancenorm()
#     elif type_ == 'pixel':
#         norm_layer = pixelnorm()
#     elif type_ is None or type_.lower() == 'none':
#         norm_layer = iden()
#     else:
#         raise NotImplementedError(type_)
#     return norm_layer
# 
# def act(type_):
#     if type_ == 'relu':
#         act_layer = tf.keras.layers.ReLU(negative_slope=0)
#     elif type_ == 'leakyrelu':
#         act_layer = tf.keras.layers.LeakyReLU(alpha=0.3)
#     elif type_ == 'elu':
#         act_layer = tf.keras.layers.ELU(alpha=1.0)
#     else:
#         raise NotImplementedError(type_)
#     return act_layer
# 
# def pool(type_):
#     kwargs = {
#         'pool_size': 2,
#         'strides': 2,
#         'padding': 'same'}
#     if type_ == 'max':
#         pool_layer = tf.keras.layers.MaxPooling2D(**kwargs)
#     elif type_ == 'avg':
#         pool_layer = tf.keras.layers.AveragePooling2D(**kwargs)
#     elif type_ is None or type_.lower() == 'none':
#         pool_layer = iden()
#     else:
#         raise NotImplementedError(type_)
#     return pool_layer
# 
# def iden():
#     return tf.keras.layers.Lambda(lambda x: x)
# 
# def deconv(kernel_size, n_ch_out, stride=1):
#     return tf.keras.layers.Conv2DTranspose(
#         n_ch_out,
#         kernel_size,
#         strides=stride,
#         padding='same')
# 
# def upconv(n_ch_out):
#     """2x upsampling a feature map.
#     """
#     layers = [
#         tf.keras.layers.UpSampling2D(size=2, interpolation='bilinear'),
#         tf.keras.layers.Conv2D(n_ch_out, 2, padding='same')]
#     return tf.keras.Sequential(layers)


import numpy as np
import tensorflow as tf
from util import logging as logutil, net as netutil
from .seq import Network as BaseNetwork
from .elements import conv, norm, act, pool, iden, deconv, upconv
        # 1x1 conv to generate an original-res. feature map
        self.layers.append(conv(1, n_feat[0], stride=1))
        self.is_contracting.append(True)
        self.spatsize_changes.append(1)
        for n in n_feat[:-1]:
            # Contracting spatially
            if n >= prev_n: # so 64 -> 64 is considered "contracting"
                self.layers.append(
                    tf.keras.Sequential([
                        conv(kernel, n, stride=stride),
                        norm(norm_type), # could be identity
                        act(act_type),
                        conv(kernel, n, stride=1),
                        norm(norm_type),
                        act(act_type),
                        pool(pool_type), # could be identity
                    ]))
                self.is_contracting.append(True)
                spatsize_change = 1 / stride
                if pool_type is not None:
                    spatsize_change *= 1 / 2
                self.spatsize_changes.append(spatsize_change)
            # Expanding spatially
            else:
                self.layers.append(
                    tf.keras.Sequential([
                        iden() if pool_type is None else upconv(n),
                        deconv(kernel, n, stride=stride),
                        norm(norm_type),
                        act(act_type),

# Path: nlt/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: nlt/networks/elements.py
# def conv(kernel_size, n_ch_out, stride=1):
#     return tf.keras.layers.Conv2D(
#         n_ch_out,
#         kernel_size,
#         strides=stride,
#         padding='same')
# 
# def norm(type_):
#     if type_ == 'batch':
#         norm_layer = tf.keras.layers.BatchNormalization(
#             momentum=0.99, epsilon=0.001)
#     elif type_ == 'layer':
#         norm_layer = tf.keras.layers.LayerNormalization(
#             epsilon=0.001, center=True, scale=True)
#     elif type_ == 'instance':
#         norm_layer = instancenorm()
#     elif type_ == 'pixel':
#         norm_layer = pixelnorm()
#     elif type_ is None or type_.lower() == 'none':
#         norm_layer = iden()
#     else:
#         raise NotImplementedError(type_)
#     return norm_layer
# 
# def act(type_):
#     if type_ == 'relu':
#         act_layer = tf.keras.layers.ReLU(negative_slope=0)
#     elif type_ == 'leakyrelu':
#         act_layer = tf.keras.layers.LeakyReLU(alpha=0.3)
#     elif type_ == 'elu':
#         act_layer = tf.keras.layers.ELU(alpha=1.0)
#     else:
#         raise NotImplementedError(type_)
#     return act_layer
# 
# def pool(type_):
#     kwargs = {
#         'pool_size': 2,
#         'strides': 2,
#         'padding': 'same'}
#     if type_ == 'max':
#         pool_layer = tf.keras.layers.MaxPooling2D(**kwargs)
#     elif type_ == 'avg':
#         pool_layer = tf.keras.layers.AveragePooling2D(**kwargs)
#     elif type_ is None or type_.lower() == 'none':
#         pool_layer = iden()
#     else:
#         raise NotImplementedError(type_)
#     return pool_layer
# 
# def iden():
#     return tf.keras.layers.Lambda(lambda x: x)
# 
# def deconv(kernel_size, n_ch_out, stride=1):
#     return tf.keras.layers.Conv2DTranspose(
#         n_ch_out,
#         kernel_size,
#         strides=stride,
#         padding='same')
# 
# def upconv(n_ch_out):
#     """2x upsampling a feature map.
#     """
#     layers = [
#         tf.keras.layers.UpSampling2D(size=2, interpolation='bilinear'),
#         tf.keras.layers.Conv2D(n_ch_out, 2, padding='same')]
#     return tf.keras.Sequential(layers)


# Path: nlt/networks/convnet.py
import numpy as np
import tensorflow as tf
from util import logging as logutil, net as netutil
from .seq import Network as BaseNetwork
from .elements import conv, norm, act, pool, iden, deconv, upconv

        # 1x1 conv to generate an original-res. feature map
        self.layers.append(conv(1, n_feat[0], stride=1))
        self.is_contracting.append(True)
        self.spatsize_changes.append(1)
        for n in n_feat[:-1]:
            # Contracting spatially
            if n >= prev_n: # so 64 -> 64 is considered "contracting"
                self.layers.append(
                    tf.keras.Sequential([
                        conv(kernel, n, stride=stride),
                        norm(norm_type), # could be identity
                        act(act_type),
                        conv(kernel, n, stride=1),
                        norm(norm_type),
                        act(act_type),
                        pool(pool_type), # could be identity
                    ]))
                self.is_contracting.append(True)
                spatsize_change = 1 / stride
                if pool_type is not None:
                    spatsize_change *= 1 / 2
                self.spatsize_changes.append(spatsize_change)
            # Expanding spatially
            else:
                self.layers.append(
                    tf.keras.Sequential([
                        iden() if pool_type is None else upconv(n),
                        deconv(kernel, n, stride=stride),
                        norm(norm_type),
                        act(act_type),

                        deconv(kernel, n, stride=1),