# ***************************************************************
# Copyright (c) 2021 Jittor. All Rights Reserved. 
# Maintainers: Dun Liang <randonlang@gmail.com>. 
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.
# ***************************************************************
import unittest
import jittor as jt
import numpy as np
from .test_core import expect_error
from .test_grad import ngrad
from itertools import permutations
from jittor.test.test_cuda import test_cuda

def gen_data(shape):
    num = np.multiply.reduce(shape)
    a = np.arange(0, num)
    return a.reshape(shape)

class TestTransposeOp(unittest.TestCase):
    def test_with_np(self):
        def check(a):
            perms = list(permutations(range(a.ndim))) + [None]
            for perm in perms:
                if perm:
                    x = np.transpose(a, perm)
                    y = jt.transpose(a, perm).data
                else:
                    x = np.transpose(a)
                    y = jt.transpose(a).data
                self.assertEqual(x.shape, y.shape)
                assert (x==y).all(), f"\n{x}\n{y}"
                
        # ia = [gen_data([2,3,4,5]), gen_data([5,3])]
        ia = [gen_data([2,2,2]), gen_data([2,3,4,5]), gen_data([5,3])]
        for a in ia: check(a)
        
    def test_grad(self):
        def check(a):
            perms = list(permutations(range(a.ndim))) + [None]
            for perm in perms:
                x = jt.array(a).float()
                if perm:
                    y = x.transpose(perm)
                else:
                    y = x.transpose()
                dx = jt.grad(y*y, x).data
                self.assertEqual(dx.shape, a.shape)
                assert (dx==a*2).all(), f"\n{dx}\n{a}\n{perm}"
        ia = [gen_data([2,2,2]), gen_data([2,3,4,5]), gen_data([5,3])]
        for a in ia: check(a)
        
    def test_matmul_grad(self):
        np.random.seed(0)
        for i in range(10):
            a = np.random.rand(2,3).astype("float32")
            b = np.random.rand(3,4).astype("float32")
            out, (da, db) = ngrad(lambda vars: np.matmul(vars[0],vars[1]).sum(), [a,b], 1e-1)
            ja = jt.array(a)
            jb = jt.array(b)
            jc = ja.matmul(jb)
            jda, jdb = jt.grad(jc, [ja,jb])
            assert ((da-jda.data)<1e-5).all(), (da, jda.data, da-jda.data)
            assert ((db-jdb.data)<1e-5).all(), (db-jdb.data)

    def test_permute(self):
        a = jt.ones([2,3,4])
        assert a.permute().shape == [4,3,2]
        assert a.permute(0,2,1).shape == [2,4,3]

    def test_transpose_3d2i(self):
        a = jt.ones([2,3,4])
        assert a.transpose(0,1).shape == (3,2,4)

    @unittest.skipIf(not jt.compiler.has_cuda, "No CUDA found")
    @jt.flag_scope(use_cuda=1)
    def test_cutt(self):
        a = jt.rand((10,2)) > 0.5
        b = a.transpose()
        assert (a.data.transpose() == b.data).all()

        a = jt.zeros((1,1))
        b = a.transpose((1,0))
        b.sync()


class TestFuseTransposeOp(unittest.TestCase):

    def test_fuse_transpose1(self):
        with jt.profile_scope() as rep:
            a = jt.rand((10,11,12))
            b = a.fuse_transpose((1,2,0))+1
            np.testing.assert_allclose(
                a.data.transpose((1,2,0))+1,
                b.data
            )
        assert len(rep) == 3

    def test_fuse_transpose2(self):
        with jt.profile_scope() as rep:
            a = jt.rand((10,11,12))
            b = (a+1).fuse_transpose((1,2,0))
            np.testing.assert_allclose(
                a.data.transpose((1,2,0))+1,
                b.data
            )
        assert len(rep) == 3

    def test_fuse_transpose3(self):
        with jt.profile_scope() as rep:
            a = jt.rand((10,11,12))
            c = jt.rand((11,12,10))
            b = a.fuse_transpose((1,2,0))+c
            np.testing.assert_allclose(
                a.data.transpose((1,2,0))+c.data,
                b.data
            )
        assert len(rep) == 3

    def test_fuse_transpose4(self):
        with jt.profile_scope() as rep:
            a = jt.rand((10,11,12))
            c = jt.rand((10,11,12))
            b = (a+c).fuse_transpose((1,2,0))
            np.testing.assert_allclose(
                (a.data+c.data).transpose((1,2,0)),
                b.data
            )
        assert len(rep) == 3

    def test_fuse_transpose5(self):
        with jt.profile_scope() as rep:
            a = jt.rand((10,11,6,7))
            c = jt.rand((10,11,6,7))
            b = (a+c).fuse_transpose((1,0,2,3))
            np.testing.assert_allclose(
                (a.data+c.data).transpose((1,0,2,3)),
                b.data
            )
        assert len(rep) == 3


@unittest.skipIf(not jt.compiler.has_cuda, "No CUDA found")
class TestFuseTransposeCudaOp(TestFuseTransposeOp):
    def setUp(self):
        jt.flags.use_cuda = 1
    def tearDown(self):
        jt.flags.use_cuda = 0

if __name__ == "__main__":
    unittest.main()