#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# doctest: +NORMALIZE_WHITESPACE
# doctest: +SKIP
'''This module provides functions for PCR. Primers with 5' tails as well as inverse PCR on
circular templates are handled correctly.

'''
import math
import itertools

import re
import textwrap
import copy

import logging
module_logger = logging.getLogger("pydna."+__name__)

from Bio.SeqRecord                  import SeqRecord
from Bio.SeqUtils                   import GC
from Bio.SeqUtils.MeltingTemp       import Tm_NN
from .utils                        import rc
from .dseqrecord                   import Dseqrecord
from ._pretty                       import pretty_str, pretty_unicode
from .tm                           import tmbresluc

def _annealing_positions(primer, template, limit=15):
    '''Finds the annealing position(s) for a primer on a template where the
    primer anneals perfectly with at least limit nucleotides in the 3' part.

    start is a position (integer)
    footprint1 and tail1 are strings.

    ::

        <- - - - - - - - - - template - - - - - - - - - - - - - >

        <------------- start ---->
     5'-...gctactacacacgtactgactgcctccaagatagagtcagtaaccacactcgat...3'
           ||||||||||||||||||||||||||||||||||||||||||||||||
                                  3'-gttctatctcagtcattggtgtATAGTG-5'

                                                        <tail>
                                     <---footprint----->
                                     <--------- primer ------>

    Parameters
    ----------
    primer : string
        The primer sequence 5'-3'

    template : string
        The template sequence 5'-3'

    limit : int = 15, optional
        footprint needs to be at least of length limit.

    Returns
    -------
    describe : list of tuples (int, string, string)
        [ (start1, footprint1, tail1), (start2, footprint2, tail2),..., ]
    '''

    if len(primer)<limit:
        return []
    prc = rc(primer)
    head = prc[:limit].upper()

    table = {"R":"(A|G)",
             "Y":"(C|T)",
             "S":"(G|C)",
             "W":"(A|T)",
             "K":"(G|T)",
             "M":"(A|C)",
             "B":"(C|G|T)",
             "D":"(A|G|T)",
             "H":"(A|C|T)",
             "V":"(A|C|G)",
             "N":"(A|G|C|T)"}

    for key in table:
        head=head.replace(key, table[key])

    positions = [m.start() for m in re.finditer('(?={})'.format(head), template, re.I)]

    if positions:
        tail = prc[limit:]
        length = len(tail)
        results = []
        for match_start in positions:
            tm = template[match_start+limit:match_start+limit+length]
            footprint = rc(template[match_start:match_start+limit]+"".join([b for a,b in itertools.takewhile(lambda x: x[0].lower()==x[1].lower(), list(zip(tail, tm)))]))
            results.append((match_start, footprint, primer[: len(primer) - len(footprint) ]))
        return results
    return []

