import unittest

import numpy as np

from inmoose.deseq2 import (
    DESeqDataSet,
    estimateDispersionsFit,
    estimateDispersionsGeneEst,
    estimateDispersionsMAP,
    makeExampleDESeqDataSet,
    nbinomWaldTest,
)


class Test(unittest.TestCase):
    def test_linear_mu(self):
        """test that the use of linear model for fitting mu works as expected"""

        dds = makeExampleDESeqDataSet(
            n=100,
            m=4,
            interceptMean=10,
            interceptSD=3,
            dispMeanRel=lambda x: 0.5,
            sizeFactors=[0.5, 1, 1, 2],
        )
        dds = dds.estimateSizeFactors()
        dds1 = DESeqDataSet(dds.copy())
        dds1 = estimateDispersionsGeneEst(dds1, linearMu=False)
        dds2 = estimateDispersionsGeneEst(dds.copy(), linearMu=True)
        mu1 = dds1.layers["mu"]
        mu2 = dds2.layers["mu"]
        cors = np.diag(np.corrcoef(mu1, mu2)[: mu1.shape[0], mu1.shape[1] :])
        self.assertTrue(np.all(cors > 1 - 1e-6))

        dds2 = estimateDispersionsFit(dds2, fitType="mean")
        dds2 = estimateDispersionsMAP(dds2)
        dds2 = nbinomWaldTest(dds2)
        dds2.results()