import pygpu

from pygpu.basic import (tril, triu)
from unittest import TestCase
from .support import (guard_devsup, gen_gpuarray, context)
import numpy


def test_tril():
    for dtype in ['float32','float64']:
        for shape in [(10, 5), (5, 10), (10, 10)]:
            for order in ['c', 'f']:
                for inplace in [True, False]:
                    yield run_tril, dtype, shape, order, inplace

@guard_devsup
def run_tril(dtype, shape, order, inplace):
    ac, ag = gen_gpuarray(shape, dtype, order=order, ctx=context)
    result = tril(ag, inplace=inplace)
    assert numpy.all(numpy.tril(ac) == result)
    if inplace:
        assert numpy.all(numpy.tril(ac) == ag)
    else:
        assert numpy.all(ac == ag)


def test_triu():
    for dtype in ['float32','float64']:
        for shape in [(10, 5), (5, 10), (10, 10)]:
            for order in ['c', 'f']:
                for inplace in [True, False]:
                    yield run_triu, dtype, shape, order, inplace

@guard_devsup
def run_triu(dtype, shape, order, inplace):
    ac, ag = gen_gpuarray(shape, dtype, order=order, ctx=context)
    result = triu(ag, inplace=inplace)
    assert numpy.all(numpy.triu(ac) == result)
    if inplace:
        assert numpy.all(numpy.triu(ac) == ag)
    else:
        assert numpy.all(ac == ag)


class test_errors(TestCase):

    def runTest(self):
        self.assertRaises(ValueError, self.run_1d_triu)
        self.assertRaises(ValueError, self.run_3d_triu)
        self.assertRaises(ValueError, self.run_1d_tril)
        self.assertRaises(ValueError, self.run_3d_tril)

        self.assertRaises(ValueError, self.run_noncontiguous_tril)
        self.assertRaises(ValueError, self.run_noncontiguous_triu)

    def run_1d_triu(self):
        ac, ag = gen_gpuarray((10, ), 'float32', ctx=context)
        triu(ag)

    def run_3d_triu(self):
        ac, ag = gen_gpuarray((10, 10, 10), 'float32', ctx=context)
        triu(ag)

    def run_1d_tril(self):
        ac, ag = gen_gpuarray((10, ), 'float32', ctx=context)
        tril(ag)

    def run_3d_tril(self):
        ac, ag = gen_gpuarray((10, 10, 10), 'float32', ctx=context)
        tril(ag)

    def run_noncontiguous_tril(self):
        a = numpy.random.rand(5, 5)
        b = pygpu.array(a, context=context)
        b = b[::-1]
        assert b.flags.c_contiguous is b.flags.f_contiguous is False
        tril(b)

    def run_noncontiguous_triu(self):
        a = numpy.random.rand(5, 5)
        b = pygpu.array(a, context=context)
        b = b[::-1]
        assert b.flags.c_contiguous is b.flags.f_contiguous is False
        triu(b)
