import unittest
import libpysal
from scipy import sparse
import numpy as np
from spreg.ml_lag_regimes import ML_Lag_Regimes
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"
        self.regimes = db.by_col("CITCOU")

    def test_model1(self):
        reg = ML_Lag_Regimes(
            self.y,
            self.x,
            self.regimes,
            w=self.w,
            name_y=self.y_name,
            name_x=self.x_names,
            name_w=self.w_name,
            name_ds=self.ds_name,
            name_regimes="CITCOU",
            regime_lag_sep=False,
        )
        betas = np.array(
            [
                [-15.00586577],
                [4.49600801],
                [-0.03180518],
                [0.34995882],
                [-4.54040395],
                [3.92187578],
                [-0.17021393],
                [0.81941371],
                [0.53850323],
            ]
        )
        np.testing.assert_allclose(reg.betas, betas, RTOL)
        u = np.array([32.73718478])
        np.testing.assert_allclose(reg.u[0], u, RTOL)
        predy = np.array([14.26281522])
        np.testing.assert_allclose(reg.predy[0], predy, RTOL)
        n = 211
        np.testing.assert_allclose(reg.n, n, RTOL)
        k = 9
        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, 0.0, 0.0, 0.0, 0.0]])
        np.testing.assert_allclose(reg.x[0].toarray(), x, RTOL)
        e = np.array([29.45407124])
        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(
            [
                47.42000914,
                2.39526578,
                0.00506895,
                0.06480022,
                69.67653371,
                3.20661492,
                0.01156766,
                0.04862014,
                0.00400775,
            ]
        )
        np.testing.assert_allclose(reg.vm.diagonal(), vm, RTOL)
        sig2 = 200.04433357145007
        np.testing.assert_allclose(reg.sig2, sig2, RTOL)
        pr2 = 0.6404460298085746
        np.testing.assert_allclose(reg.pr2, pr2)
        std_err = np.array(
            [
                6.88621878,
                1.54766462,
                0.07119654,
                0.25455888,
                8.34724707,
                1.79070235,
                0.10755305,
                0.22049975,
                0.0633068,
            ]
        )
        np.testing.assert_allclose(reg.std_err, std_err, RTOL)
        logll = -864.98505596489736
        np.testing.assert_allclose(reg.logll, logll, RTOL)
        aic = 1747.9701119297947
        np.testing.assert_allclose(reg.aic, aic, RTOL)
        schwarz = 1778.1368351310794
        np.testing.assert_allclose(reg.schwarz, schwarz, RTOL)
        chow_r = np.array(
            [
                [1.00180776, 0.31687348],
                [0.05904944, 0.8080047],
                [1.16987812, 0.27942629],
                [1.95931177, 0.16158694],
            ]
        )
        np.testing.assert_allclose(reg.chow.regi, chow_r, RTOL)
        chow_j = 21.648337464039283
        np.testing.assert_allclose(reg.chow.joint[0], chow_j, RTOL)

    def test_model2(self):
        reg = ML_Lag_Regimes(
            self.y,
            self.x,
            self.regimes,
            w=self.w,
            name_y=self.y_name,
            name_x=self.x_names,
            name_w=self.w_name,
            name_ds=self.ds_name,
            name_regimes="CITCOU",
            regime_lag_sep=True,
        )
        betas = np.array(
            [
                [-0.71589799],
                [4.40910538],
                [-0.08652467],
                [0.46266265],
                [0.1627765],
                [-5.00594358],
                [2.91060349],
                [-0.18207394],
                [0.71129227],
                [0.66753263],
            ]
        )
        np.testing.assert_allclose(reg.betas, betas, RTOL)
        vm = np.array(
            [
                55.3593679,
                -7.22927797,
                -0.19487326,
                0.6030953,
                -0.52249569,
                0.0,
                0.0,
                0.0,
                0.0,
                0.0,
            ]
        )
        np.testing.assert_allclose(reg.vm[0], vm, RTOL)
        u = np.array([34.03630518])
        np.testing.assert_allclose(reg.u[0], u, RTOL)
        predy = np.array([12.96369482])
        np.testing.assert_allclose(reg.predy[0], predy, RTOL)
        e = np.array([32.46466912])
        np.testing.assert_allclose(reg.e_pred[0], e, RTOL)
        chow_r = np.array(
            [
                [0.15654726, 0.69235548],
                [0.43533847, 0.509381],
                [0.60552514, 0.43647766],
                [0.59214981, 0.441589],
                [11.69437282, 0.00062689],
            ]
        )
        np.testing.assert_allclose(reg.chow.regi, chow_r, RTOL)
        chow_j = 21.978012275873063
        np.testing.assert_allclose(reg.chow.joint[0], chow_j, RTOL)


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