import unittest

import numpy as np

from inmoose.deseq2 import DESeq, makeExampleDESeqDataSet, nbinomWaldTest
from inmoose.utils import Factor


class Test(unittest.TestCase):
    def test_addMLE(self):
        """test that adding MLE works as expected"""
        dds = makeExampleDESeqDataSet(n=200, m=12, betaSD=1, seed=42)
        dds.obs["condition"] = Factor(np.repeat(["A", "B", "C"], 4))
        dds.design = "~condition"
        dds = DESeq(dds, betaPrior=True)
        ddsNP = nbinomWaldTest(dds.copy(), betaPrior=False)

        res1 = dds.results(contrast=["condition", "C", "A"], addMLE=True)
        res2 = ddsNP.results(contrast=["condition", "C", "A"])
        self.assertTrue(np.allclose(res1.lfcMLE, res2.log2FoldChange, equal_nan=True))

        res1 = dds.results(contrast=["condition", "A", "B"], addMLE=True)
        res2 = ddsNP.results(contrast=["condition", "A", "B"])
        self.assertTrue(np.allclose(res1.lfcMLE, res2.log2FoldChange, equal_nan=True))

        res1 = dds.results(contrast=["condition", "C", "B"], addMLE=True)
        res2 = ddsNP.results(contrast=["condition", "C", "B"])
        self.assertTrue(np.allclose(res1.lfcMLE, res2.log2FoldChange, equal_nan=True))