"""
Implementation of Nussinov's algorithm for finding maximum number of base-pair interactions.

The user can either specify an explicit strings, or a value n such that
a randomly generate string of length n is used, which by default will be length 15.
Examples:

   python secondary.py ACUAGGAAUGGAAGCC
   python secondary.py 15
   python secondary.py

a minimum separation gap can be specified optionally after an explicit first argument:

   python secondary.py ACUAGGAAUGGAAGCC 3
   python secondary.py 15 3

the default   
"""
import sys
from random import choice


def w(a,b):
    if {a,b} in ( {'A','U'}, {'C','G'} ):
        return 1   # complementary bases
    else:
        return 0

def buildTable(rna, gap):
    opt = {}
    for g in range(-1,len(rna)):
        # solve for all i,j with j = i + g
        for i in range(len(rna)-g):
            j = i + g
            if g <= gap:
                best = 0    # cannot make any pairs
            else:
                option1 = opt[i+1,j-1] + w(rna[i], rna[j])
                option2 = opt[i+1,j]
                option3 = opt[i,j-1]
                best = max(option1,option2,option3)
                for k in range(i+1,j):
                    if opt[i,k] + opt[k+1,j] > best:
                        best = opt[i,k] + opt[k+1,j]
            opt[i,j] = best
    return opt


def pairs(rna, gap, opt, i, j):
    """Return list of pairs for an optimal pairing."""
    if j - i <= gap:
        return []      # no pairs possible
    elif opt[i,j] == opt[i+1,j]:
        return pairs(rna, gap, opt, i+1, j)
    elif opt[i,j] == opt[i,j-1]:
        return pairs(rna, gap, opt, i, j-1)
    elif opt[i,j] == opt[i+1,j-1] + w(rna[i],rna[j]):
        return [ (i,j) ] + pairs(rna, gap, opt, i+1, j-1)
    else:
        for k in range(i+1,j):
            if opt[i,j] == opt[i,k] + opt[k+1,j]:
                return pairs(rna,gap,opt,i,k) + pairs(rna,gap,opt,k+1,j)


def parseArgs():
    # set defaults
    rna = genRandom(15)
    gap = 0

    if len(sys.argv) >= 2:
        if sys.argv[1].isdigit():
            n = int(sys.argv[1])
            rna = genRandom(n)
        else:
            rna = sys.argv[1]

        if len(sys.argv) >= 3:
            gap = max(0,int(sys.argv[2]))

    return rna,gap

def genRandom(n):
    return ''.join(choice('ACGU') for _ in range(n))


if __name__ == '__main__':
    rna,gap = parseArgs()
    print(rna)

    opt = buildTable(rna,gap)

    print('       ' + '  '.join(rna))
    for i in range(len(rna)):
        print('%2d '%i + rna[i],end=' ')
        for j in range(len(rna)):
            if j < i:
                msg = ''
            else:
                msg = opt[i,j]
            print('%3s' % msg,end='')
        print()
    

    print()
    print(opt[0,len(rna)-1])
    print(pairs(rna,gap,opt,0,len(rna)-1))