import sys
import unittest

import numpy as np
import pandas as pd

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


class Test(unittest.TestCase):
    def test_weights(self):
        """test that weights work"""
        dds = makeExampleDESeqDataSet(n=10, seed=42)
        w = np.ones(dds.shape)
        w[0, 0] = 0
        # check that weight to 0 is like a remove sample
        dds = DESeq(dds, quiet=True)
        dds2 = dds.copy()
        dds2.layers["weights"] = w
        dds2 = nbinomWaldTest(dds2)
        dds3 = dds.copy()[1:, :]
        dds3 = nbinomWaldTest(dds3)

        # in terms of LFC, SE and deviance
        self.assertAlmostEqual(
            dds2.results().log2FoldChange.iloc[0],
            dds3.results().log2FoldChange.iloc[0],
            delta=1e-5,
        )
        self.assertAlmostEqual(
            dds2.results().lfcSE.iloc[0], dds3.results().lfcSE.iloc[0], delta=1e-6
        )
        self.assertAlmostEqual(
            dds2.var["deviance"].iloc[0], dds3.var["deviance"].iloc[0], delta=1e-7
        )

        # check weights working in the optim code
        nf = np.repeat(dds.sizeFactors.values[:, None], dds.n_vars, axis=1)
        o = fitNbinomGLMsOptim(
            obj=dds,
            modelMatrix=dds.design,
            lambda_=np.repeat(1e-6, 2),
            colsForOptim=[0],
            colStable=np.repeat(True, dds.n_vars),
            normalizationFactors=nf,
            alpha_hat=dds.var["dispersion"],
            weights=w,
            useWeights=True,
            betaMatrix=pd.DataFrame(np.zeros((dds.n_vars, 2))),
            betaSE=pd.DataFrame(np.zeros((dds.n_vars, 2))),
            betaConv=np.repeat(False, dds.n_vars),
            beta_mat=np.zeros((dds.n_vars, 2)),
            mu=np.zeros(dds.shape),
            logLike=np.zeros(dds.n_vars),
        )
        self.assertAlmostEqual(
            dds3.results().log2FoldChange.iloc[0],
            o["betaMatrix"].iloc[0, 1],
            delta=1e-4,
        )

    def test_weights_with_beta_prior(self):
        """test that weights can be used with betaPrior=True"""
        dds = makeExampleDESeqDataSet(n=10, seed=42)
        w = np.ones(dds.shape)
        w[0, 0] = 0
        dds.layers["weights"] = w
        dds = DESeq(dds, betaPrior=True, quiet=True)

        # check weights working for intercept only
        dds.design = "~1"
        dds = DESeq(dds, quiet=True)
        dds2 = dds
        dds2.layers["weights"] = w
        dds2 = nbinomWaldTest(dds2)
        dds3 = dds[1:, :]
        dds3 = nbinomWaldTest(dds3)

        self.assertEqual(
            dds2.results().log2FoldChange.iloc[0], dds3.results().log2FoldChange.iloc[0]
        )
        self.assertEqual(dds2.results().lfcSE.iloc[0], dds3.results().lfcSE.iloc[0])
        self.assertEqual(dds2.var["deviance"].iloc[0], dds3.var["deviance"].iloc[0])

    def test_weights_downweight_outliers(self):
        """test that weights downweight outlier in dispersion estimation"""
        dds = makeExampleDESeqDataSet(n=10, seed=42)
        dds.counts()[0, 0] = 100
        dds.sizeFactors = np.ones(dds.n_obs)
        dds = dds.estimateDispersions()
        dds2 = dds.copy()
        w = np.ones(dds.shape)
        w[0, 0] = 0
        dds2.layers["weights"] = w
        dds2 = dds2.estimateDispersions()
        dds3 = dds.copy()[1:, :]
        dds3 = dds3.estimateDispersions()

        self.assertAlmostEqual(
            dds2.var["dispGeneEst"].iloc[0], dds3.var["dispGeneEst"].iloc[0], delta=1e-1
        )
        # MAP estimates won't be equal because of different dispersion prior widths
        self.assertGreater(dds.var["dispMAP"].iloc[0], dds2.var["dispMAP"].iloc[0])

    def test_weights_warning(self):
        """test that weights failing check gives warning, passes them through"""
        dds = makeExampleDESeqDataSet(n=10, seed=42)
        w = np.ones(dds.shape)
        w[0:6, 0] = 0
        dds.layers["weights"] = w
        with self.assertLogs("inmoose", level="WARNING") as logChecker:
            dds = DESeq(dds)
        self.assertRegex(
            # account for https://github.com/python/cpython/issues/86109
            logChecker.output[0]
            if sys.version_info >= (3, 10)
            else logChecker.output[3],
            "for 1 genes, the weights as supplied won't allow parameter estimation",
        )
        self.assertTrue(dds.var["allZero"].iloc[0])
        self.assertTrue(dds.var["weightsFail"].iloc[0])
        dds.results()

    @unittest.skip("not sure what is tested here")
    def test_weights_CR(self):
        """test that weights with and without CR term included"""
        alpha = 0.25

        def dmr(x):
            return alpha

        dds = makeExampleDESeqDataSet(
            n=50, m=100, betaSD=1, interceptMean=10, interceptSD=0.5, dispMeanRel=dmr
        )
        dds.obs["group"] = Factor(np.repeat(np.arange(50), 2))
        dds.design = "~0 + group + condition"
        w = np.ones(dds.shape)
        o = 35
        w[0:o, :] = 1e-6
        w[50 : (50 + o), :] = 1e-6
        dds.layers["weights"] = w
        dds.counts()[1:o, :] = 1
        dds.counts()[50 : (50 + o), :] = 1
        dds.sizeFactors = 1
        dds = dds.estimateDispersions(fitType="mean")
        dds.estimateDispersions(fitType="mean", useCR=False)