#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import unittest
from singlecellmultiomics.bamProcessing.bamCountRegions import count_regions
import tempfile
import pysam
from collections import Counter,defaultdict
import pandas as pd
import os

class TestRegionCount(unittest.TestCase):


    def test_region_counter(self):
        regions = {
            'regA':('chr1', 100,200),
            'regB':('chr1', 400,500),
            'regC':('chr1', 450,500),
        }

        reads = {
            # read name, cell index, library, mate, contig, cut location, strand : should be assigned  to:
            ('readA' , 1, 'LIBA', 1, 'chr1', 100, '+'): 'regA',
            ('readB' , 1, 'LIBA', 1, 'chr1', 150, '+'): 'regA',
            ('readC' , 1, 'LIBA', 1, 'chr1', 199, '+'): 'regA',
            ('readD' , 1, 'LIBA', 1, 'chr1', 200, '+'): None,
            ('readE' , 1, 'LIBA', 1, 'chr1', 210, '+'): None,
            ('readE' , 1, 'LIBA', 1, 'chr1', 399, '+'): None,
            ('readF' , 1, 'LIBA', 1, 'chr1', 400, '+'): 'regB',
            ('readG' , 1, 'LIBA', 1, 'chr1', 450, '+'): ['regB','regC'],
            ('readH' , 1, 'LIBB', 1, 'chr1', 450, '+'): ['regB','regC']

        }

        # Create bed file
        temp_bed_file = tempfile.NamedTemporaryFile(mode='w')
        for reg, (contig,start,end) in regions.items():
            temp_bed_file.write(f'{contig}\t{start}\t{end}\t{reg}\n')
        temp_bed_file.flush()

        # Create BAM file
        refseq = 'TTAATCATGAAACCGTGGAGGCAAATCGGAGTGTAAGGCTTGACTTTTAAGGGTTGAGATTTTCGAGAGGGATTCCTACGTTGCGTAGGTTCATGGGGGG'*10

        temp_bam_path = './temp_counting.bam'
        with pysam.AlignmentFile(
            temp_bam_path,
            'wb',
            reference_names=['chr1'],
            reference_lengths=[500]) as bam:

            for (read_name, cell_index, library, mate, contig, cut_location, strand) in reads:
                rlen = 20
                read_A = pysam.AlignedSegment(bam.header)
                read_A.reference_name = contig
                read_A.reference_start = cut_location
                # Before last A is a bogus G>A conversion to test strandness:
                read_A.query_sequence = refseq[cut_location:cut_location+rlen]
                read_A.cigarstring = f'{len(read_A.query_sequence)}M'
                read_A.qual = 'A'*len(read_A.query_sequence)
                read_A.mapping_quality = 60
                read_A.query_name = read_name
                read_A.set_tag('SM', f'{library}_{cell_index}')
                read_A.set_tag('LY', library)
                read_A.set_tag('DS', cut_location)
                read_A.set_tag('bi', cell_index)
                read_A.is_read1 = (mate==1)
                read_A.is_read2 = (mate!=1)

                read_A.is_paired = False
                read_A.is_proper_pair = True

                bam.write(read_A)
        pysam.index(temp_bam_path)

        ## Call the function to test:
        df = count_regions(temp_bed_file.name, bam_paths=[temp_bam_path,])
        temp_bed_file.close()


        targets = defaultdict(Counter)
        for (read_name, cell_index, library, mate, contig, cut_location, strand), goals in reads.items():
            if goals is None:
                continue
            if type(goals)!=list:
                goals = [goals]
            for goal in goals:
                targets[f'{library}_{cell_index}'][goal] += 1
        dfob = pd.DataFrame(targets).fillna(0)

        for feature, row in dfob.iterrows():
            for sample,value in row.items():
                self.assertTrue( df.loc[sample][feature] == value )

        os.remove(temp_bam_path)
        os.remove(temp_bam_path+'.bai')

if __name__ == '__main__':
    unittest.main()