from __future__ import division
from math import fabs
from Thermo import DNAThermoModel
import unittest


class ProbePair():
    """Represent a pair of Primers

    Paramenters:
    ------------
    id: String
        Intifier of the pair of primers

    a : Primer class
        Probe a

    b: Primer class
        Probe b

    met_diff: Float
              Difference in melting temperature between the two primers

    mean_amp_len: Float
                  Mean length of the putative amplicons generated for this
                  pair of primers

    inter_energy: Float
                  Gibbs free energy for the interaction between the pairs

    a_fscore: Float
              F-score for primer a

    b_fscore: Float
              F-score for primer b
    """

    def __init__(self, id, pb_a, pb_b):
        self.id = id
        self.a = pb_a
        self.b = pb_b
        if pb_a.mt is not None and pb_b.mt is not None:
            self.met_diff = round(fabs(pb_a.mt - pb_b.mt), 3)
        else:
            self.met_diff = 0
        self.amplicons = dict()

        self.mean_amp_len = None
        self.inter_energy = None
        self.a_fscore = None
        self.b_fscore = None

    def are_candidates(self, targets, min_targets_frac,
                       min_amp_size, max_amp_size,
                       na, mg, prim_con, temp, dghm_threshold):
        """Check the two Probes are candidates to be Primers

        Args:
        -----
        targets: List<Strings>
                 List of targets sequences to be amplified

        min_target_frac: Float
                         Minimum fraction of targets that the two Probes must
                         detect in order to be consired Primers

        min_amp_size: Integer
                      Amplicons generated by the two probes must be of length
                      greater or equal to this value

        max_amp_size: Integer
                      Amplicons generated by the two probes must be of length
                      less or equal to this value
        na: Float
            Na concentration in Molar

        mg: Float
            Mg concentration in Molar

        prim_con: Float
                  Primer concentration in Molar

        temp: Float
              Operation temperature in C

        dghm_threshold: Float
                        Gibbs free energy for Probes interaction in kcal/mol.
                        Pairs of probes with Gibbs free energy below
                        dghm_threshold will be discarded

        Returns:
        --------
        Integer: 0 if this Pair of probes is suitable for Primers, -1 if the
                 Pair doesn't meet the minimum target fraction, -2 if the
                 amplicons the Pair generates are weird and -3 if the Probes
                 are able to interact one each other

        """

        # get the list of hits sequences each probe is able to detect
        sp_a = self.a.get_hits_list()
        sp_b = self.b.get_hits_list()

        # all hit sequences both probes are able to detect
        detected_seqs = self.__shared_elems__(sp_a, sp_b)

        # targets sequences both probes are able to detect
        targets_detected = self.__shared_elems__(detected_seqs,
                                                 targets)

        # fraction of targets detected by the pair of probes
        tgt_frac = len(targets_detected) / len(targets)

        if tgt_frac < min_targets_frac:
            return -1
        else:
            if self.def_amplicons(min_amp_size, max_amp_size) == 1:
                return -2
            else:
                if self.can_interact(na, mg, prim_con, temp, dghm_threshold):
                    return -3
                else:
                    self.get_fscores(targets)
                    return 0

    def def_amplicons(self, min_len, max_len):
        """Defines the amplicons the Pairs is able to generate

        Args:
        -----
        min_len: Integer
                 Minium length allowed for amplicons
        max_len: Integer
                 Maximum length allowed for amplicons
        Return:
        Void

        List: (seq_id, [start, end])
               The amplicons generated are assigned to the amplicons parameter
               of this class

        Notes:
        ------
        This function uses the alignment info of each Probe to determine the
        amplicons. For each target sequence, generates a list of coordinates,
        of the type [start, end], indicating the start and the end positions
        of a given amplicon generated by the Probes in the target id. The info
        is stored in the class parameter amplicons, in the form
        amplicons[seq_id] = [[start, end], ..]
        """

        # get the list of hits sequences each probe is able to detect
        sp_a = self.a.get_hits_list()
        sp_b = self.b.get_hits_list()

        # all hit sequences both probes are able to detect
        detected_seqs = self.__shared_elems__(sp_a, sp_b)

        for seq_id in detected_seqs:
            amps = []
            algs_a = self.a.get_alignments(seq_id, by_strand=True)
            algs_b = self.b.get_alignments(seq_id, by_strand=True)
            combs = [[algs_a[0], algs_b[1]], [algs_b[0], algs_a[1]]]

            # there are just to combinations, (PA +,PB -) and (PB +, PA -)
            for i in range(0, 2):

                pos_aligns = combs[i][0]
                neg_aligns = combs[i][1]

                for pos in pos_aligns:
                    for neg in neg_aligns:
                        if pos[0] < neg[0] and pos[1] < neg[0]:
                            tmp_amp = (pos[0], neg[1])
                            if i == 0:
                                tmp_amp = (pos[0], neg[1])
                            else:
                                tmp_amp = (pos[0], neg[1])
                            amps.append(tmp_amp)
            ampl_len = [(x[1] - x[0] + 1) for x in amps]

            # if the probes generate more than one amplicon in the same
            # sequence or the amplicon is greater than allowed, then the probe
            # pair is discarded
            if len(amps) != 1 or \
               not min_len <= ampl_len[0] <= max_len:
                return 1

            self.amplicons[seq_id] = amps

        # get the mean amp length
        amp_values = self.amplicons.values()
        amp_len_sum = sum([x[0][1] - x[0][0] + 1 for x in amp_values])
        mean_amp_len = amp_len_sum / len(self.amplicons)

        self.mean_amp_len = round(mean_amp_len, 3)

        return 0

    def can_interact(self, na, mg, prim_con, temp, energy_threshold):
        """Determines if the Probes are able to interact one each other

        Args:
        -----
        na: Float
            Molar Na concentration

        mg: Float
            Molar Mg concentration

        prim_con: Float
                  Molar primer concentration

        temp: Float
              Temperature in C

        energy_threshold: Float
                          Energy threshold in kcal/mol to say the Probes are
                          interacting

        Return:
        -------
        Boolean: True if the Probes have an interaction energy below the
                 threshold, False other way
        Notes:
        To say the Probes are interacting, the free Gibbs energy for hetero
        dimer formation is measured
        """
        thermo_model = DNAThermoModel()
        self.inter_energy = thermo_model.calc_heterodimer(str(self.a.seq),
                                                          str(self.b.seq),
                                                          na, mg, prim_con,
                                                          temp)
        self.inter_energy = round(self.inter_energy, 3)
        if self.inter_energy < energy_threshold:
            return True
        else:
            return False

    def get_fscores(self, targets):
        """Measures the FScore for each probe

        Args:
        -----
        targets: List<String>
                 List of target sequences to be amplified

        Return:
        -------
        [Float, Float]: A list with the fscore values [fscore_a, f_score_b]
        """
        self.a_fscore = self.__get_ind_fscore(self.a, targets)
        self.b_fscore = self.__get_ind_fscore(self.b, targets)
        return (self.a_fscore, self.b_fscore)

    def __get_ind_fscore(self, probe, targets):
        """Get the Fscore for a given Probe
        Args:
        -----
            pb : Probe class
                 Probe to be evaluated
            targets : List<String>
                      ids of targets sequences
        Return:
        -------
            Float: Fscore value
        """
        sp = probe.get_hits_list()
        inter = len(self.__shared_elems__(sp, targets))
        pp = inter/len(sp)
        rp = inter/len(targets)
        v1 = (1/pp) + (1/rp)
        v2 = v1 * 0.5
        fscore = 1/v2
        return round(fscore, 3)

    def __shared_elems__(self, list_a, list_b):
        """obtains the shared elements of two lists
        Args:
        -----
        list_a: List<String>
                List of elemnts a
        list_b: List<String>
                List of elements b
        Return:
        -------
        List<String>: list of shared elements
        """
        return list(set(list_a) & set(list_b))

    def to_string(self):
        """Transform a ProbePair class into String
        Return:
        -------
        String: representation of the class in String form
        """
        res = "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s" % (self.id,
                                                      self.a.id,
                                                      self.b.id,
                                                      str(self.met_diff),
                                                      str(self.inter_energy),
                                                      str(self.a_fscore),
                                                      str(self.b_fscore),
                                                      self.a.seq,
                                                      self.b.seq)
        return res


