"""
Tests for Linear Regression
"""

import numpy as np
from pyai.ml.linear_regression import (
    linear_regression_normal_equation,
    linear_regression_sgd,
    plot_results,
    mean_squared_error,
)


def test_linear_regression():
    """Test Linear Regression"""

    # Toy dataset with only one feature
    samples = np.array(
        [
            [3.3],
            [4.4],
            [5.5],
            [6.71],
            [6.93],
            [4.168],
            [9.779],
            [6.182],
            [7.59],
            [2.167],
            [7.042],
            [10.791],
            [5.313],
            [7.997],
            [3.1],
        ]
    )

    # Expected results
    targets = np.array(
        [
            [1.7],
            [2.76],
            [2.09],
            [3.19],
            [1.694],
            [1.573],
            [3.366],
            [2.596],
            [2.53],
            [1.221],
            [2.827],
            [3.465],
            [1.65],
            [2.904],
            [1.3],
        ]
    )

    n_samples = samples.shape[0]

    # Add x0 = 1 to each sample for bias weight
    x_train = np.c_[np.ones((n_samples, 1)), samples]
    assert x_train.shape == (n_samples, samples.shape[1] + 1)

    y_pred_ne = linear_regression_normal_equation(x_train, targets)
    y_pred_sgd = linear_regression_sgd(x_train, targets)

    mse_ne = mean_squared_error(targets, y_pred_ne)
    mse_sgd = mean_squared_error(targets, y_pred_sgd)

    assert mse_ne <= 0.16892
    # Given its stochastic nature and the small number of samples,
    # results are sometimes poor with SGD
    assert mse_sgd < 2
    # When applicable, Normal Equation offers the best possible results
    assert mse_ne <= mse_sgd

    result_list = {
        "Normal Equation": y_pred_ne,
        "SGD": y_pred_sgd,
    }
    plot_results(samples, targets, result_list)
