import unittest

import itertools
import numpy

import chainer
from chainer.backends import cuda
from chainer import initializers
from chainer import links
from chainer import testing
from chainer.testing import attr


@testing.parameterize(*testing.product({
    'dtype': [numpy.float16, numpy.float32, numpy.float64],
    'initial_cost': ['random', None],
}))
class TestCRF1d(unittest.TestCase):
    def _calc_score(self, batch, ys):
        cost = self.link.cost.array
        return sum(x[batch, y] for x, y in zip(self.xs, ys)) + \
            sum(cost[y1, y2] for y1, y2 in zip(ys[:-1], ys[1:]))

    def _crf1d(self, cost_data, xs_data, ys_data):
        z = numpy.zeros((self.batches[0],), numpy.float32)
        for b, length in enumerate(self.lengths):
            for ys in itertools.product(range(self.n_label), repeat=length):
                z[b] += numpy.exp(chainer.cuda.to_cpu(self._calc_score(b, ys)))

        score = numpy.zeros((self.batches[0],), numpy.float32)
        for b, length in enumerate(self.lengths):
            ys = [self.ys[i][b] for i in range(length)]
            score[b] = self._calc_score(b, ys)

        loss = -(score - numpy.log(z))
        return numpy.sum(loss) / self.batches[0]

    def setUp(self):
        self._config_user = chainer.using_config('dtype', self.dtype)
        self._config_user.__enter__()
        self.n_label = 3

        self.lengths = [3, 3]
        self.batches = [2, 2, 2]

        self.xs = [numpy.random.uniform(-1, 1, (b, 3)).astype(self.dtype)
                   for b in self.batches]
        self.ys = [numpy.random.randint(
            0, self.n_label, (b,)).astype(numpy.int32)
            for b in self.batches]

        self.link = links.CRF1d(n_label=self.n_label)
        self.cost_shape = (self.n_label, self.n_label)

        if self.dtype == numpy.float16:
            self.check_forward_options = {'atol': 5e-3}
        else:
            self.check_forward_options = {'atol': 1e-4}

    def tearDown(self):
        self._config_user.__exit__(None, None, None)

    def check_forward(self, x_data, t_data):
        x = self.link(x_data, t_data)
        t = self._crf1d(self.link.cost.array, x_data, t_data)
        testing.assert_allclose(x.array, t,
                                **self.check_forward_options)

    def test_forward_cpu(self):
        self.check_forward(self.xs, self.ys)

    @attr.gpu
    def test_forward_gpu(self):
        self.link.to_gpu()
        self.check_forward(cuda.to_gpu(self.xs), cuda.to_gpu(self.ys))


@testing.parameterize(*testing.product({
    'dtype': [numpy.float16, numpy.float32, numpy.float64],
    'initializer': ['random', None]
}))
class TestInitialization(unittest.TestCase):

    def setUp(self):
        self.n_label = 3
        self.initial_cost = numpy.empty((self.n_label, self.n_label),
                                        dtype=self.dtype)

        if self.initializer is None:
            initializer = initializers.constant.Zero()

        elif self.initializer == 'random':
            initializer = initializers.GlorotUniform()

        initializer(self.initial_cost)
        with chainer.using_config('dtype', self.dtype):
            self.link = links.CRF1d(self.n_label,
                                    initial_cost=self.initial_cost)

    def check_param(self):
        link = self.link
        dtype = self.dtype
        assert link.cost.dtype == dtype
        testing.assert_allclose(link.cost.array,
                                self.initial_cost,
                                atol=0, rtol=0)

    def test_param_cpu(self):
        self.check_param()

    @attr.gpu
    def test_param_gpu(self):
        self.link.to_gpu()
        self.check_param()


testing.run_module(__name__, __file__)

@@ -39,16 +39,38 @@
 rt numpy
+%0Afrom six import moves
 %0A%0Aimport
@@ -373,16 +373,48 @@
  None%5D,%0A
+    'transpose': %5BTrue, False%5D,%0A
 %7D))%0Aclas
@@ -2153,12 +2153,91 @@
     
-x
+if self.transpose:%0A            # Make transposed arrays manually%0A            xs
  = 
+%5B
 self
@@ -2245,23 +2245,491 @@
 link
-(x_data, t_data
+.xp.empty((l, 3), dtype=self.dtype)%0A                  for l in self.lengths%5D%0A            ts = %5Bself.link.xp.empty((l,), dtype=numpy.int32)%0A                  for l in self.lengths%5D%0A            for i, batch in enumerate(self.batches):%0A                for j in moves.range(batch):%0A                    xs%5Bj%5D%5Bi%5D = x_data%5Bi%5D%5Bj%5D%0A                    ts%5Bj%5D%5Bi%5D = t_data%5Bi%5D%5Bj%5D%0A        else:%0A            xs = x_data%0A            ts = t_data%0A%0A        x = self.link(xs, ts, transpose=self.transpose
 )%0A  

Write test for transposed inputs