"""
Tests for metrics
"""

import math
import numpy as np
from pyai.metrics import euclidean_distance, mean_squared_error


def test_mse():
    """Test Mean Squared Error."""

    y_true = np.array([3, -0.5, 2, 7])
    y_pred = np.array([2.5, 0.0, 2, 8])

    error = mean_squared_error(y_true, y_pred)
    assert math.isclose(error, (0.25 + 0.25 + 1) / 4)

    y_true = np.array([[1, 0, 0, 1], [0, 1, 1, 1], [1, 1, 0, 1]])
    y_pred = np.array([[0, 0, 0, 1], [1, 0, 1, 1], [0, 0, 0, 1]])

    error = mean_squared_error(y_true, y_pred)
    assert math.isclose(error, (1.0 / 3 + 2.0 / 3 + 2.0 / 3) / 4)


def test_euclidean_distance():
    """Test euclidean distance."""

    a = np.array([1, 1.5])
    b = np.array([3, 3])

    dist = euclidean_distance(a, b)
    assert math.isclose(dist, np.linalg.norm(a - b))

    a = np.random.rand(3, 3)
    b = np.random.rand(3, 3)

    dist = euclidean_distance(a, b)
    assert math.isclose(dist, np.linalg.norm(a - b))
