"""Tests for finding a positive semi-definite correlation or covariance matrix

Created on Mon May 27 12:07:02 2013

Author: Josef Perktold
"""

from statsmodels.compat.platform import PLATFORM_OSX

import warnings

import numpy as np
from numpy.testing import assert_allclose, assert_almost_equal
import pytest
from scipy import sparse

from statsmodels.stats.correlation_tools import (
    FactoredPSDMatrix,
    _project_correlation_factors,
    _spg_optim,
    corr_clipped,
    corr_nearest,
    corr_nearest_factor,
    corr_thresholded,
    cov_nearest,
    cov_nearest_factor_homog,
)
from statsmodels.tools.testing import Holder


def norm_f(x, y):
    """Frobenious norm (squared sum) of difference between two arrays"""
    d = ((x - y) ** 2).sum()
    return np.sqrt(d)


# R library Matrix results
cov1_r = Holder()
# > nc  <- nearPD(pr, conv.tol = 1e-7, keepDiag = TRUE, doDykstra =FALSE, corr=TRUE)
# > cat_items(nc, prefix="cov1_r.")
cov1_r.mat = """<S4 object of class structure("dpoMatrix", package = "Matrix")>"""
cov1_r.eigenvalues = np.array(
    [
        4.197315628646795,
        0.7540460243978023,
        0.5077608149667492,
        0.3801267599652769,
        0.1607508970775889,
        4.197315628646795e-08,
    ]
)
cov1_r.corr = """TRUE"""
cov1_r.normF = 0.0743805226512533
cov1_r.iterations = 11
cov1_r.rel_tol = 8.288594638441735e-08
cov1_r.converged = """TRUE"""
# > mkarray2(as.matrix(nc$mat), name="cov1_r.mat")
cov1_r.mat = np.array(
    [
        1,
        0.487968018215892,
        0.642651880010906,
        0.4906386709070835,
        0.6440990530811909,
        0.8087111845493985,
        0.487968018215892,
        1,
        0.5141147294352735,
        0.2506688108312097,
        0.672351311297074,
        0.725832055882795,
        0.642651880010906,
        0.5141147294352735,
        1,
        0.596827778712154,
        0.5821917790519067,
        0.7449631633814129,
        0.4906386709070835,
        0.2506688108312097,
        0.596827778712154,
        1,
        0.729882058012399,
        0.772150225146826,
        0.6440990530811909,
        0.672351311297074,
        0.5821917790519067,
        0.729882058012399,
        1,
        0.813191720191944,
        0.8087111845493985,
        0.725832055882795,
        0.7449631633814129,
        0.772150225146826,
        0.813191720191944,
        1,
    ]
).reshape(6, 6, order="F")


cov_r = Holder()
# nc  <- nearPD(pr+0.01*diag(6), conv.tol = 1e-7, keepDiag = TRUE, doDykstra =FALSE, corr=FALSE)
# > cat_items(nc, prefix="cov_r.")
# cov_r.mat = '''<S4 object of class structure("dpoMatrix", package = "Matrix")>'''
cov_r.eigenvalues = np.array(
    [
        4.209897516692652,
        0.7668341923072066,
        0.518956980021938,
        0.390838551407132,
        0.1734728460460068,
        4.209897516692652e-08,
    ]
)
cov_r.corr = """FALSE"""
cov_r.normF = 0.0623948693159157
cov_r.iterations = 11
cov_r.rel_tol = 5.83987595937896e-08
cov_r.converged = """TRUE"""

# > mkarray2(as.matrix(nc$mat), name="cov_r.mat")
cov_r.mat = np.array(
    [
        1.01,
        0.486207476951913,
        0.6428524769306785,
        0.4886092840296514,
        0.645175579158233,
        0.811533860074678,
        0.486207476951913,
        1.01,
        0.514394615153752,
        0.2478398278204047,
        0.673852495852274,
        0.7297661648968664,
        0.6428524769306785,
        0.514394615153752,
        1.01,
        0.5971503271420517,
        0.582018469844712,
        0.7445177382760834,
        0.4886092840296514,
        0.2478398278204047,
        0.5971503271420517,
        1.01,
        0.73161232298669,
        0.7766852947049376,
        0.645175579158233,
        0.673852495852274,
        0.582018469844712,
        0.73161232298669,
        1.01,
        0.8107916469252828,
        0.811533860074678,
        0.7297661648968664,
        0.7445177382760834,
        0.7766852947049376,
        0.8107916469252828,
        1.01,
    ]
).reshape(6, 6, order="F")


