import unittest

from online_fdr.investing.lord.dependent import LordDependent
from online_fdr.investing.lord.mem_decay import LORDMemoryDecay
from online_fdr.investing.lord.plus_plus import LordPlusPlus
from online_fdr.investing.lord.discard import LordDiscard
from online_fdr.investing.lord.three import LordThree
from online_fdr.utils.testing import get_test_data


class TestSuiteLord(unittest.TestCase):

    DATA: dict = get_test_data()

    def test_lord_three(self):
        """
        Deviates from the current C++ implementation in 'onlineFDR'.
        Resembles the original R implementation in 'onlineFDR'.
        See tests/reference/onlineFDR_lond.R for further information.
        """
        lord = LordThree(alpha=0.05, wealth=0.025, reward=0.025)

        alpha, decision = [], []
        for i, p_value in enumerate(self.DATA["p_value"]):
            result = lord.test_one(p_value)
            alpha.append(round(lord.alpha, ndigits=6))
            decision.append(result)

        self.assertEqual(
            alpha,
            [
                0.001338,
                0.002604,
                0.003803,
                0.000827,
                0.000704,
                0.000586,
                0.004824,
                0.001049,
                0.000893,
                0.000743,
                0.005760,
                0.001253,
                0.001067,
                0.006665,
                0.001450,
            ],
        )

        self.assertEqual(
            decision,
            [
                True,
                True,
                False,
                False,
                False,
                True,
                False,
                False,
                False,
                True,
                False,
                False,
                True,
                False,
                False,
            ],
        )

    def test_lord_plus_plus(self):
        lord = LordPlusPlus(alpha=0.05, wealth=0.025)

        alpha, decision = [], []
        for i, p_value in enumerate(self.DATA["p_value"]):
            result = lord.test_one(p_value)
            alpha.append(round(lord.alpha, ndigits=6))
            decision.append(result)

        self.assertEqual(
            alpha,
            [
                0.001338,
                0.001629,
                0.003215,
                0.001036,
                0.000876,
                0.000738,
                0.003309,
                0.001136,
                0.000987,
                0.000854,
                0.003426,
                0.001251,
                0.0011,
                0.003639,
                0.001438,
            ],
        )

        self.assertEqual(
            decision,
            [
                True,
                True,
                False,
                False,
                False,
                True,
                False,
                False,
                False,
                True,
                False,
                False,
                True,
                False,
                False,
            ],
        )

    def test_lord_discard(self):
        lord = LordDiscard(alpha=0.05, wealth=0.025, tau=0.5)

        alpha, decision = [], []
        for i, p_value in enumerate(self.DATA["p_value"]):
            result = lord.test_one(p_value)
            alpha.append(round(lord.alpha, 6) if lord.alpha is not None\
        else None)  # fmt: skip
            decision.append(result)

        self.assertEqual(
            alpha,
            [
                0.001338,
                0.000291,
                0.000248,
                0.000206,
                0.000175,
                0.000151,
                None,  # Discarded
                0.001471,
                0.00041,
                0.000355,
                None,
                0.001641,
                0.000555,
                None,  # Discarded
                None,  # Discarded
            ],
        )

        self.assertEqual(
            decision,
            [
                True,
                False,
                False,
                False,
                False,
                True,
                False,  # Discarded
                False,
                False,
                True,
                False,  # Discarded
                False,
                True,
                False,  # Discarded
                False,  # Discarded
            ],
        )

    def test_lord_dependent(self):
        lord = LordDependent(alpha=0.05, wealth=0.025, reward=0.025)

        alpha, decision = [], []
        for i, p_value in enumerate(self.DATA["p_value"]):
            result = lord.test_one(p_value)
            alpha.append(round(lord.alpha, ndigits=6))
            decision.append(result)

        self.assertEqual(
            alpha,
            [
                0.010458,
                0.008270,
                0.001971,
                0.000736,
                0.000376,
                0.000227,
                0.000211,
                0.000151,
                0.000114,
                0.000089,
                0.000094,
                0.000077,
                0.000065,
                0.000055,
                0.000048,
            ],
        )

        self.assertEqual(
            decision,
            [
                True,
                True,
                False,
                False,
                False,
                True,
                False,
                False,
                False,
                True,
                False,
                False,
                False,
                False,
                False,
            ],
        )

    def test_lord_memory_decay(self):
        """Disclaimer: No reference solution available!"""
        lord = LORDMemoryDecay(alpha=0.05, delta=0.99, eta=0.001)

        alpha, decision = [], []
        for i, p_value in enumerate(self.DATA["p_value"]):
            result = lord.test_one(p_value)
            alpha.append(round(lord.alpha, ndigits=6))
            decision.append(result)

        self.assertEqual(
            alpha,
            [
                3e-06,     # First test
                0.00265,   # After first rejection
                0.00322,   # With decay contributions
                0.001052,
                0.000877,
                0.000729,
                0.003267,  # After more rejections
                0.001103,
                0.000948,
                0.000811,
                0.003354,
                0.001192,
                0.001036,
                0.003546,
                0.001358,
            ],
        )

        self.assertEqual(
            decision,
            [
                True,   # p=0.000000 <= alpha
                True,   # p=0.001430 <= alpha  
                False,
                False,
                False,
                True,   # p=0.000036 <= alpha
                False,
                False,
                False,
                True,   # p=0.000000 <= alpha
                False,
                False,
                True,   # p=0.000000 <= alpha
                False,
                False,
            ],
        )


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