class ProbePairUnittest(unittest.TestCase):
    """
    unittest for ProbePair class
    """

    def test_def_amplicons(self):
        from Probes import Probe
        pb1 = Probe("id1", "ATCG")
        pb2 = Probe("id2", "GTTA")

        pb1.add_alignment("seq1", 4, 7, "+", 1)
        pb1.add_alignment("seq2", 5, 8, "+", 0)
        pb1.add_alignment("seq3", 17, 20, "-", 1)
        pb1.add_alignment("seq4", 12, 15, "-", 2)
        pb1.add_alignment("seq5", 4, 7, "+", 1)

        pb2.add_alignment("seq1", 12, 15, "-", 1)
        pb2.add_alignment("seq2", 15, 18, "-", 0)
        pb2.add_alignment("seq3", 2, 5, "+", 1)
        pb2.add_alignment("seq4", 2, 5, "+", 2)
        pb2.add_alignment("seq4", 5, 8, "+", 2)
        pb2.add_alignment("seq5", 12, 15, "-", 1)

        pp = ProbePair("pair", pb1, pb2)
        self.assertEqual(pp.def_amplicons(10, 14), 1)

        pb1 = Probe("id1", "ATCG")
        pb2 = Probe("id2", "GTTA")

        pb1.add_alignment("seq1", 4, 7, "+", 1)
        pb1.add_alignment("seq2", 5, 8, "+", 0)
        pb1.add_alignment("seq5", 4, 7, "+", 1)

        pb2.add_alignment("seq1", 12, 15, "-", 1)
        pb2.add_alignment("seq2", 15, 18, "-", 0)
        pb2.add_alignment("seq5", 12, 15, "-", 1)

        pp = ProbePair("pair", pb1, pb2)
        self.assertEqual(pp.def_amplicons(10, 14), 0)

    def test_can_interact(self):
        from Probes import Probe
        pb1 = Probe("id1", "ATCG")
        pb2 = Probe("id2", "GTTA")
        pp = ProbePair("pair", pb1, pb2)
        self.assertEqual(pp.can_interact(50/1e3, 0, 50/1e9, 37, 0.63), True)
        self.assertEqual(pp.can_interact(50/1e3, 0, 50/1e9, 37, 0.61), False)

    def test_fscore(self):
        """
        """
        from Probes import Probe
        pb1 = Probe("id1", "ATCG")
        pb2 = Probe("id2", "GTTA")

        pb1.add_alignment("seq1", 4, 7, "+", 1)
        pb1.add_alignment("seq2", 5, 8, "+", 0)
        pb1.add_alignment("seq3", 17, 20, "-", 1)
        pb1.add_alignment("seq4", 12, 15, "-", 2)
        pb1.add_alignment("seq5", 4, 7, "+", 1)

        pb2.add_alignment("seq1", 12, 15, "-", 1)
        pb2.add_alignment("seq2", 15, 18, "-", 0)
        pb2.add_alignment("seq3", 2, 5, "+", 1)
        pb2.add_alignment("seq4", 2, 5, "+", 2)
        pb2.add_alignment("seq4", 5, 8, "+", 2)
        pb2.add_alignment("seq5", 12, 15, "-", 1)

        pp = ProbePair("pair", pb1, pb2)

        self.assertEqual(pp.get_fscores(["seq1", "seq2",
                                         "seq3", "seq4", "seq5"]), (1, 1))
        self.assertEqual(pp.get_fscores(["seq1", "seq2",
                                         "seq3", "seq4"]),
                         (0.889, 0.889))

    def test_are_candidates(self):
        from Probes import Probe
        targets = ["seq1", "seq2", "seq5"]
        pb1 = Probe("id1", "ATCG")
        pb2 = Probe("id2", "GTTA")

        pb1.add_alignment("seq1", 4, 7, "+", 1)
        pb1.add_alignment("seq2", 5, 8, "+", 0)
        pb1.add_alignment("seq5", 4, 7, "+", 1)

        pb2.add_alignment("seq1", 12, 15, "-", 1)
        pb2.add_alignment("seq2", 15, 18, "-", 0)
        pb2.add_alignment("seq5", 12, 15, "-", 1)

        pp = ProbePair("pair", pb1, pb2)
        self.assertEqual(pp.are_candidates(targets, 0.9, 10, 14,
                                           50/1e3, 0, 50/1e9, 37, 0.61), True)
        self.assertEqual(pp.are_candidates(targets, 0.9, 10, 14,
                                           50/1e3, 0, 50/1e9, 37, 0.63), False)

        pb1.add_alignment("seq3", 17, 20, "-", 1)
        pb1.add_alignment("seq4", 12, 15, "-", 2)
        pb2.add_alignment("seq3", 2, 5, "+", 1)
        pb2.add_alignment("seq4", 2, 5, "+", 2)
        pb2.add_alignment("seq4", 5, 8, "+", 2)

        self.assertEqual(pp.are_candidates(targets, 0.9, 10, 14,
                                           50/1e3, 0, 50/1e9, 37, 0.61), False)
