import unittest

import numpy

import chainer
from chainer import cuda
from chainer import gradient_check
from chainer import links
from chainer import testing
from chainer.testing import attr


def _sigmoid(x):
    xp = cuda.get_array_module(x)
    half = x.dtype.type(0.5)
    return xp.tanh(x * half) * half + half


def _peephole(func, c, h, x):
    xp = cuda.get_array_module(x)
    with cuda.get_device_from_array(x):
        lstm_in = x.dot(func.upward.W.data.T)
        lstm_in += h.dot(func.lateral.W.data.T)
        lstm_in = xp.reshape(lstm_in, (len(lstm_in),
                                       lstm_in.shape[1] // 4,
                                       4))
        a, i, f, o = xp.split(lstm_in, 4, 2)
        a = xp.reshape(a, (len(a), a.shape[1]))
        i = xp.reshape(i, (len(i), i.shape[1]))
        f = xp.reshape(f, (len(f), f.shape[1]))
        o = xp.reshape(o, (len(o), o.shape[1]))
        peep_in_i = c.dot(func.peep_i.W.data.T)
        peep_in_f = c.dot(func.peep_f.W.data.T)
        a = xp.tanh(a)
        i = _sigmoid(i + peep_in_i)
        f = _sigmoid(f + peep_in_f)
        c_next = a * i + f * c
        peep_in_o = c_next.dot(func.peep_o.W.data.T)
        o = _sigmoid(o + peep_in_o)
        y = o * xp.tanh(c_next)
    return c_next, y


@testing.parameterize(
    {'in_size': 10, 'out_size': 10},
    {'in_size': 10, 'out_size': 40},
)
class TestPeephole(unittest.TestCase):

    def setUp(self):
        self.link = links.StatefulPeepholeLSTM(self.in_size, self.out_size)
        upward = self.link.upward.W.data
        upward[...] = numpy.random.uniform(-1, 1, upward.shape)
        lateral = self.link.lateral.W.data
        lateral[...] = numpy.random.uniform(-1, 1, lateral.shape)
        peep_i = self.link.peep_i.W.data
        peep_i[...] = numpy.random.uniform(-1, 1, peep_i.shape)
        peep_f = self.link.peep_f.W.data
        peep_f[...] = numpy.random.uniform(-1, 1, peep_f.shape)
        peep_o = self.link.peep_o.W.data