
from unittest import TestCase

import numpy as np
from numpy.testing import assert_allclose
from pytest import mark

from fastats.core.ast_transforms.convert_to_jit import convert_to_jit
from fastats.linear_algebra import qr, qr_classical_gram_schmidt
from tests.data.datasets import SKLearnDataSets

qr_jit = convert_to_jit(qr)
qr_classical_gram_schmidt_jit = convert_to_jit(qr_classical_gram_schmidt)


class QRTestMixin:

    @staticmethod
    def assert_orthonormal(Q):
        n = Q.shape[1]
        assert_allclose(Q.T @ Q, np.eye(n), atol=1e-10)

    @staticmethod
    def check_versus_expectations(Q, Q_expected, R, R_expected, A):
        assert_allclose(Q, Q_expected)
        assert_allclose(R, R_expected)
        assert_allclose(Q @ R, A)

    def test_ucla_4x3(self):
        """
        QR decomposition of a 4x3 matrix, taken from literature directory
        'ucla_qr_factorization.pdf':
        """
        A = np.array([[-1, -1, 1],
                      [1, 3, 3],
                      [-1, -1, 5],
                      [1, 3, 7]])

        Q_expected = np.array([[-0.5, 0.5, -0.5],
                               [0.5, 0.5, -0.5],
                               [-0.5, 0.5, 0.5],
                               [0.5, 0.5, 0.5]])

        R_expected = np.array([[2, 4, 2],
                               [0, 2, 8],
                               [0, 0, 4]])

        Q, R = self.fn(A)
        self.check_versus_expectations(Q, Q_expected, R, R_expected, A)
        self.assert_orthonormal(Q)

    def test_wikipedia_3x3(self):
        """
        QR decomposition of a 3x3 matrix, per the following:
        https://en.wikipedia.org/wiki/QR_decomposition
        """
        A = np.array([[12, -51, 4],
                      [6, 167, -68],
                      [-4, 24, -41]])

        Q_expected = np.array([[6/7, -69/175, -58/175],
                               [3/7, 158/175, 6/175],
                               [-2/7, 6/35, -33/35]])

        R_expected = np.array([[14, 21, -14],
                               [0, 175, -70],
                               [0, 0, 35]])

        Q, R = self.fn(A)
        self.check_versus_expectations(Q, Q_expected, R, R_expected, A)
        self.assert_orthonormal(Q)


class QRTests(QRTestMixin, TestCase):

    def setUp(self):
        self.fn = qr


class QRClassicalGSTests(QRTestMixin, TestCase):

    def setUp(self):
        self.fn = qr_classical_gram_schmidt


class QRJitTests(QRTestMixin, TestCase):

    def setUp(self):
        self.fn = qr_jit


class QRClassicalGSJitTests(QRTestMixin, TestCase):

    def setUp(self):
        self.fn = qr_classical_gram_schmidt_jit


def standardise(Q, R):
    """
    QR decomposition may not be unique; here we chose to enforce
    positive R diagonals to facilitate comparison of solutions
    """
    D = np.diag(np.sign(np.diag(R)))
    return Q @ D, D @ R


def check_versus_numpy(A, fn, atol=1e-7):
    Q_expected, R_expected = standardise(*np.linalg.qr(A))
    Q, R = standardise(*fn(A))
    assert_allclose(Q, Q_expected, atol=atol)
    assert_allclose(R, R_expected, atol=atol)


@mark.parametrize('A', SKLearnDataSets)
def test_sklearn_dataset(A):
    data = A.value
    A = data.T @ data
    check_versus_numpy(A, qr_jit)


def test_stability_of_classical_gram_schmidt():
    dataset = SKLearnDataSets.WINE
    data = dataset.value
    A = data.T @ data

    check_versus_numpy(A, qr_classical_gram_schmidt, atol=1e-1)
    check_versus_numpy(A, qr, atol=1e-10)  # <- much tighter tol

    # notice that maximum atol for the modified Gram Schmidt
    # method is significantly higher for this data set due to
    # better numerical stability


if __name__ == '__main__':
    import pytest
    pytest.main([__file__])

@@ -246,16 +246,61 @@
 schmidt%0A
+from fastats.scaling.scaling import standard%0A
 from tes
@@ -2953,19 +2953,8 @@
 , fn
-, atol=1e-7
 ):%0A 
@@ -3075,27 +3075,16 @@
 expected
-, atol=atol
 )%0A    as
@@ -3114,19 +3114,8 @@
 cted
-, atol=atol
 )%0A%0A%0A
@@ -3133,17 +3133,23 @@
 etrize('
-A
+dataset
 ', SKLea
@@ -3189,183 +3189,38 @@
 set(
-A):%0A    data = A.value%0A    A = data.T @ data%0A    check_versus_numpy(A, qr_jit)%0A%0A%0Adef test_stability_of_classical_gram_schmidt():%0A    dataset = SKLearnDataSets.WINE%0A    data = 
+dataset):%0A    data = standard(
 data
@@ -3228,16 +3228,17 @@
 et.value
+)
 %0A    A =
@@ -3252,17 +3252,16 @@
  @ data%0A
-%0A
     chec
@@ -3285,264 +3285,12 @@
  qr_
-classical_gram_schmidt, atol=1e-1)%0A    check_versus_numpy(A, qr, atol=1e-10)  # %3C- much tighter tol%0A%0A    # notice that maximum atol for the modified Gram Schmidt%0A    # method is significantly higher for this data set due to%0A    # better numerical stability
+jit)
 %0A%0A%0Ai

update tests