from ..utils import loader
import numpy as np


class Layer(object):

    def __init__(self, *args):
        self._signature = list(args)
        self.type = list(args)[0]
        self.number = list(args)[1]

        self.w = dict()  # weights
        self.h = dict()  # placeholders
        self.wshape = dict()  # weight shape
        self.wsize = dict()  # weight size
        self.setup(*args[2:])  # set attr up
        self.present()
        for var in self.wshape:
            shp = self.wshape[var]
            size = np.prod(shp)
            self.wsize[var] = size

    def load(self, src_loader):
        var_lay = src_loader.VAR_LAYER
        if self.type not in var_lay: return

        src_type = type(src_loader)
        if src_type is loader.weights_loader:
            wdict = self.load_weights(src_loader)
        else:
            wdict = self.load_ckpt(src_loader)
        if wdict is not None:
            self.recollect(wdict)

    def load_weights(self, src_loader):
        val = src_loader([self.presenter])
        if val is None:
            return None
        else:
            return val.w

    def load_ckpt(self, src_loader):
        result = dict()
        presenter = self.presenter
        for var in presenter.wshape:
            name = presenter.varsig(var)
            shape = presenter.wshape[var]
            key = [name, shape]
            val = src_loader(key)
            result[var] = val
        return result

    @property
    def signature(self):
        return self._signature

    # For comparing two layers
    def __eq__(self, other):
        return self.signature == other.signature

    def __ne__(self, other):
        return not self.__eq__(other)

    def varsig(self, var):
        if var not in self.wshape:
            return None
        sig = str(self.number)
        sig += '-' + self.type
        sig += '/' + var
        return sig

    def recollect(self, w):
        self.w = w

    def present(self):
        self.presenter = self

    def setup(self, *args):
        pass

    def finalize(self):
        pass