#!/usr/bin/env python3

# -*- coding: utf-8 -*-

import unittest

from singlecellmultiomics.methylation import MethylationCountMatrix

"""

These tests check if the methylation module is working correctly

"""

class TestMethylationCountMatrix(unittest.TestCase):

    def test_matrix(self):

        mA = MethylationCountMatrix()

        mA['sample_A', ('chr1', 10, 20)][0] += 1

        mA['sample_A', ('chr2', 10, 20)][0] += 2

        self.assertTrue( len(mA.counts) == 1 )

        self.assertEqual(mA.counts['sample_A'].get(('chr1', 10, 20))[0], 1)

        self.assertEqual(mA.counts['sample_A'].get(('chr1', 10, 20))[1], 0)

        self.assertEqual(mA.counts['sample_A'].get(('chr2', 10, 20))[0], 2)

        self.assertEqual(mA.counts['sample_A'].get(('chr2', 10, 20))[1], 0)

        self.assertEqual(mA.get_bulk_frame('pd').loc[('chr1', 10, 20)].unmethylated , 1.0)

        self.assertEqual(mA.get_bulk_frame('pd').loc[('chr1', 10, 20)].methylated , 0)

        self.assertEqual(mA.get_bulk_frame('pd').loc[('chr1', 10, 20)].beta , 0 )

        mB = MethylationCountMatrix()

        mB['sample_A', ('chr2', 10, 20)][1] += 2

        mB['sample_B', ('chr2', 10, 20)][1] += 1

        mB['sample_B', ('chr2', 10, 20)][0] = 2

        mA.update(mB)

        self.assertEqual(len(mA.counts), 2)

        self.assertEqual(mA.get_frame('beta')[('chr2', 10, 20)]['sample_A'], 1)

        self.assertEqual(mA.get_frame('beta')[('chr2', 10, 20)]['sample_B'], 1 / 3)

        self.assertEqual(mA.get_bulk_frame()['n_samples'].max(), 2)

        self.assertEqual(mA.get_bulk_frame()['n_samples'].min(), 1)

        self.assertEqual(mA.get_bulk_frame()['methylated'].sum(), 3)

        self.assertEqual(mA.get_bulk_frame()['unmethylated'].sum(), 3)

if __name__ == '__main__':

    unittest.main()