import numpy as np
import scipy.stats as st
import pdb

import tensorflow as tf


def layer(op):
    def layer_decorated(self, *args, **kwargs):
        # 입력하지 않으면 이름을 'name'으로 자동 지정
        name = kwargs.setdefault('name', self.get_unique_name(op.__name__))
        # layer input 구체화
        if len(self.terminals) == 0:
            raise RuntimeError('No Input Variables Found For Layers %s.'%name)
        elif len(self.terminals) == 1:
            layer_input = self.terminals[0]
        else:
            layer_input = list(self.terminals)
        # 실행하고 결과값 가져오기
        layer_output = op(self, layer_input, *args, **kwargs)
        # Add to layer LUT
        self.layers[name] = layer_output
        # 아웃풋이 다음 레이어의 인풋이 됨
        self.feed(layer_output)
        # Return self for chained calls 
        return self

    return layer_decorated


class Smoother(object):
    def __init__(self, inputs, filter_size, sigma, heat_map_size=0):
        self.inputs = inputs
        self.terminals = []
        self.layers = dict(inputs)
        self.filter_size = filter_size
        self.sigma = sigma
        self.heat_map_size = heat_map_size
        self.setup()

    def setup(self):
        self.feed('data').conv(name = 'smoothing')

    def get_unique_name(self, prefix):
        ident = sum(t.startswith(prefix) for t, _ in self.layers.items()) + 1
        return '%s_%d' %(prefix, ident)

    def feed(self, *args):
        assert len(args) != 0
        self.terminals = []
        for fed_layer in args:
            if isinstance(fed_layer, str):
                try:
                    fed_layer = self.layers[fed_layer]
                except KeyError:
                    raise KeyError('Unknown layer name fed: %s'% fed_layer)
            self.terminals.append(fed_layer)
        return self

    def gauss_kernel(self, kernlen = 21, nsig = 3, channels = 1):
        interval = (2 * nsig + 1.)/(kernlen)
        x = np.linspace( -nsig -interval /2., nsig + interval /2., kernlen+1)
        kern1d = np.diff(st.norm.cdf(x))
        kernel_raw = np.sqrt(np.outer(kern1d, kern1d))
        kernel = kernel_raw/kernel_raw.sum()
        out_filter = np.array(kernel, dtype = np.float32)
        out_filter = out_filter.reshape((kernlen, kernlen, 1, 1))
        out_filter = np.repeat(out_filter, channels, axis = 2)
        return out_filter

    def make_gauss_var(self, name, size, sigma, c_i):
        # with tf.device('/cpu:0):
        kernel = self.gauss_kernel(size, sigma, c_i)
        var = tf.Variable(tf.convert_to_tensor(kernel), name=name)
        return var

    def get_output(self):
        '''Retrun the smoother output'''
        return self.terminals[-1]

    @layer
    def conv(self,
             input,
             name,
             padding = 'SAME'):
        # 인풋에서 채널 수 가져오기
        if self.heat_map_size != 0:
            c_i = self.heat_map_size
        else:
            c_i = input.get_shape().as_list()[3]
        # 입력 데이터에 Conv
        convolve = lambda i, k: tf.nn.depthwise_conv2d(i, k, [1,1,1,1], padding = padding)
        with tf.compat.v1.variable_scope(name) as scope:
            kernel = self.make_gauss_var('gauss_weight', self.filter_size, self.sigma, c_i)
            output = convolve(input, kernel)
        return output