import unittest

import numpy as np

from patsy import dmatrix

from inmoose.limma import makeContrasts

class Test(unittest.TestCase):

    def test_makeContrasts(self):

        self.assertTrue(

            np.array_equal(

                makeContrasts(contrasts=["B-A", "C-B", "C-A"], levels=["A", "B", "C"]),

                [[-1, 0, -1], [1, -1, 0], [0, 1, 1]],

            )

        )

        self.assertTrue(

            np.array_equal(

                makeContrasts(contrasts="A-(B+C)/2", levels=["A", "B", "C"]),

                [[1.0], [-0.5], [-0.5]],

            )

        )

        design = dmatrix("0+A+B", {"A": ["a1", "a2", "a1"], "B": ["b1", "b2", "b3"]})

        self.assertTrue(

            np.array_equal(

                makeContrasts(contrasts="A[a1]-A[a2]", levels=design),

                [[1.0], [-1.0], [0.0], [0.0]],

            )

        )

        # test support for underscores in contrasts

        self.assertTrue(

            np.array_equal(

                makeContrasts(

                    contrasts=["T_T[B_1]-T[A_0]", "T[C_2]-T_T[B_1]", "T[C_2]-T[A_0]"],

                    levels=["T[A_0]", "T_T[B_1]", "T[C_2]"],

                ),

                [[-1, 0, -1], [0, 1, 1], [1, -1, 0]],

            )

        )