#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# nnutil2 - Tensorflow utilities for training neural networks
# Copyright (c) 2019, Abdó Roig-Maranges <abdo.roig@gmail.com>
#
# This file is part of 'nnutil2'.
#
# This file may be modified and distributed under the terms of the 3-clause BSD
# license. See the LICENSE file for details.

from typing import List

import tensorflow as tf

from ..util import kwargs_for
from .layer import Layer

class Segment(Layer):
    """A sequential collection of layers"""
    def __init__(self, layers: List[Layer] = [], activation=None, **kwargs):
        super(Segment, self).__init__(**kwargs)

        self._segment_layers = layers
        self._segment_activation = tf.keras.activations.get(activation)
        self._segment_states = []

    def get_config(self):
        config = {
            'layers': [ly.get_config() for ly in self._layers],
            'activation': self._segment_activation
        }

        base_config = super(Segment, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def call(self, inputs, **kwargs):
        x = inputs
        self._segment_states.append(x)

        for l in self._segment_layers:
            layer_kwargs = kwargs_for(kwargs, l.call)
            x = l(x, **layer_kwargs)
            self._segment_states.append(x)

        if self._segment_activation is not None:
            x = self._segment_activation(x)
            self._segment_states.append(x)

        return x

    def compute_output_shape(self, input_shape):
        shape = input_shape
        for l in self._segment_layers:
            shape = l.compute_output_shape(shape)
        return shape

    @property
    def flat_layers(self):
        layers = []

        def add_layers(ly):
            if isinstance(ly, Segment):
                for ly2 in ly.layers:
                    add_layers(ly2)
            else:
                layers.append(ly)

        add_layers(self)
        return layers

    @property
    def layers(self):
        return self._segment_layers

    @property
    def states(self):
        return self._segment_states
