import unittest

from inmoose.edgepy import pnbinom, qnbinom

class test_nbinom(unittest.TestCase):

    def test_nbinom(self):

        for i in range(10):

            hi = pnbinom(i, size=5, prob=0.2, lower_tail=False)

            lo = pnbinom(i, size=5, prob=0.2, lower_tail=True)

            self.assertEqual(hi, 1.0 - lo)

            self.assertEqual(i, qnbinom(hi, size=5, prob=0.2, lower_tail=False))

            self.assertEqual(i, qnbinom(lo, size=5, prob=0.2, lower_tail=True))

            hi = pnbinom(i, size=5, mu=4, lower_tail=False)

            lo = pnbinom(i, size=5, mu=4, lower_tail=True)

            self.assertAlmostEqual(hi, 1.0 - lo)

            self.assertEqual(i, qnbinom(hi, size=5, mu=4, lower_tail=False))

            self.assertEqual(i, qnbinom(lo, size=5, mu=4, lower_tail=True))

        ctxt = self.assertRaisesRegex(

            ValueError, expected_regex="exactly one of prob and mu must be provided"

        )

        with ctxt:

            pnbinom(42, size=5, prob=0.5, mu=4)

        with ctxt:

            qnbinom(42, size=5, prob=0.5, mu=4)

        with ctxt:

            pnbinom(42, size=5)

        with ctxt:

            qnbinom(42, size=5)