import unittest
import libpysal
from scipy import sparse
import numpy as np
from spreg.ml_lag import ML_Lag
from libpysal.common import RTOL
from warnings import filterwarnings

filterwarnings("ignore", category=sparse.SparseEfficiencyWarning)


class TestMLError(unittest.TestCase):
    def setUp(self):
        db = libpysal.io.open(libpysal.examples.get_path("baltim.dbf"), "r")
        self.ds_name = "baltim.dbf"
        self.y_name = "PRICE"
        self.y = np.array(db.by_col(self.y_name)).T
        self.y.shape = (len(self.y), 1)
        self.x_names = ["NROOM", "AGE", "SQFT"]
        self.x = np.array([db.by_col(var) for var in self.x_names]).T
        ww = libpysal.io.open(libpysal.examples.get_path("baltim_q.gal"))
        self.w = ww.read()
        ww.close()
        self.w_name = "baltim_q.gal"
        self.w.transform = "r"

    def _estimate_and_compare(self, **kwargs):
        reg = ML_Lag(
            self.y,
            self.x,
            w=self.w,
            name_y=self.y_name,
            name_x=self.x_names,
            name_w=self.w_name,
            name_ds=self.ds_name,
            **kwargs
        )
        betas = np.array(
            [[-6.04040164], [3.48995114], [-0.20103955], [0.65462382], [0.62351143]]
        )
        np.testing.assert_allclose(reg.betas, betas, RTOL)
        u = np.array([47.51218398])
        np.testing.assert_allclose(reg.u[0], u, RTOL)
        predy = np.array([-0.51218398])
        np.testing.assert_allclose(reg.predy[0], predy, RTOL)
        n = 211
        np.testing.assert_allclose(reg.n, n, RTOL)
        k = 5
        np.testing.assert_allclose(reg.k, k, RTOL)
        y = np.array([47.0])
        np.testing.assert_allclose(reg.y[0], y, RTOL)
        x = np.array([1.0, 4.0, 148.0, 11.25])
        np.testing.assert_allclose(reg.x[0], x, RTOL)
        e = np.array([41.99251608])
        np.testing.assert_allclose(reg.e_pred[0], e, RTOL)
        my = 44.307180094786695
        np.testing.assert_allclose(reg.mean_y, my)
        sy = 23.606076835380495
        np.testing.assert_allclose(reg.std_y, sy)
        vm = np.array([28.57288755, 1.42341656, 0.00288068, 0.02956392, 0.00332139])
        np.testing.assert_allclose(reg.vm.diagonal(), vm, RTOL)
        sig2 = 216.27525647243797
        np.testing.assert_allclose(reg.sig2, sig2, RTOL)
        pr2 = 0.6133020721559487
        np.testing.assert_allclose(reg.pr2, pr2)
        std_err = np.array([5.34536131, 1.19307022, 0.05367198, 0.17194162, 0.05763147])
        np.testing.assert_allclose(reg.std_err, std_err, RTOL)
        logll = -875.92771143484833
        np.testing.assert_allclose(reg.logll, logll, RTOL)
        aic = 1761.8554228696967
        np.testing.assert_allclose(reg.aic, aic, RTOL)
        schwarz = 1778.614713537077
        np.testing.assert_allclose(reg.schwarz, schwarz, RTOL)

    def test_dense(self):
        self._estimate_and_compare(method="FULL")

    def test_ord(self):
        reg = ML_Lag(
            self.y,
            self.x,
            w=self.w,
            name_y=self.y_name,
            name_x=self.x_names,
            name_w=self.w_name,
            name_ds=self.ds_name,
            method="ORD",
        )
        betas = np.array(
            [[-6.04040164], [3.48995114], [-0.20103955], [0.65462382], [0.62351143]]
        )
        np.testing.assert_allclose(reg.betas, betas, RTOL)
        u = np.array([47.51218398])
        np.testing.assert_allclose(reg.u[0], u, RTOL)
        predy = np.array([-0.51218398])
        np.testing.assert_allclose(reg.predy[0], predy, RTOL)
        n = 211
        np.testing.assert_allclose(reg.n, n, RTOL)
        k = 5
        np.testing.assert_allclose(reg.k, k, RTOL)
        y = np.array([47.0])
        np.testing.assert_allclose(reg.y[0], y, RTOL)
        x = np.array([1.0, 4.0, 148.0, 11.25])
        np.testing.assert_allclose(reg.x[0], x, RTOL)
        e = np.array([41.99251608])
        np.testing.assert_allclose(reg.e_pred[0], e, RTOL)
        my = 44.307180094786695
        np.testing.assert_allclose(reg.mean_y, my)
        sy = 23.606076835380495
        np.testing.assert_allclose(reg.std_y, sy)
        vm = np.array([28.63404, 1.423698, 0.002884738, 0.02957845, 0.003379166])
        np.testing.assert_allclose(reg.vm.diagonal(), vm, RTOL)
        sig2 = 216.27525647243797
        np.testing.assert_allclose(reg.sig2, sig2, RTOL)
        pr2 = 0.6133020721559487
        np.testing.assert_allclose(reg.pr2, pr2)
        std_err = np.array([5.351078, 1.193188, 0.05371, 0.171984, 0.058131])
        np.testing.assert_allclose(reg.std_err, std_err, RTOL)
        logll = -875.92771143484833
        np.testing.assert_allclose(reg.logll, logll, RTOL)
        aic = 1761.8554228696967
        np.testing.assert_allclose(reg.aic, aic, RTOL)
        schwarz = 1778.614713537077
        np.testing.assert_allclose(reg.schwarz, schwarz, RTOL)

    def test_LU(self):
        self._estimate_and_compare(method="LU")


if __name__ == "__main__":
    unittest.main()
