import tf_slim as slim

from .baseop import BaseOp

import tensorflow as tf

from distutils.version import StrictVersion

class route(BaseOp):

    def forward(self):

        routes = self.lay.routes

        routes_out = list()

        for r in routes:

            this = self.inp

            while this.lay.number != r:

                this = this.inp

                assert this is not None, \

                    'Routing to non-existence {}'.format(r)

            routes_out += [this.out]

        self.out = tf.concat(routes_out, 3)

    def speak(self):

        msg = 'concat {}'

        return msg.format(self.lay.routes)

class connected(BaseOp):

    def forward(self):

        self.out = tf.nn.xw_plus_b(

            self.inp.out,

            self.lay.w['weights'],

            self.lay.w['biases'],

            name=self.scope)

    def speak(self):

        layer = self.lay

        args = [layer.inp, layer.out]

        args += [layer.activation]

        msg = 'full {} x {}  {}'

        return msg.format(*args)

class select(connected):

    """a weird connected layer"""

    def speak(self):

        layer = self.lay

        args = [layer.inp, layer.out]

        args += [layer.activation]

        msg = 'sele {} x {}  {}'

        return msg.format(*args)

class extract(connected):

    """a weird connected layer"""

    def speak(self):

        layer = self.lay

        args = [len(layer.inp), len(layer.out)]

        args += [layer.activation]

        msg = 'extr {} x {}  {}'

        return msg.format(*args)

class flatten(BaseOp):

    def forward(self):

        temp = tf.transpose(

            self.inp.out, [0, 3, 1, 2])

        self.out = slim.flatten(

            temp, scope=self.scope)

    def speak(self): return 'flat'

class softmax(BaseOp):

    def forward(self):

        self.out = tf.nn.softmax(self.inp.out)

    def speak(self): return 'softmax()'

class avgpool(BaseOp):

    def forward(self):

        self.out = tf.reduce_mean(

            self.inp.out, [1, 2],

            name=self.scope

        )

    def speak(self): return 'avgpool()'

class dropout(BaseOp):

    def forward(self):

        if self.lay.h['pdrop'] is None:

            self.lay.h['pdrop'] = 1.0

        self.out = tf.nn.dropout(

            self.inp.out,

            self.lay.h['pdrop'],

            name=self.scope

        )

    def speak(self): return 'drop'

class crop(BaseOp):

    def forward(self):

        self.out = self.inp.out * 2. - 1.

    def speak(self):

        return 'scale to (-1, 1)'

class maxpool(BaseOp):

    def forward(self):

        self.out = tf.nn.max_pool(

            self.inp.out, padding='SAME',

            ksize=[1] + [self.lay.ksize] * 2 + [1],

            strides=[1] + [self.lay.stride] * 2 + [1],

            name=self.scope

        )

    def speak(self):

        l = self.lay

        return 'maxp {}x{}p{}_{}'.format(

            l.ksize, l.ksize, l.pad, l.stride)

class leaky(BaseOp):

    def forward(self):

        self.out = tf.maximum(

            .1 * self.inp.out,

            self.inp.out,

            name=self.scope

        )

    def verbalise(self): pass

class identity(BaseOp):

    def __init__(self, inp):

        self.inp = None

        self.out = inp