from .base import Model
from .exceptions import ShapeError


class Affine(Model):
    name = 'affine'
    nr_out = None
    nr_in = None

    @property
    def nr_weight(self):
        return (self.nr_out * self.nr_in) + self.nr_out

    def set_weights(self, data=None, initialize=True):
        if data is None:
            self.data = self.ops.allocate_pool(self.nr_weight,
                            name=(self.name, 'pool'))
        else:
            self.data = data
        self.W = self.data.allocate_shape((self.nr_out, self.nr_in))
        self.b = self.data.allocate_shape((self.nr_out,))
        if initialize:
            self.ops.xavier_uniform_init(self.W, inplace=True)

    def set_gradient(self, data=None, initialize=False):
        if data is None:
            self.d_data = self.ops.allocate_pool(self.nr_weight,
                            name=(self.name, 'pool'))
        else:
            self.d_data = data
        self.d_W = self.d_data.allocate_shape((self.nr_out, self.nr_in))
        self.d_b = self.d_data.allocate_shape((self.nr_out,))

    def predict_batch(self, input_BI):
        return self.ops.affine(self.W, self.b, input_BI)

    def begin_update(self, input_BI, dropout=0.0):
        self.check_shape(input_BI, True)
        output_BO = self.ops.affine(self.W, self.b, input_BI)
        output_BO, bp_dropout = self.ops.dropout(output_BO, dropout)
        return output_BO, bp_dropout(self._get_finish_update(input_BI))
   
    def _get_finish_update(self, acts_BI):
        def finish_update(d_acts_BO, optimizer=None, **kwargs):
            self.d_b += d_acts_BO.sum(axis=0)
            self.d_W = self.ops.batch_outer(d_acts_BO, acts_BI)
            
            if optimizer is not None:
                optimizer(self.W, self.d_W, key=('W', self.name), **kwargs)
                optimizer(self.b, self.d_b, key=('b', self.name), **kwargs)
            return self.ops.batch_dot(d_acts_BO, self.W.T)
        return finish_update


class ReLu(Affine):
    name = 'relu'
    def predict_batch(self, input_BI):
        dotted = self.ops.xp.tensordot(input_BI, self.W, axes=[[1], [1]])
        dotted += self.b
        return self.ops.clip_low(dotted, 0, inplace=True)

    def begin_update(self, input_BI, dropout=0.0):
        output_BO = self.ops.affine(self.W, self.b, input_BI)
        mask = self.ops.get_dropout(output_BO.shape, dropout)
        if mask is not None:
            mask *= output_BO > 0
            output_BO *= mask
        return output_BO, self._get_finish_update(input_BI, mask)


class Softmax(Affine):
    name = 'softmax'
    def predict_batch(self, input_bi):
        output_bo = self.ops.affine(self.W, self.b, input_bi)
        return self.ops.softmax(output_bo, axis=-1, inplace=True)


class Maxout(Affine):
    name = 'maxout'
    def predict_batch(self, input_bi):
        take_which = self.ops.take_which
        argmax = self.ops.argmax
        affine = self.ops.affine
        return take_which(argmax(affine(input_bi, self.W, self.b)))

    def begin_update(self, input_BI, dropout=0.0):
        W_OCI = self.W
        b_OC = self.b
        output_BOC = self.ops.affine(W_OCI, b_OC, input_BI)
        which_BO = self.ops.argmax(output_BOC, axis=-1)
        best_BO = self.ops.take_which(output_BOC, which_BO)
        mask_BO = self.ops.get_dropout(best_BO.shape, dropout)
        finish_update = self._get_finish_update(input_BI, which_BO, mask_BO)
        if mask_BO is not None:
            best_BO *= mask_BO
        return best_BO, finish_update

    def _get_finish_update(self, acts_BI, which_BO, mask_BO):
        def finish_update(d_acts_BO):
            raise NotImplementedError
            #d_acts_BO *= mask_BO
            ## TODO
            #self.d_b += d_acts_BO.sum(axis=0)
            #self.d_W += d_W_OCI
            #return d_acts_BI
        return backward

@@ -2284,57 +2284,61 @@
 elf.
-ops.affine(self.W, self.b, input_BI)%0A        mask
+predict_batch(input_BI)%0A        output_BO, bp_dropout
  = s
@@ -2337,36 +2337,32 @@
 pout = self.ops.
-get_
 dropout(output_B
@@ -2358,30 +2358,24 @@
 ut(output_BO
-.shape
 , dropout)%0A 
@@ -2377,101 +2377,8 @@
 ut)%0A
-        if mask is not None:%0A            mask *= output_BO %3E 0%0A            output_BO *= mask%0A
     
@@ -2431,22 +2431,16 @@
 input_BI
-, mask
 )%0A%0A%0Aclas

Refactor ReLu begin_update