import numpy as np

from mla import Tensor, gradcheck


def _rand(*shape):
    rng = np.random.default_rng(0)
    return Tensor(rng.standard_normal(shape))


def test_add():
    ok, err = gradcheck(lambda a, b: a.add(b), [_rand(3, 4), _rand(3, 4)])
    assert ok, err


def test_add_broadcast():
    ok, err = gradcheck(lambda a, b: a.add(b), [_rand(3, 4), _rand(4)])
    assert ok, err


def test_mul():
    ok, err = gradcheck(lambda a, b: a.mul(b), [_rand(3, 4), _rand(3, 4)])
    assert ok, err


def test_mul_broadcast():
    ok, err = gradcheck(lambda a, b: a.mul(b), [_rand(2, 3, 4), _rand(4)])
    assert ok, err


def test_matmul():
    ok, err = gradcheck(lambda a, b: a.matmul(b), [_rand(3, 5), _rand(5, 2)])
    assert ok, err


def test_matmul_batched():
    ok, err = gradcheck(lambda a, b: a.matmul(b), [_rand(2, 3, 5), _rand(5, 2)])
    assert ok, err


def test_reshape():
    ok, err = gradcheck(lambda a: a.reshape(4, 3), [_rand(2, 6)])
    assert ok, err


def test_transpose():
    ok, err = gradcheck(lambda a: a.transpose((0, 2, 1)), [_rand(2, 3, 4)])
    assert ok, err


def test_sum_all():
    ok, err = gradcheck(lambda a: a.sum(), [_rand(3, 4)])
    assert ok, err


def test_sum_axis():
    ok, err = gradcheck(lambda a: a.sum(axis=1), [_rand(3, 4)])
    assert ok, err


def test_gather():
    idx = np.array([0, 2, 2, 4, 1])
    ok, err = gradcheck(lambda w: w.gather(idx), [_rand(5, 3)])
    assert ok, err


def test_chain_add_mul_sum():
    ok, err = gradcheck(lambda a, b: a.mul(b).add(a).sum(), [_rand(3, 4), _rand(3, 4)])
    assert ok, err