def test_corr_psd():
    # test positive definite matrix is unchanged
    x = np.array([[1, -0.2, -0.9], [-0.2, 1, -0.2], [-0.9, -0.2, 1]])

    y = corr_nearest(x, n_fact=100)
    # print np.max(np.abs(x - y))
    assert_almost_equal(x, y, decimal=14)

    y = corr_clipped(x)
    assert_almost_equal(x, y, decimal=14)

    y = cov_nearest(x, n_fact=100)
    assert_almost_equal(x, y, decimal=14)

    x2 = x + 0.001 * np.eye(3)
    y = cov_nearest(x2, n_fact=100)
    assert_almost_equal(x2, y, decimal=14)


class CheckCorrPSDMixin:

    def test_nearest(self):
        x = self.x
        res_r = self.res
        y = corr_nearest(x, threshold=1e-7, n_fact=100)
        # print np.max(np.abs(x - y))
        assert_almost_equal(y, res_r.mat, decimal=3)
        d = norm_f(x, y)
        assert_allclose(d, res_r.normF, rtol=0.0015)
        evals = np.linalg.eigvalsh(y)
        # print 'evals', evals / res_r.eigenvalues[::-1] - 1
        assert_allclose(evals, res_r.eigenvalues[::-1], rtol=0.003, atol=1e-7)
        # print evals[0] / 1e-7 - 1
        assert_allclose(evals[0], 1e-7, rtol=1e-6)

    def test_clipped(self):
        x = self.x
        res_r = self.res
        y = corr_clipped(x, threshold=1e-7)
        # print np.max(np.abs(x - y)), np.max(np.abs((x - y) / y))
        assert_almost_equal(y, res_r.mat, decimal=1)
        d = norm_f(x, y)
        assert_allclose(d, res_r.normF, rtol=0.15)

        evals = np.linalg.eigvalsh(y)
        assert_allclose(evals, res_r.eigenvalues[::-1], rtol=0.1, atol=1e-7)
        assert_allclose(evals[0], 1e-7, rtol=0.02)

    def test_cov_nearest(self):
        x = self.x
        res_r = self.res
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            y = cov_nearest(x, method="nearest", threshold=1e-7)
        # print np.max(np.abs(x - y))
        assert_almost_equal(y, res_r.mat, decimal=2)
        d = norm_f(x, y)
        assert_allclose(d, res_r.normF, rtol=0.0015)


class TestCovPSD:

    @classmethod
    def setup_class(cls):
        x = np.array(
            [
                1,
                0.477,
                0.644,
                0.478,
                0.651,
                0.826,
                0.477,
                1,
                0.516,
                0.233,
                0.682,
                0.75,
                0.644,
                0.516,
                1,
                0.599,
                0.581,
                0.742,
                0.478,
                0.233,
                0.599,
                1,
                0.741,
                0.8,
                0.651,
                0.682,
                0.581,
                0.741,
                1,
                0.798,
                0.826,
                0.75,
                0.742,
                0.8,
                0.798,
                1,
            ]
        ).reshape(6, 6)
        cls.x = x + 0.01 * np.eye(6)
        cls.res = cov_r

    def test_cov_nearest(self):
        x = self.x
        res_r = self.res
        y = cov_nearest(x, method="nearest")
        # print np.max(np.abs(x - y))
        assert_almost_equal(y, res_r.mat, decimal=3)
        d = norm_f(x, y)
        assert_allclose(d, res_r.normF, rtol=0.001)

        y = cov_nearest(x, method="clipped")
        # print np.max(np.abs(x - y))
        assert_almost_equal(y, res_r.mat, decimal=2)
        d = norm_f(x, y)
        assert_allclose(d, res_r.normF, rtol=0.15)


