import unittest

import numpy as np

from inmoose.limma import squeezeVar


class Test(unittest.TestCase):
    def test_squeezeVar(self):
        with self.assertRaisesRegex(ValueError, expected_regex="var is empty"):
            squeezeVar(np.array([]), df=3)

        res = squeezeVar(np.array([1]), df=3)
        self.assertEqual(res["df_prior"], 0)
        self.assertTrue(np.array_equal(res["var_prior"], [1]))
        self.assertTrue(np.array_equal(res["var_post"], [1]))

        res = squeezeVar(np.array([0.5, 1.3, 7.8]), df=3)
        self.assertAlmostEqual(res["df_prior"], 2.830777, places=6)
        self.assertAlmostEqual(res["var_prior"], 1.676927, places=6)
        self.assertTrue(
            np.allclose(res["var_post"], [1.071385, 1.482994, 4.827317], atol=1e-6)
        )

        res = squeezeVar(
            np.array([0.5, 0.5, 1.3, 1.3, 7.8, 7.8]),
            df=20,
            covariate=[1, 2, 3, 4, 5, 6],
        )
        self.assertAlmostEqual(res["df_prior"], 15.44626, places=5)
        self.assertTrue(
            np.allclose(
                res["var_prior"],
                [0.4322342, 0.6426360, 1.0331707, 1.9359150, 4.2277336, 9.9836662],
                atol=1e-6,
            )
        )
        self.assertTrue(
            np.allclose(
                res["var_post"],
                [0.4704700, 0.5621558, 1.1837250, 1.5771098, 6.2433295, 8.7515664],
                atol=1e-6,
            )
        )