from SearchTrees import *
from Score import *


def BruteForceMotifSearchAgain(DNA,t,n,l):
    # searches all possible alignments of 't'
    # DNA fragments (of length 'n') for the
    # 'l'-mer with the maximum consensus score
    s = [1 for i in xrange(t)]
    bestScore = Score(s, DNA, l)
    bestMotif = [x for x in s]
    while (True):
        s = NextLeaf(s,t,n-l+1)
        if (sum(s) == t):
            break
        score = Score(s, DNA, l)
        if (score > bestScore):
            print s, " = ", score
            bestScore = score
            bestMotif = [x for x in s]
    return bestMotif


def SimpleMotifSearch(DNA,t,n,l):
    # traverses the full search tree of
    # all possible alignments of 't'
    # DNA fragments (of length 'n') for the
    # 'l'-mer with the maximum consensus score
    s = [0 for i in xrange(t)]
    bestScore = 0
    i = 0
    while (i > 0 or bestScore == 0):
        if (i < t):
            s, i = NextVertex(s,i,t,n-l+1)
        else:
            score = Score(s, DNA, l)
            if (score > bestScore):
                print s, " = ", score
                bestScore = score
                bestMotif = [x for x in s]
            s, i = NextVertex(s,i,t,n-l+1)
    return bestMotif


def BranchAndBoundMotifSearch(DNA,t,n,l):
    # traverses a search tree of alignments
    # of 't' DNA fragments (of length 'n')
    # while pruning hopless excursions in
    # search for the 'l'-mer with the
    # maximum consensus score
    s = [0 for i in xrange(t)]
    bestScore = 0
    i = 0
    while (i > 0 or bestScore == 0):
        if (i < t):
            optimisticScore = Score(s, DNA, l) + (t-i)*l
            if (optimisticScore < bestScore):
                s, i = Bypass(s,i,t,n-l+1)
            else:
                s, i = NextVertex(s,i,t,n-l+1)
        else:
            score = Score(s, DNA, l)
            if (score >= bestScore):
                print s, " = ", score
                bestScore = score
                bestMotif = [x for x in s]
            s, i = NextVertex(s,i,t,n-l+1)
    return bestMotif


if __name__ == "__main__":
    # Example from onlline notes
    DNA = ["cctgatagacgctatctggctatccaggtacttaggtcctctgtgcgaatctatgcgtttccaaccat",
           "agtactggtgtacatttgatccatacgtacaccggcaacctgaaacaaacgctcagaaccagaagtgc",
           "aaacgttagtgcaccctctttcttcgtggctctggccaacgagggctgatgtataagacgaaaatttt",
           "agcctccgatgtaagtcatagctgtaactattacctgccacccctattacatcttacgtccatataca",
           "ctgttatacaacgcgtcatggcggggtatgcgttttggtcgtcgtacgctcgatcgttaccgtacggc"]

    for s in DNA:
        print s

    print "Branch-and-Bound Search = ",
    s = BranchAndBoundMotifSearch(DNA,5,68,8)
    print s, "=", Consensus(s,DNA,8)
    print "iterations =", getAndClearScoreCount()

    # Testing code
    DNA = [ "acgcctggaattcttcgtgcccacgagtaatg",
            "gactagtgacaaagtccccgcgatcgtaaccc",
            "ttctccacggccatcaaggaccctcctctgga",
            "cacacacactaggcggattcaccacgaagtta",
            "ctatccccacgttctgatatttttgggagacg" ]

    for s in DNA:
        print s

    print "Branch-and-Bound Search = ",
    s = BranchAndBoundMotifSearch(DNA,5,32,8)
    print s, "=", Consensus(s,DNA,8)
    print "iterations =", getAndClearScoreCount()

    DNA = [ "acgagtacccttcttc",
            "gccaggaccgttaatg",
            "gactacggaccctgtg",
            "aaggaccctgtaaccc",
            "tgtcggaccctcatca" ]

    for s in DNA:
        print s

    print "Brute Force:"
    s = BruteForceMotifSearchAgain(DNA,5,16,8)
    print s, "=", Consensus(s,DNA,8)
    print "iterations =", getAndClearScoreCount()

    print "Simple Brute Force:"
    s = SimpleMotifSearch(DNA,5,16,8)
    print s, "=", Consensus(s,DNA,8)
    print "iterations =", getAndClearScoreCount()

    print "Branch-and-Bound Search = ",
    s = BranchAndBoundMotifSearch(DNA,5,16,8)
    print s, "=", Consensus(s,DNA,8)
    print "iterations =", getAndClearScoreCount()

    raw_input("Press [Enter] to quit")