import numpy as np

from util import auxiliary


class TestAuxiliary(object):

    def test_regression(self):
        np.random.seed(923)
        datasize = 1000
        X = np.random.random((datasize, 3))
        missX = np.zeros((datasize,3), dtype=np.bool_)
        Y = X@np.array([[.3], [.4],[.5]]) + (np.random.random((datasize,1)) - .5)
        # panel_data = np.savetxt('regress_test.csv', np.column_stack((Y,X)),  delimiter=',')
        missY = np.zeros((datasize), dtype=np.bool_)
        result = auxiliary.regress(X, Y[:, 0], missX=missX, missY=missY)
        resids = Y - X@result[:,None]
        std_errors = np.sqrt(np.diag(np.linalg.inv((X.T)@X) * (np.var(resids))))
        assert np.all(np.abs(result[:,None] - np.array([[.3], [.4],[.5]])) < (1.2 * std_errors))