class Amplicon(Dseqrecord):
    '''The Amplicon class holds information about a PCR reaction involving two
    primers and one template. This class is used by the Anneal class and is not
    meant to be instantiated directly.

    Parameters
    ----------
    forward_primer : SeqRecord(Biopython)
        SeqRecord object holding the forward (sense) primer

    reverse_primer : SeqRecord(Biopython)
        SeqRecord object holding the reverse (antisense) primer

    template : Dseqrecord
        Dseqrecord object holding the template (circular or linear)

    saltc : float, optional
        saltc = monovalent cations (mM) (Na,K..)
        default value is 50mM
        This is used for Tm calculations.

    forward_primer_concentration : float, optional
        primer concentration (nM)
        default set to 1000nM = 1µM
        This is used for Tm calculations.

    rc : float, optional
        primer concentration (nM)
        default set to 1000nM = 1µM
        This is used for Tm calculations.
    '''

    def __init__(    self,
                     record,
                     *args,
                     template=None,
                     forward_primer=None,
                     reverse_primer=None,
                     saltc=None,
                     fprimerc=1000.0,
                     rprimerc=1000.0,
                     **kwargs):

        super().__init__(record, *args, **kwargs)
        self.template = template
        self.forward_primer = forward_primer
        self.reverse_primer = reverse_primer
        self.fwd_primer = self.forward_primer
        self.rev_primer = self.reverse_primer
        self.fprimerc = fprimerc
        self.rprimerc = rprimerc
        self.saltc = saltc

    def __getitem__(self, sl):
        answer = copy.copy(self)
        answer.seq = answer.seq.__getitem__(sl)
        answer.seq.alphabet = self.seq.alphabet
        sr = SeqRecord("n"*len(self))
        sr.features = self.features
        answer.features = SeqRecord.__getitem__(sr, sl).features
        return answer

    def __repr__(self):
        '''returns a short string representation of the object'''
        return "Amplicon({})".format(self.__len__())

    def _repr_pretty_(self, p, cycle):
            p.text("Amplicon({})".format(self.__len__()))
            
    def _repr_html_(self):
        return "Amplicon({})".format(self.__len__())

    def flankup(self, flankuplength=50):
        '''Returns a Dseqrecord object containing flankuplength bases upstream of the forward primer footprint,
       Truncated if the template is not long enough.

       ::

        <--- flankup --->

                  5TAATAAactactgactatct3
                         ||||||||||||||
        acgcattcagctactgtactactgactatctatcg

       '''
        return self.template.seq[self.forward_primer.position-flankuplength-len(self.forward_primer.footprint):self.forward_primer.position-len(self.forward_primer.footprint)]

    def flankdn(self, flankdnlength=50):
        '''Returns a Dseqrecord object containing flankdnlength bases downstream of the reverse primer footprint.
       Truncated if the template is not long enough.

       ::

                                       <---- flankdn ------>

                        3actactgactatctTAATAA5
                         ||||||||||||||
        acgcattcagctactgtactactgactatctatcgtacatgtactatcgtat


       '''
        return self.template.seq[self.reverse_primer.position+len(self.reverse_primer.footprint):self.reverse_primer.position+flankdnlength+len(self.reverse_primer.footprint)]



    def figure(self):
        '''
        This method returns a simple figure of the two primers binding to a part
        of the template.

        ::

         5gctactacacacgtactgactg3
          |||||||||||||||||||||| tm 52.6 (dbd) 58.3
         5gctactacacacgtactgactg...caagatagagtcagtaaccaca3
         3cgatgatgtgtgcatgactgac...gttctatctcagtcattggtgt5
                                   |||||||||||||||||||||| tm 49.1 (dbd) 57.7
                                  3gttctatctcagtcattggtgt5



        Returns
        -------
        figure:string
             A string containing a text representation of the primers
             annealing on the template (see example above).


        Notes
        -----
        tm in the figure above is the melting temperature (tm) for each primer calculated according to
        SantaLucia 1998 [#]_.

        dbd is the tm calculation for enzymes with dsDNA binding domains like Pfu-Sso7d [#]_. See [#]_
        for more information.

        References
        ----------

        .. [#] J. SantaLucia, “A Unified View of Polymer, Dumbbell, and Oligonucleotide DNA Nearest-neighbor Thermodynamics,” Proceedings of the National Academy of Sciences 95, no. 4 (1998): 1460.
        .. [#] M. Nørholm, “A Mutant Pfu DNA Polymerprimerase Designed for Advanced Uracil-excision DNA Engineering,” BMC Biotechnology 10, no. 1 (2010): 21, doi:10.1186/1472-6750-10-21.
        .. [#] http://www.thermoscientificbio.com/webtools/tmc/

        '''

        tmf = Tm_NN(str(self.forward_primer.footprint),
                    dnac1=self.fprimerc,
                    Na=self.saltc)
        tmr = Tm_NN(str(self.reverse_primer.footprint),
                    dnac1=self.fprimerc,
                    Na=self.saltc)

        tmf_dbd = tmbresluc(str(self.forward_primer.footprint), primerc=self.fprimerc)
        tmr_dbd = tmbresluc(str(self.reverse_primer.footprint), primerc=self.rprimerc)

        f =   '''
            {sp1}5{faz}...{raz}3
             {sp3}{rap} tm {tmr} (dbd) {tmr_dbd}
            {sp3}3{rp}5
            5{fp}3
             {fap:>{fplength}} tm {tmf} (dbd) {tmf_dbd}
            {sp2}3{fzc}...{rzc}5
            '''.format( fp       = self.forward_primer.seq,
                        fap      = "|"*len(self.forward_primer.footprint),
                        fplength = len(self.forward_primer.seq),
                        tmf      = round(tmf,1),
                        tmr      = round(tmr,1),
                        tmf_dbd  = round(tmf_dbd,1),
                        tmr_dbd  = round(tmr_dbd,1),
                        rp       = self.reverse_primer.seq[::-1],
                        rap      = "|"*len(self.reverse_primer.footprint),
                        rplength = len(self.reverse_primer.seq),
                        faz      = self.forward_primer.footprint,
                        raz      = self.reverse_primer.footprint.reverse_complement(),
                        fzc      = self.forward_primer.footprint.complement(),
                        rzc      = self.reverse_primer.footprint[::-1],
                        sp1      = " "*(len(self.forward_primer.seq)-len(self.forward_primer.footprint)),
                        sp2      = " "*(len(self.forward_primer.seq)-len(self.forward_primer.footprint)),
                        sp3      = " "*(3+len(self.forward_primer.seq))
                       )
        return pretty_str(textwrap.dedent(f).strip("\n"))


    def program(self):

        '''Returns a string containing a text representation of a proposed
       PCR program using Taq or similar polymerase.

       ::

        Taq (rate 30 nt/s)
        Three-step|         30 cycles     |      |Tm formula: Biopython Tm_NN
        94.0°C    |94.0°C                 |      |SaltC 50mM
        __________|_____          72.0°C  |72.0°C|
        04min00s  |30s  \         ________|______|
                  |      \ 46.0°C/ 0min 1s|10min |
                  |       \_____/         |      |
                  |         30s           |      |4-8°C

       '''


        # Primer melting temperatures are calculated with the Tm_NN formula from
        # biopython
        # simple salt concentration correction is used and the template concentration
        # is ignored. dnac1 = primer concentration
        
        tmf = Tm_NN(str(self.forward_primer.footprint),
                    dnac1=self.fprimerc,
                    Na=self.saltc)
        tmr = Tm_NN(str(self.reverse_primer.footprint),
                    dnac1=self.fprimerc,
                    Na=self.saltc)

        # Ta calculation according to
        # Rychlik, Spencer, and Rhoads, 1990, Optimization of the anneal
        # ing temperature for DNA amplification in vitro
        # http://www.ncbi.nlm.nih.gov/pubmed/2243783
        # The formula described uses the length and GC content of the product and
        # salt concentration (monovalent cations).

        #GC_prod=GC(str(self.seq))

        tmp = 81.5 + 0.41*GC(str(self.seq)) + 16.6*math.log10(self.saltc/1000.0) - 675/len(self)
        tml = min(tmf,tmr)
        ta = 0.3*tml+0.7*tmp-14.9

        # Taq polymerase extension rate is set to 30 nt/s
        # see https://www.thermofisher.com/pt/en/home/life-science/pcr/pcr-enzymes-master-mixes/taq-dna-polymerase-enzymes/taq-dna-polymerase.html
        taq_extension_rate = 30  # seconds/kB PCR product length
        extension_time_taq = int(round(taq_extension_rate * len(self) / 1000)) # seconds
        f  = textwrap.dedent('''
                                 Taq (rate {rate} nt/s) 35 cycles             |{size}bp
                                 95.0°C    |95.0°C                 |      |Tm formula: Biopython Tm_NN
                                 |_________|_____          72.0°C  |72.0°C|SaltC {saltc:2}mM
                                 | 03min00s|30s  \         ________|______|Primer1C {forward_primer_concentration:3}µM
                                 |         |      \ {ta}°C/{0:2}min{1:2}s| 5min |Primer2C {reverse_primer_concentration:3}µM
                                 |         |       \_____/         |      |GC {GC_prod}%
                                 |         |         30s           |      |4-12°C'''.format(rate=taq_extension_rate,
                                                                                            forward_primer_concentration=self.fprimerc/1000,
                                                                                            reverse_primer_concentration=self.rprimerc/1000,
                                                                                            ta=round(ta,1),
                                                                                            saltc=self.saltc,
                                                                                            *divmod(extension_time_taq,60),
                                                                                            size= len(self.seq),
                                                                                            GC_prod= int(self.gc()) ))
        return pretty_unicode(f)

    def taq_program(self):
        return self._program()

    def dbd_program(self):
        '''Returns a string containing a text representation of a proposed
       PCR program using a polymerase with a DNA binding domain such as Pfu-Sso7d.

       ::

        Pfu-Sso7d (rate 15s/kb)             |{size}bp
        Three-step|          30 cycles   |      |Tm formula: Pydna tmbresluc
        98.0°C    |98.0°C                |      |SaltC 50mM
        __________|_____          72.0°C |72.0°C|Primer1C   1µM
        00min30s  |10s  \ 61.0°C ________|______|Primer2C   1µM
                  |      \______/ 0min 0s|10min |
                  |        10s           |      |4-12°C

       '''
        PfuSso7d_extension_rate = 15                #seconds/kB PCR product
        extension_time_PfuSso7d = PfuSso7d_extension_rate * len(self) / 1000  # seconds

        # The program returned is eaither a two step or three step progrem
        # This depends on the tm and length of the primers in the
        # original instructions from finnzyme. These do not seem to be

        tmf_dbd = tmbresluc(str(self.forward_primer.footprint), primerc=self.fprimerc)
        tmr_dbd = tmbresluc(str(self.reverse_primer.footprint), primerc=self.rprimerc)

        # Ta calculation for enzymes with dsDNA binding domains like Pfu-Sso7d
        # https://www.finnzymes.fi/tm_determination.html

        length_of_f = len(self.forward_primer.footprint)
        length_of_r = len(self.reverse_primer.footprint)

        if (length_of_f>20 and length_of_r>20 and tmf_dbd>=69.0 and tmr_dbd>=69.0) or (tmf_dbd>=72.0 and tmr_dbd>=72.0):
            f=textwrap.dedent(  '''
                                    Pfu-Sso7d (rate {rate}s/kb)
                                    Two-step|    30 cycles |      |{size}bp
                                    98.0°C  |98.0C         |      |Tm formula: Pydna tmbresluc
                                    _____ __|_____         |      |SaltC {saltc:2}mM
                                    00min30s|10s  \        |      |Primer1C {forward_primer_concentration:3}µM
                                            |      \ 72.0°C|72.0°C|Primer2C {reverse_primer_concentration:3}µM
                                            |       \______|______|GC {GC_prod}%
                                            |      {0:2}min{1:2}s|10min |4-12°C
                                 '''.format(rate = PfuSso7d_extension_rate,
                                            forward_primer_concentration = self.fprimerc/1000,
                                            reverse_primer_concentration = self.rprimerc/1000,
                                            saltc = self.saltc,
                                            *map(int,divmod(extension_time_PfuSso7d,60)),
                                            GC_prod= int(self.gc()),
                                            size = len(self.seq) ))
        else:

            if (length_of_f>20 and length_of_r>20):
                ta = min(tmf_dbd,tmr_dbd)+3
            else:
                ta = min(tmf_dbd,tmr_dbd)

            f=textwrap.dedent(  '''
                                    Pfu-Sso7d (rate {rate}s/kb)                 |{size}bp
                                    Three-step|          30 cycles   |      |Tm formula: Pydna tmbresluc
                                    98.0°C    |98.0°C                |      |SaltC {saltc:2}mM
                                    __________|_____          72.0°C |72.0°C|Primer1C {forward_primer_concentration:3}µM
                                    00min30s  |10s  \ {ta:.1f}°C ________|______|Primer2C {reverse_primer_concentration:3}µM
                                              |      \______/{0:2}min{1:2}s|10min |GC {GC_prod}%
                                              |        10s           |      |4-12°C
                                 '''.format(rate = PfuSso7d_extension_rate,
                                            size= len(self.seq),
                                            ta   = round(ta),
                                            forward_primer_concentration   = self.fprimerc/1000,
                                            reverse_primer_concentration   = self.rprimerc/1000,
                                            saltc= self.saltc,
                                            GC_prod= int(self.gc()),
                                            *map(int, divmod(extension_time_PfuSso7d,60)) ))
        return pretty_str(f)

    def pfu_sso7d_program(self):
        return self.dbd_program()
