#!/usr/bin/env python
import numpy
import pysam
import warnings
from pkg_resources import resource_filename
from nose.tools import assert_true, assert_equal, assert_greater_equal
from plastid.genomics.map_factories import FivePrimeMapFactory,\
                                       ThreePrimeMapFactory,\
                                       CenterMapFactory,\
                                       VariableFivePrimeMapFactory,\
                                       StratifiedVariableFivePrimeMapFactory
from plastid.genomics.roitools import GenomicSegment, SegmentChain
from plastid.genomics.genome_array import BAMGenomeArray
from plastid.util.services.mini2to3 import cStringIO


class TestBAM_MappingRules(object):
    @classmethod
    def setUpClass(cls):
        min_ = 25
        max_ = 40
        cls.strands = ("+", "-")
        cls.segs = {X: GenomicSegment("mock", 0, 2000, X) for X in cls.strands}

        cls.reads = {
            Y: [cls.make_alignment(0, X, Y) for X in range(min_, max_)]
            for Y in ("+", "-")
        }
        cls.expected = {}
        for mapping in ("fiveprime", "threeprime", "center"):
            for param in (0, 10):
                for strand in cls.strands:
                    cls.expected[(mapping, param, strand)] = numpy.zeros(2000)

        cls.expected[("fiveprime", 0, "+")][0] = max_ - min_
        cls.expected[("fiveprime", 10, "+")][10] = max_ - min_
        cls.expected[("fiveprime", 0, "-")][min_ - 1:max_ - 1] = 1
        cls.expected[("fiveprime", 10, "-")][min_ - 11:max_ - 11] = 1

        cls.expected[("threeprime", 0, "-")][0] = max_ - min_
        cls.expected[("threeprime", 10, "-")][10] = max_ - min_
        cls.expected[("threeprime", 0, "+")][min_ - 1:max_ - 1] = 1
        cls.expected[("threeprime", 10, "+")][min_ - 11:max_ - 11] = 1

        for my_len in range(min_, max_):
            cls.expected[("center", 0, "+")][:my_len] += 1.0 / my_len
            cls.expected[("center", 0, "-")][:my_len] += 1.0 / my_len
            cls.expected[("center", 10, "+")][10:my_len - 10] += 1.0 / (my_len - 2 * 10)
            cls.expected[("center", 10, "-")][10:my_len - 10] += 1.0 / (my_len - 2 * 10)

        cls.map_factories = {
            "fiveprime": FivePrimeMapFactory,
            "threeprime": ThreePrimeMapFactory,
            "fiveprime_variable": VariableFivePrimeMapFactory,
            "center": CenterMapFactory
        }

    @staticmethod
    def make_alignment(start_pos, end_pos, strand):
        read = pysam.AlignedSegment()
        read.reference_start = start_pos
        my_length = end_pos - start_pos
        read.query_sequence = "N" * my_length
        read.reference_id = 0
        read.is_reverse = True if strand == "-" else False
        read.cigarstring = "%sM" % my_length
        return read

    def check_mapping(self, map_name, map_param, strand):
        fn = self.map_factories[map_name](map_param)
        expected = self.expected[(map_name, map_param, strand)]
        reads_out, count_array = self.map_factories[map_name](map_param)(
            self.reads[strand], self.segs[strand]
        )
        msg1 = "failed to return reads with %s mapping with param %s on strand '%s'." % (
            map_name, map_param, strand
        )
        msg2 = "failed %s mapping vector with param %s on strand '%s'." % (
            map_name, map_param, strand
        )
        assert_equal(reads_out, self.reads[strand], msg1)
        assert_true((count_array == expected).all(), msg2)

    def test_fiveprime_threeprime_center(self):
        for mapping in ("fiveprime", "threeprime", "center"):
            for offset in 0, 10:
                for strand in self.strands:
                    yield self.check_mapping, mapping, offset, strand

    def test_fiveprime_variable(self):
        offset_dicts = {}
        fancy = {X: X // 2 for X in range(25, 40)}
        fancy_reads = {X: numpy.zeros(2000) for X in self.strands}
        for read in self.reads["+"]:
            fancy_reads["+"][read.positions[fancy[len(read.positions)]]] += 1
        for read in self.reads["-"]:
            fancy_reads["-"][read.positions[-fancy[len(read.positions)] - 1]] += 1

        fancy_plus = numpy.zeros(2000)
        fancy_minus = numpy.zeros(2000)

        offset_dicts[("default_only", "+")] = ({"default": 0}, self.expected[("fiveprime", 0, "+")])
        offset_dicts[("default_only", "-")] = ({"default": 0}, self.expected[("fiveprime", 0, "-")])
        offset_dicts[("fancy", "+")] = (fancy, fancy_reads["+"])
        offset_dicts[("fancy", "-")] = (fancy, fancy_reads["-"])

        for (name, strand), (dict_, expected) in offset_dicts.items():
            fn = VariableFivePrimeMapFactory(dict_)
            reads_out, count_array = fn(self.reads[strand], self.segs[strand])
            msg = "Failed fiveprime variable mapping for test %s(%s)" % (name, strand)
            assert_true((count_array == expected).all(), msg)

    def test_fiveprime_variable_from_file(self):
        offset_dicts = {}
        fancy = {X: X // 2 for X in range(25, 40)}
        fancy_reads = {X: numpy.zeros(2000) for X in self.strands}
        for read in self.reads["+"]:
            fancy_reads["+"][read.positions[fancy[len(read.positions)]]] += 1
        for read in self.reads["-"]:
            fancy_reads["-"][read.positions[-fancy[len(read.positions)] - 1]] += 1

        fancy_plus = numpy.zeros(2000)
        fancy_minus = numpy.zeros(2000)

        offset_dicts[("default_only", "+")] = ({"default": 0}, self.expected[("fiveprime", 0, "+")])
        offset_dicts[("default_only", "-")] = ({"default": 0}, self.expected[("fiveprime", 0, "-")])
        offset_dicts[("fancy", "+")] = (fancy, fancy_reads["+"])
        offset_dicts[("fancy", "-")] = (fancy, fancy_reads["-"])

        for (name, strand), (dict_, expected) in offset_dicts.items():
            fhlike = cStringIO.StringIO("\n".join(["%s\t%s" % (K, V) for K, V in dict_.items()]))
            fn = VariableFivePrimeMapFactory.from_file(fhlike)
            reads_out, count_array = fn(self.reads[strand], self.segs[strand])
            msg = "Failed fiveprime variable mapping for test %s(%s)" % (name, strand)
            assert_true((count_array == expected).all(), msg)

    def check_unmappable_raises_warnings(self, test_name, map_param, strand):
        map_factory = self.map_factories[test_name]
        msg = "Map function '%s' failed to raise warning for unmappable reads on strand '%s'" % (
            test_name, strand
        )
        with warnings.catch_warnings(record=True) as warns:
            warnings.simplefilter("always")
            fn = map_factory(map_param)
            reads_out, count_array = fn(self.reads[strand], self.segs[strand])

        assert_greater_equal(len(warns), 0, msg)

    def test_unmappable_raises_warnings(self):
        func_params = {
            "fiveprime": 30,
            "threeprime": 30,
            "center": 15,
            "fiveprime_variable": {
                25: 10,
                "default": 28
            }
        }
        for test_name, map_param in func_params.items():
            for strand in self.strands:
                yield self.check_unmappable_raises_warnings, test_name, map_param, strand

    def check_unmappable_not_mapped_in_vector_or_returned(
            self, test_name, map_param, expected_num_reads, strand
    ):
        msg1 = "Map function '%s' failed to filter out short reads on strand '%s' from returned reads" % (
            test_name, strand
        )
        msg2 = "Map function '%s' failed to filter out short reads on strand '%s' from count vector" % (
            test_name, strand
        )
        fn = self.map_factories[test_name](map_param)
        reads_out, count_array = fn(self.reads[strand], self.segs[strand])
        assert_equal(len(reads_out), expected_num_reads, msg1)
        assert_equal(count_array.sum(), expected_num_reads, msg2)

    def test_unmappable_not_mapped_in_vector_or_returned(self):
        func_params = {
            "fiveprime"  : 30,
            "threeprime" : 30,
            "center"     : 15,
            "fiveprime_variable": {
                25       : 10,
                "default": 28
            } # yapf: disable
        }
        expected_num = {
            "fiveprime":
            len([X for X in self.reads["+"] if len(X.positions) > func_params["fiveprime"]]),
            "threeprime":
            len([X for X in self.reads["+"] if len(X.positions) > func_params["threeprime"]]),
            "center":
            len([X for X in self.reads["+"] if len(X.positions) > 2 * func_params["center"]]),
            "fiveprime_variable":
            len([X for X in self.reads["+"] if len(X.positions) > 28 or len(X.positions) == 25]),
        }
        for test_name, map_param in func_params.items():
            expected_num_reads = expected_num[test_name]
            for strand in self.strands:
                yield self.check_unmappable_not_mapped_in_vector_or_returned, test_name, map_param, expected_num_reads, strand

    def test_variable_stratified_mapping_plus(self):
        offsets = {26: 6, 27: 22, 28: 13, 29: 4, 30: 5}

        chains = {
            "fw":
            SegmentChain(
                GenomicSegment('chrII', 392959, 393180, '+'),
                GenomicSegment('chrII', 393510, 394742, '+'),
                GenomicSegment('chrII', 394860, 394901, '+'),
                ID='YBR078W_mRNA'
            ),
            "rc":
            SegmentChain(
                GenomicSegment('chrVIII', 189061, 189749, '-'),
                GenomicSegment('chrVIII', 189850, 190017, '-'),
                ID='YHR041C_mRNA'
            )
        }
        expected = {
            "fw":
            numpy.loadtxt(
                resource_filename("plastid", "test/data/stratmap/strat_fw_vec.txt"), delimiter="\t"
            ),
            "rc":
            numpy.loadtxt(
                resource_filename("plastid", "test/data/stratmap/strat_rc_vec.txt"), delimiter="\t"
            ),
        }
        ga = BAMGenomeArray([resource_filename("plastid", "test/data/stratmap/strat.bam")])
        ga.set_mapping(StratifiedVariableFivePrimeMapFactory(offsets, 26, 30))

    def test_variable_stratified_mapping_minus(self):
        pass
