import unittest

import numpy as np

from inmoose.deseq2 import DESeq, makeExampleDESeqDataSet

from inmoose.utils import Factor

class Test(unittest.TestCase):

    def test_results(self):

        """test that results work as expected and throw errors"""

        # test contrasts

        dds = makeExampleDESeqDataSet(n=200, m=12, seed=42)

        dds.obs["condition"] = Factor(np.repeat([1, 2, 3], 4))

        dds.obs["group"] = Factor(np.repeat([[1, 2]], 6, axis=0).flatten())

        dds.obs["foo"] = np.repeat(["lo", "hi"], 6)

        dds.counts()[:, 0] = np.repeat([100, 200, 800], 4)

        dds.design = "~ group + condition"

        # calling results too early

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="could not find results in obj. first run DESeq()",

        ):

            dds.results()

        dds.sizeFactors = np.ones(dds.n_obs)

        dds = DESeq(dds)

        res = dds.results()

        # TODO

        # show_res = res.show()

        summary = res.summary()

        print(summary)

        summary_ref = """

out of 200 with nonzero total read count

adjusted p-value < 0.1

LFC > 0 (up)       : 1, 0.50%

LFC < 0 (down)     : 0, 0.00%

outliers [1]       : 0, 0.00%

low counts [2]     : 0, 0.00%

(mean count < 0)

[1] see 'cooksCutoff' argument of results()

[2] see 'independentFiltering' argument of results()

"""

        self.assertEqual(summary, summary_ref)

        # various results error checking

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="the LRT requires the user to run nbinomLRT or DESeq",

        ):

            dds.results(test="LRT")

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="when testing altHypothesis='lessAbs', set the argument lfcThreshold to a positive value",

        ):

            dds.results(altHypothesis="lessAbs")

        with self.assertRaisesRegex(

            ValueError, expected_regex="'name' should be a string"

        ):

            dds.results(name=["Intercept", "group1"])

        with self.assertRaisesRegex(ValueError, expected_regex="foo is not a factor"):

            dds.results(contrast=["foo", "B", "A"])

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="as 1 is the reference level, was expecting condition_4_vs_1 to be present in",

        ):

            dds.results(contrast=["condition", "4", "1"])

        with self.assertRaisesRegex(

            ValueError, expected_regex="invalid value for test: foo"

        ):

            dds.results(test="foo")

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="numeric contrast vector should have one element for every element of",

        ):

            dds.results(contrast=False)

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="'contrast', as a pair of lists, should have length 2",

        ):

            dds.results(contrast=["a", "b", "c", "d"])

        with self.assertRaisesRegex(

            ValueError, expected_regex="1 and 1 should be different level names"

        ):

            dds.results(contrast=["condition", "1", "1"])

        dds.results(independentFiltering=False)

        dds.results(contrast=["condition_2_vs_1"])

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="condition_3_vs_1 and condition_3_vs_1 should be different level names",

        ):

            dds.results(

                contrast=["condition_2_vs_1", "condition_3_vs_1", "condition_3_vs_1"]

            )

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="'contrast', as a pair of lists, should have lists of strings as elements",

        ):

            dds.results(contrast=["condition_2_vs_1", 1])

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="all elements of the 2-element contrast should be elements of",

        ):

            dds.results(contrast=["condition_2_vs_1", "foo"])

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="elements in the 2-element contrast should only appear in the numerator",

        ):

            dds.results(contrast=["condition_2_vs_1", "condition_2_vs_1"])

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="all elements of the 2-element contrast should be elements of",

        ):

            dds.results(contrast=["", ""])

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="numeric contrast vector should have one element for every element of",

        ):

            dds.results(contrast=np.repeat(0, 6))

        with self.assertRaisesRegex(ValueError, expected_regex="foo is not a factor"):

            dds.results(contrast=["foo", "lo", "hi"])

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "1", "3"]).log2FoldChange.iloc[0],

            -3,

            delta=1e-6,

        )

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "1", "2"]).log2FoldChange.iloc[0],

            -1,

            delta=1e-6,

        )

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "2", "3"]).log2FoldChange.iloc[0],

            -2,

            delta=1e-6,

        )

        # test a number of contrast as list options

        self.assertAlmostEqual(

            dds.results(

                contrast=["condition_3_vs_1", "condition_2_vs_1"]

            ).log2FoldChange.iloc[0],

            2,

            delta=1e-6,

        )

        dds.results(

            contrast=["condition_3_vs_1", "condition_2_vs_1"], listValues=[0.5, -0.5]

        )

        dds.results(contrast=["condition_3_vs_1", []])

        dds.results(contrast=["condition_3_vs_1", []], listValues=[0.5, -0.5])

        dds.results(contrast=[[], "condition_2_vs_1"])

        dds.results(contrast=[[], "condition_2_vs_1"], listValues=[0.5, -0.5])

        # test no prior on intercept

        self.assertTrue(np.array_equal(dds.betaPriorVar, np.repeat(1e6, 4)))

        # test thresholding

        dds.results(lfcThreshold=np.log2(1.5))

        dds.results(lfcThreshold=1, altHypothesis="lessAbs")

        dds.results(lfcThreshold=1, altHypothesis="greater")

        dds.results(lfcThreshold=1, altHypothesis="less")

        dds3 = DESeq(dds, betaPrior=True)

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="testing altHypothesis='lessAbs' requires setting the DESeq\(\) argument betaPrior=False",

        ):

            dds3.results(lfcThreshold=1, altHypothesis="lessAbs")

    def test_results_zero_intercept(self):

        """test results on designs with zero intercept"""

        dds = makeExampleDESeqDataSet(n=100, m=12, seed=42)

        dds.obs["condition"] = Factor(np.repeat([1, 2, 3], 4))

        dds.obs["group"] = Factor(np.repeat([[1, 2]], 6, axis=0).flatten())

        dds.X[:, 0] = np.repeat([100, 200, 400], 4)

        dds.design = "~ 0 + condition"

        dds = DESeq(dds, betaPrior=False)

        self.assertAlmostEqual(dds.results().log2FoldChange.iloc[0], 2, delta=0.1)

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "2", "1"]).log2FoldChange.iloc[0],

            1.25,

            delta=0.1,

        )

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "3", "2"]).log2FoldChange.iloc[0],

            0.68,

            delta=0.1,

        )

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "1", "3"]).log2FoldChange.iloc[0],

            -2,

            delta=0.1,

        )

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "1", "2"]).log2FoldChange.iloc[0],

            -1.25,

            delta=0.1,

        )

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "2", "3"]).log2FoldChange.iloc[0],

            -0.68,

            delta=0.1,

        )

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="condition\[4\] and condition\[1\] are expected to be in",

        ):

            dds.results(contrast=["condition", "4", "1"])

        dds.design = "~ 0 + group + condition"

        dds = DESeq(dds, betaPrior=False)

        self.assertAlmostEqual(dds.results().log2FoldChange.iloc[0], 2, delta=0.1)

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "3", "1"]).log2FoldChange.iloc[0],

            2,

            delta=0.1,

        )

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "2", "1"]).log2FoldChange.iloc[0],

            1.25,

            delta=0.1,

        )

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "3", "2"]).log2FoldChange.iloc[0],

            0.68,

            delta=0.1,

        )

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "1", "3"]).log2FoldChange.iloc[0],

            -2,

            delta=0.1,

        )

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "1", "2"]).log2FoldChange.iloc[0],

            -1.25,

            delta=0.1,

        )

        self.assertAlmostEqual(

            dds.results(contrast=["condition", "2", "3"]).log2FoldChange.iloc[0],

            -0.68,

            delta=0.1,

        )

    @unittest.skip("LRT is not implemented yet")

    def test_results_likelihood_ratio_test(self):

        """test results with likelihood ratio test"""

        dds = makeExampleDESeqDataSet(n=100)

        dds.obs["group"] = Factor([1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2])

        dds.design = "~ group + condition"

        dds = DESeq(dds, test="LRT", reduced="~group")

        self.assertFalse(

            np.all(

                dds.results(name="condition_B_vs_A").stat

                == dds.results(name="condition_B_vs_A", test="Wald").stat

            )

        )

        # LFC are already MLE

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="addMLE=TRUE is only for when a beta prior was used",

        ):

            dds.results(addMLE=True)

        with self.assertRaisesRegex(

            ValueError,

            expected_regex="tests of log fold change above or below a theshold must be Wald tests",

        ):

            dds.results(lfcThreshold=1, test="LRT")

        self.assertTrue(

            np.all(

                dds.results(test="LRT", contrast=["group", "1", "2"]).log2FoldChange

                == -dds.results(test="LRT", contrast=["group", "2", "1"]).log2FoldChange

            )

        )

    def test_results_basics(self):

        """test that results basics regarding format, saveCols, tidy, MLE, remove are working"""

        dds = makeExampleDESeqDataSet(n=100)

        dds.var["score"] = np.arange(1, 101)

        dds = DESeq(dds)

        # try saving metadata columns

        res = dds.results(saveCols="score")  # string

        # check tidy-ness (unimplemented)

        with self.assertRaises(NotImplementedError):

            res = dds.results(tidy=True)

            self.assertTrue(res.columns[0] == "rows")

        # test MLE and 'name'

        dds2 = DESeq(dds, betaPrior=True)

        dds2.results(addMLE=True)

        with self.assertRaises(ValueError):

            dds2.results(name="condition_B_vs_A", addMLE=True)

        # test remove results

        dds = dds.removeResults()

        self.assertTrue(dds.var.description.filter("results").empty)

    @unittest.skip("not sure what to test")

    def test_results_custom_filters(self):

        """test that custom filters can be provided to results()"""

        dds = makeExampleDESeqDataSet(n=200, m=4, betaSD=np.repeat([0, 2], [150, 50]))

        dds = DESeq(dds)

        _res = dds.results()

        _method = "BH"

        _alpha = 0.1

        raise NotImplementedError()