class TestCorrPSD1(CheckCorrPSDMixin):

    @classmethod
    def setup_class(cls):
        x = np.array(
            [
                1,
                0.477,
                0.644,
                0.478,
                0.651,
                0.826,
                0.477,
                1,
                0.516,
                0.233,
                0.682,
                0.75,
                0.644,
                0.516,
                1,
                0.599,
                0.581,
                0.742,
                0.478,
                0.233,
                0.599,
                1,
                0.741,
                0.8,
                0.651,
                0.682,
                0.581,
                0.741,
                1,
                0.798,
                0.826,
                0.75,
                0.742,
                0.8,
                0.798,
                1,
            ]
        ).reshape(6, 6)
        cls.x = x
        cls.res = cov1_r


@pytest.mark.parametrize("threshold", [0, 1e-15, 1e-10, 1e-6])
def test_corrpsd_threshold(threshold):
    x = np.array([[1, -0.9, -0.9], [-0.9, 1, -0.9], [-0.9, -0.9, 1]])

    y = corr_nearest(x, n_fact=100, threshold=threshold)
    evals = np.linalg.eigvalsh(y)
    assert_allclose(evals[0], threshold, rtol=1e-6, atol=1e-15)

    y = corr_clipped(x, threshold=threshold)
    evals = np.linalg.eigvalsh(y)
    assert_allclose(evals[0], threshold, rtol=0.25, atol=1e-15)

    y = cov_nearest(x, method="nearest", n_fact=100, threshold=threshold)
    evals = np.linalg.eigvalsh(y)
    assert_allclose(evals[0], threshold, rtol=1e-6, atol=1e-15)

    y = cov_nearest(x, n_fact=100, threshold=threshold)
    evals = np.linalg.eigvalsh(y)
    assert_allclose(evals[0], threshold, rtol=0.25, atol=1e-15)


