# This script runs the test functions defined in cpp/example.hpp.
#
# Author: Malte J. Ziebarth (malte.ziebarth@tum.de)
#
# Copyright (C) 2025 Technische Universität München

from funkybadger.example import test_function, test_function_2, test_function_3, \
    test_function_4, test_function_5
from funkybadger._funktionsklempner.quantityarray import QuantityArray

import numpy as np

def test0():
    """
    Tests test_function
    """
    name = "testname"
    i = 0
    result = test_function(name, i)
    assert result == 1.7281
    print("test0 result =", test_function(name, i))


def test1():
    """
    Tests test_function_2
    """
    result = test_function_2()
    assert result == "test test"
    print("test1 result =",test_function_2())


def test2():
    """
    Tests test_function_3
    """
    x = np.random.random(10)
    y = test_function_3(x)
    assert np.allclose(2.0 * x, y)


def test3():
    """
    Test test_function_4
    """
    array = np.random.random(38)
    qa = QuantityArray(array, 'm')
    out = test_function_4(qa)

    # Computes the potential energy of lifting
    # an object of 1 kg by 'qa':
    print("out.unit:", out.unit)
    assert out.unit == "J"

    for xin,xout in zip(qa.array, out.array):
        assert xout == 9.81 * xin


def test4():
    """
    Test test_function_5
    """
    rng = np.random.default_rng(283928)
    a0 = rng.random(5)
    a1 = rng.random(7)
    a2 = rng.random((2,3))
    qa0 = QuantityArray(a0, 'm')
    qa1 = QuantityArray(a1, 'kg')
    qa2 = QuantityArray(a2, 'kg')
    out = test_function_5(qa0, qa1)

    # Check that shape and unit matches our expectation.
    # Inside test_function_5, we compute 'm*g*h', potential
    # energy.
    assert out.shape == (a0.size, a1.size)
    assert out.unit == "J"

    # Check the quantitative results:
    assert np.allclose(out.array, 9.81 * a0[:,np.newaxis] * a1[np.newaxis,:])

    # Check the advanced shape computation:
    out2 = test_function_5(qa0, qa2)
    assert out2.shape == (a0.size, *a2.shape)
    assert out2.unit == "J"

    assert np.allclose(
        out2.array, 9.81 * a0[:,np.newaxis, np.newaxis] * a2[np.newaxis,:,:]
    )



test0()
test1()
test2()
test3()
test4()