import unittest
import libpysal
from libpysal.examples import load_example
from libpysal.weights import Queen
import numpy as np
from scipy import sparse
from spreg.ml_error import ML_Error
from libpysal.common import RTOL, ATOL
from warnings import filterwarnings

filterwarnings("ignore", category=sparse.SparseEfficiencyWarning)
filterwarnings("ignore", message="^Method 'bounded' does not support")


class TestMLError(unittest.TestCase):
    def setUp(self):
        south = load_example("South")
        db = libpysal.io.open(south.get_path("south.dbf"), "r")
        self.y_name = "HR90"
        self.y = np.array(db.by_col(self.y_name))
        self.y.shape = (len(self.y), 1)
        self.x_names = ["RD90", "PS90", "UE90", "DV90"]
        self.x = np.array([db.by_col(var) for var in self.x_names]).T
        self.w = Queen.from_shapefile(south.get_path("south.shp"))
        self.w.transform = "r"

    def _estimate_and_compare(self, method="FULL", RTOL=RTOL):
        reg = ML_Error(
            self.y,
            self.x,
            w=self.w,
            name_y=self.y_name,
            name_x=self.x_names,
            name_w="south_q.gal",
            method=method,
        )
        betas = np.array([[6.1492], [4.4024], [1.7784], [-0.3781], [0.4858], [0.2991]])
        np.testing.assert_allclose(reg.betas, betas, RTOL + 0.0001)
        u = np.array([-5.97649777])
        np.testing.assert_allclose(reg.u[0], u, RTOL)
        predy = np.array([6.92258051])
        np.testing.assert_allclose(reg.predy[0], predy, RTOL)
        n = 1412
        np.testing.assert_allclose(reg.n, n, RTOL)
        k = 5
        np.testing.assert_allclose(reg.k, k, RTOL)
        y = np.array([0.94608274])
        np.testing.assert_allclose(reg.y[0], y, RTOL)
        x = np.array([1.0, -0.39902838, 0.89645344, 6.85780705, 7.2636377])
        np.testing.assert_allclose(reg.x[0], x, RTOL)
        e = np.array([-4.92843327])
        np.testing.assert_allclose(reg.e_filtered[0], e, RTOL)
        my = 9.5492931620846928
        np.testing.assert_allclose(reg.mean_y, my)
        sy = 7.0388508798387219
        np.testing.assert_allclose(reg.std_y, sy)
        vm = np.array(
            [1.06476526, 0.05548248, 0.04544514, 0.00614425, 0.01481356, 0.00143001]
        )
        np.testing.assert_allclose(reg.vm.diagonal(), vm, RTOL)
        sig2 = np.array([[32.40685441]])
        np.testing.assert_allclose(reg.sig2, sig2, RTOL)
        pr2 = 0.3057664820364818
        np.testing.assert_allclose(reg.pr2, pr2)
        std_err = np.array(
            [1.03187463, 0.23554719, 0.21317867, 0.07838525, 0.12171098, 0.03781546]
        )
        np.testing.assert_allclose(reg.std_err, std_err, RTOL)
        z_stat = [
            (5.9592751097983534, 2.5335926307459251e-09),
            (18.690182928021841, 5.9508619446611137e-78),
            (8.3421632936950338, 7.2943630281051907e-17),
            (-4.8232686291115678, 1.4122456582517099e-06),
            (3.9913060809142995, 6.5710406838016854e-05),
            (7.9088780724028922, 2.5971882547279339e-15),
        ]
        np.testing.assert_allclose(reg.z_stat, z_stat, RTOL, atol=ATOL)
        logll = -4471.407066887894
        np.testing.assert_allclose(reg.logll, logll, RTOL)
        aic = 8952.8141337757879
        np.testing.assert_allclose(reg.aic, aic, RTOL)
        schwarz = 8979.0779458660545
        np.testing.assert_allclose(reg.schwarz, schwarz, RTOL)

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

    def test_LU(self):
        self._estimate_and_compare(method="LU", RTOL=RTOL * 10)

    def test_ord(self):
        reg = ML_Error(
            self.y,
            self.x,
            w=self.w,
            name_y=self.y_name,
            name_x=self.x_names,
            name_w="south_q.gal",
            method="ORD",
        )
        betas = np.array([[6.1492], [4.4024], [1.7784], [-0.3781], [0.4858], [0.2991]])
        np.testing.assert_allclose(reg.betas, betas, RTOL + 0.0001)
        u = np.array([-5.97649777])
        np.testing.assert_allclose(reg.u[0], u, RTOL)
        predy = np.array([6.92258051])
        np.testing.assert_allclose(reg.predy[0], predy, RTOL)
        n = 1412
        np.testing.assert_allclose(reg.n, n, RTOL)
        k = 5
        np.testing.assert_allclose(reg.k, k, RTOL)
        y = np.array([0.94608274])
        np.testing.assert_allclose(reg.y[0], y, RTOL)
        x = np.array([1.0, -0.39902838, 0.89645344, 6.85780705, 7.2636377])
        np.testing.assert_allclose(reg.x[0], x, RTOL)
        e = np.array([-4.92843327])
        np.testing.assert_allclose(reg.e_filtered[0], e, RTOL)
        my = 9.5492931620846928
        np.testing.assert_allclose(reg.mean_y, my)
        sy = 7.0388508798387219
        np.testing.assert_allclose(reg.std_y, sy)
        vm = np.array(
            [1.06476526, 0.05548248, 0.04544514, 0.00614425, 0.01481356, 0.001501]
        )
        np.testing.assert_allclose(reg.vm.diagonal(), vm, RTOL * 10)
        sig2 = np.array([[32.40685441]])
        np.testing.assert_allclose(reg.sig2, sig2, RTOL)
        pr2 = 0.3057664820364818
        np.testing.assert_allclose(reg.pr2, pr2)
        std_err = np.array(
            [1.03187463, 0.23554719, 0.21317867, 0.07838525, 0.12171098, 0.038744]
        )
        np.testing.assert_allclose(reg.std_err, std_err, RTOL * 10)
        z_stat = [
            (5.95927510, 2.5335927e-09),
            (18.6901829, 5.9508630e-78),
            (8.34216329, 7.2943634e-17),
            (-4.8232686, 1.4122457e-06),
            (3.99130608, 6.5710407e-05),
            (7.71923784, 1.1702739e-14),
        ]
        np.testing.assert_allclose(reg.z_stat, z_stat, rtol=RTOL, atol=ATOL)
        logll = -4471.407066887894
        np.testing.assert_allclose(reg.logll, logll, RTOL)
        aic = 8952.8141337757879
        np.testing.assert_allclose(reg.aic, aic, RTOL)
        schwarz = 8979.0779458660545
        np.testing.assert_allclose(reg.schwarz, schwarz, RTOL)


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