from typing import Any

import numpy as np

import regression


NUM_DASHES = 40


def compare_expected_actual(expected: Any, actual: Any) -> None:
    print("-----")
    print(f"Expected: {expected}")
    print(f"  Actual: {actual}")
    if expected == actual:
        print("PASS")
    else:
        print("FAIL")


def compare_expected_actual_array(
        expected: np.ndarray, actual: np.ndarray
        ) -> None:
    print("-----")
    print(f"Expected: {expected}")
    print(f"  Actual: {actual}")
    if np.array_equal(expected, actual):
        print("PASS")
    else:
        print("FAIL")


def compare_expected_actual_array_close(
        expected: np.ndarray, actual: np.ndarray
        ) -> None:
    print("-----")
    print(f"Expected: {expected}")
    print(f"  Actual: {actual}")
    if np.allclose(expected, actual):
        print("PASS")
    else:
        print("FAIL")


def test_evaluator_accuracy() -> None:
    print("-" * NUM_DASHES)
    print("Testing Evaluator.accuracy()")

    evaluator = regression.Evaluator()

    y_true = np.array([1, 0, 1, 1, 0])
    y_pred = np.array([1, 1, 1, 0, 0])
    accuracy = evaluator.accuracy(y_true, y_pred)
    compare_expected_actual(0.6, accuracy)

    y_true = np.array([0, 1, 0, 1])
    y_pred = np.array([0, 1, 0, 1])
    accuracy = evaluator.accuracy(y_true, y_pred)
    compare_expected_actual(1.0, accuracy)


def test_evaluator_mean_squared_error() -> None:
    print("-" * NUM_DASHES)
    print("Testing Evaluator.mean_squared_error()")

    evaluator = regression.Evaluator()

    y_true = np.array([3, 5, 2])
    y_pred = np.array([1, 5, 4])
    mse = evaluator.mean_squared_error(y_true, y_pred)
    compare_expected_actual(8 / 3, mse)

    y_true = np.array([10, 8, 6, 4])
    y_pred = np.array([10, 8, 6, 4])
    mse = evaluator.mean_squared_error(y_true, y_pred)
    compare_expected_actual(0.0, mse)


def test_linear_regressor_fit() -> None:
    print("-" * NUM_DASHES)
    print("Testing LinearRegressor.fit()")

    regressor = regression.LinearRegressor()
    x = regression.add_intercept(np.array([[1], [2], [3]], dtype=float))
    y = np.array([3, 5, 7], dtype=float)
    regressor.fit(x, y)
    expected = np.array([1.0, 2.0])
    compare_expected_actual_array_close(expected, regressor.w)

    regressor = regression.LinearRegressor()
    x = regression.add_intercept(
        np.array([[1, 0], [0, 1], [2, 1]], dtype=float)
    )
    y = np.array([6, 5, 9], dtype=float)
    regressor.fit(x, y)
    expected = np.array([4.0, 2.0, 1.0])
    compare_expected_actual_array_close(expected, regressor.w)


def test_linear_regressor_predict() -> None:
    print("-" * NUM_DASHES)
    print("Testing LinearRegressor.predict()")

    regressor = regression.LinearRegressor()
    regressor.w = np.array([1.0, 2.0])
    x = regression.add_intercept(np.array([[4], [0], [7]], dtype=float))
    preds = regressor.predict(x)
    expected = np.array([9, 1, 15])
    compare_expected_actual_array(expected, preds)

    regressor = regression.LinearRegressor()
    regressor.w = np.array([2.5, 1.2])
    x = regression.add_intercept(np.array([[1], [2], [5]], dtype=float))
    preds = regressor.predict(x)
    expected = np.array([3, 4, 8])
    compare_expected_actual_array(expected, preds)


def run_all_tests() -> None:
    test_evaluator_accuracy()
    test_evaluator_mean_squared_error()
    test_linear_regressor_fit()
    test_linear_regressor_predict()


if __name__ == "__main__":
    run_all_tests()