class Test_Factor:

    def test_corr_nearest_factor_arrpack(self):

        # regression results for svds call
        u2 = np.array(
            [
                [
                    6.39407581e-19,
                    9.15225947e-03,
                    1.82631698e-02,
                    2.72917181e-02,
                    3.61975557e-02,
                    4.49413101e-02,
                    5.34848732e-02,
                    6.17916613e-02,
                    6.98268388e-02,
                    7.75575058e-02,
                    8.49528448e-02,
                    9.19842264e-02,
                    9.86252769e-02,
                    1.04851906e-01,
                    1.10642305e-01,
                    1.15976906e-01,
                    1.20838331e-01,
                    1.25211306e-01,
                    1.29082570e-01,
                    1.32440778e-01,
                    1.35276397e-01,
                    1.37581605e-01,
                    1.39350201e-01,
                    1.40577526e-01,
                    1.41260396e-01,
                    1.41397057e-01,
                    1.40987160e-01,
                    1.40031756e-01,
                    1.38533306e-01,
                    1.36495727e-01,
                    1.33924439e-01,
                    1.30826443e-01,
                    1.27210404e-01,
                    1.23086750e-01,
                    1.18467769e-01,
                    1.13367717e-01,
                    1.07802909e-01,
                    1.01791811e-01,
                    9.53551023e-02,
                    8.85157320e-02,
                    8.12989329e-02,
                    7.37322125e-02,
                    6.58453049e-02,
                    5.76700847e-02,
                    4.92404406e-02,
                    4.05921079e-02,
                    3.17624629e-02,
                    2.27902803e-02,
                    1.37154584e-02,
                    4.57871801e-03,
                    -4.57871801e-03,
                    -1.37154584e-02,
                    -2.27902803e-02,
                    -3.17624629e-02,
                    -4.05921079e-02,
                    -4.92404406e-02,
                    -5.76700847e-02,
                    -6.58453049e-02,
                    -7.37322125e-02,
                    -8.12989329e-02,
                    -8.85157320e-02,
                    -9.53551023e-02,
                    -1.01791811e-01,
                    -1.07802909e-01,
                    -1.13367717e-01,
                    -1.18467769e-01,
                    -1.23086750e-01,
                    -1.27210404e-01,
                    -1.30826443e-01,
                    -1.33924439e-01,
                    -1.36495727e-01,
                    -1.38533306e-01,
                    -1.40031756e-01,
                    -1.40987160e-01,
                    -1.41397057e-01,
                    -1.41260396e-01,
                    -1.40577526e-01,
                    -1.39350201e-01,
                    -1.37581605e-01,
                    -1.35276397e-01,
                    -1.32440778e-01,
                    -1.29082570e-01,
                    -1.25211306e-01,
                    -1.20838331e-01,
                    -1.15976906e-01,
                    -1.10642305e-01,
                    -1.04851906e-01,
                    -9.86252769e-02,
                    -9.19842264e-02,
                    -8.49528448e-02,
                    -7.75575058e-02,
                    -6.98268388e-02,
                    -6.17916613e-02,
                    -5.34848732e-02,
                    -4.49413101e-02,
                    -3.61975557e-02,
                    -2.72917181e-02,
                    -1.82631698e-02,
                    -9.15225947e-03,
                    -3.51829569e-17,
                ]
            ]
        ).T
        s2 = np.array([24.88812183])

        d = 100
        dm = 1

        # Construct a test matrix with exact factor structure
        X = np.zeros((d, dm), dtype=np.float64)
        x = np.linspace(0, 2 * np.pi, d)
        for j in range(dm):
            X[:, j] = np.sin(x * (j + 1))
        _project_correlation_factors(X)
        X *= 0.7
        mat = np.dot(X, X.T)
        np.fill_diagonal(mat, 1.0)

        from scipy.sparse.linalg import svds

        u, s, vt = svds(mat, dm)

        # difference in sign
        dsign = np.sign(u[1]) * np.sign(u2[1])

        assert_allclose(u, dsign * u2, rtol=1e-6, atol=1e-14)
        assert_allclose(s, s2, rtol=1e-6)

    @pytest.mark.parametrize("dm", [1, 2])
    def test_corr_nearest_factor(self, dm):

        objvals = [
            np.array([6241.8, 6241.8, 579.4, 264.6, 264.3]),
            np.array([2104.9, 2104.9, 710.5, 266.3, 286.1]),
        ]

        d = 100

        # Construct a test matrix with exact factor structure
        X = np.zeros((d, dm), dtype=np.float64)
        x = np.linspace(0, 2 * np.pi, d)
        np.random.seed(10)
        for j in range(dm):
            X[:, j] = np.sin(x * (j + 1)) + 1e-10 * np.random.randn(d)

        _project_correlation_factors(X)
        assert np.isfinite(X).all()
        X *= 0.7
        mat = np.dot(X, X.T)
        np.fill_diagonal(mat, 1.0)

        # Try to recover the structure
        rslt = corr_nearest_factor(mat, dm, maxiter=10000)
        err_msg = "rank=%d, niter=%d" % (dm, len(rslt.objective_values))
        assert_allclose(
            rslt.objective_values[:5], objvals[dm - 1], rtol=0.5, err_msg=err_msg
        )
        assert rslt.Converged

        mat1 = rslt.corr.to_matrix()
        assert_allclose(mat, mat1, rtol=0.25, atol=1e-3, err_msg=err_msg)

    @pytest.mark.slow
    @pytest.mark.parametrize("dm", [1, 2])
    @pytest.mark.skipif(PLATFORM_OSX, reason="Frequent random failure on OSX")
    def test_corr_nearest_factor_sparse(self, dm):
        # Test that result is the same if the input is dense or sparse
        d = 200

        # Generate a test matrix of factors
        X = np.zeros((d, dm), dtype=np.float64)
        x = np.linspace(0, 2 * np.pi, d)
        rs = np.random.RandomState(10)
        for j in range(dm):
            X[:, j] = np.sin(x * (j + 1)) + rs.randn(d)

        # Get the correlation matrix
        _project_correlation_factors(X)
        X *= 0.7
        mat = np.dot(X, X.T)
        np.fill_diagonal(mat, 1)

        # Threshold it
        mat.flat[np.abs(mat.flat) < 0.35] = 0.0
        smat = sparse.csr_matrix(mat)

        dense_rslt = corr_nearest_factor(mat, dm, maxiter=10000)
        sparse_rslt = corr_nearest_factor(smat, dm, maxiter=10000)

        mat_dense = dense_rslt.corr.to_matrix()
        mat_sparse = sparse_rslt.corr.to_matrix()

        try:
            assert dense_rslt.Converged is sparse_rslt.Converged
            assert dense_rslt.Converged is True
        except AssertionError:
            pytest.xfail(
                reason="Convergence failed for one of dense/sparse. Frequent failure."
            )

        assert_allclose(mat_dense, mat_sparse, rtol=0.25, atol=1e-3)

    # Test on a quadratic function.
    @pytest.mark.skipif(PLATFORM_OSX, reason="Frequent random failure on OSX")
    def test_spg_optim(self, reset_randomstate):

        dm = 100

        ind = np.arange(dm)
        indmat = np.abs(ind[:, None] - ind[None, :])
        M = 0.8**indmat

        def obj(x):
            return np.dot(x, np.dot(M, x))

        def grad(x):
            return 2 * np.dot(M, x)

        def project(x):
            return x

        x = np.random.normal(size=dm)
        rslt = _spg_optim(obj, grad, x, project)
        xnew = rslt.params
        assert rslt.Converged is True
        assert_almost_equal(obj(xnew), 0, decimal=3)

    def test_decorrelate(self, reset_randomstate):

        d = 30
        dg = np.linspace(1, 2, d)
        root = np.random.normal(size=(d, 4))
        fac = FactoredPSDMatrix(dg, root)
        mat = fac.to_matrix()
        rmat = np.linalg.cholesky(mat)
        dcr = fac.decorrelate(rmat)
        idm = np.dot(dcr, dcr.T)
        assert_almost_equal(idm, np.eye(d))

        rhs = np.random.normal(size=(d, 5))
        mat2 = np.dot(rhs.T, np.linalg.solve(mat, rhs))
        mat3 = fac.decorrelate(rhs)
        mat3 = np.dot(mat3.T, mat3)
        assert_almost_equal(mat2, mat3)

    def test_logdet(self, reset_randomstate):

        d = 30
        dg = np.linspace(1, 2, d)
        root = np.random.normal(size=(d, 4))
        fac = FactoredPSDMatrix(dg, root)
        mat = fac.to_matrix()

        _, ld = np.linalg.slogdet(mat)
        ld2 = fac.logdet()

        assert_almost_equal(ld, ld2)

    def test_solve(self, reset_randomstate):

        d = 30
        dg = np.linspace(1, 2, d)
        root = np.random.normal(size=(d, 2))
        fac = FactoredPSDMatrix(dg, root)
        rhs = np.random.normal(size=(d, 5))
        sr1 = fac.solve(rhs)
        mat = fac.to_matrix()
        sr2 = np.linalg.solve(mat, rhs)
        assert_almost_equal(sr1, sr2)

    @pytest.mark.parametrize("dm", [1, 2])
    def test_cov_nearest_factor_homog(self, dm):

        d = 100

        # Construct a test matrix with exact factor structure
        X = np.zeros((d, dm), dtype=np.float64)
        x = np.linspace(0, 2 * np.pi, d)
        for j in range(dm):
            X[:, j] = np.sin(x * (j + 1))
        mat = np.dot(X, X.T)
        np.fill_diagonal(mat, np.diag(mat) + 3.1)

        # Try to recover the structure
        rslt = cov_nearest_factor_homog(mat, dm)
        mat1 = rslt.to_matrix()

        assert_allclose(mat, mat1, rtol=0.25, atol=1e-3)

    @pytest.mark.parametrize("dm", [1, 2])
    def test_cov_nearest_factor_homog_sparse(self, dm):
        # Check that dense and sparse inputs give the same result

        d = 100

        # Construct a test matrix with exact factor structure
        X = np.zeros((d, dm), dtype=np.float64)
        x = np.linspace(0, 2 * np.pi, d)
        for j in range(dm):
            X[:, j] = np.sin(x * (j + 1))
        mat = np.dot(X, X.T)
        np.fill_diagonal(mat, np.diag(mat) + 3.1)

        # Fit to dense
        rslt = cov_nearest_factor_homog(mat, dm)
        mat1 = rslt.to_matrix()

        # Fit to sparse
        smat = sparse.csr_matrix(mat)
        rslt = cov_nearest_factor_homog(smat, dm)
        mat2 = rslt.to_matrix()

        assert_allclose(mat1, mat2, rtol=0.25, atol=1e-3)

    def test_corr_thresholded(self, reset_randomstate):

        X = np.random.normal(size=(2000, 10))
        tcor = corr_thresholded(X, 0.2, max_elt=4e6)

        fcor = np.corrcoef(X)
        fcor *= np.abs(fcor) >= 0.2

        assert_allclose(tcor.todense(), fcor, rtol=0.25, atol=1e-3)
