#! /usr/bin/env python
"""
Usage: ./good_adapter_combinations.py n m < adapter_list
Arguments:
    n:  size of the desired adapter sets (i.e. how many samples to multiplex)
    m:  how many adapter combinations to show
Reads adapter list in format 'index nr|index sequence'
Ranks adapter combinations by (1) how many positions/cycles
have close to maximal sequnce entropy (90%) (2) how
many adapter pairs have 4 or more differences.
"""

import sys
import itertools
import array
import math

def entropy_ratio(seq_list, max_h):
    hf = []
    for s in itertools.izip(*seq_list):
        hf.append(entropy(s) / max_h)
    return hf

def entropy(s):
    nf = float(len(s))
    h = 0
    for nt in ["A", "C", "G", "T"] :
        f = s.count(nt) / nf
        if f > 0:
            h += - f * math.log(f, 2)
    return h

def dist(a, b):
    return sum(1 for a, b in itertools.izip(a, b) if a != b)

def pairwise_diff(seq_list):
    d = []
    for pair in itertools.combinations(seq_list, 2):
        d.append(dist(*pair))
    return d

def max_entropy(n):
    max_h = 0
    for s in itertools.combinations_with_replacement("ACGT", n):
        h = entropy(s)
        if h > max_h:
            max_h = h
    return max_h
 
def rank_good_combinations(kit, n, max_h, adapter_length):
    combo_list = []
    for combo in itertools.combinations(kit, n):
        seq_list = [b for a, b in combo]
        hf = entropy_ratio(seq_list, max_h)
        f_hf = sum(1 for x in hf if x > 0.9) / float(adapter_length)
        d = pairwise_diff(seq_list)
        f_d = sum(1 for x in d if x >= 4) / float(len(d))
        combo_list.append((f_hf, f_d, combo))
    combo_list.sort(key = lambda x: (x[0], x[1]))
    combo_list.reverse()
    return combo_list
 
 
if __name__ == "__main__":
    if (len(sys.argv) != 3):
        print >>sys.stderr, __doc__
        sys.exit(1)
    n = int(sys.argv[1])
    m = int(sys.argv[2])
    #read in the list of adapters from stdin
    adapters = [x.strip().split("|") for x in sys.stdin]
    adapter_len_set = set(len(b) for a, b in adapters)
    if len(adapter_len_set) != 1:
        print >>sys.stderr, "adapters of differing lenghts found in input"
        sys.exit(1)
    adapter_len = adapter_len_set.pop()
    max_h = max_entropy(n)
    print "{:*^70}".format("  ADAPTER SET  ")
    print "\n".join("{}\t{}".format(a, b) for a, b in adapters)
    print "max sequence entropy at each cycle: {}".format(max_h)
    print "top {} sequences sorted by ".format(m)
    print "  (1) the fraction of positions with entropy > 90% of max. entropy"
    print "  (2) the fraction of pairs with 4 or more differences"
    print "*" * 70
    ranked_combinations = rank_good_combinations(adapters, n, max_h, adapter_len)
    for i in range(m):
        c = ranked_combinations[i]
        print "-" * 15
        print "{:4.0%} of positions had close to max. entropy".format(c[0])
        print "{:4.0%} of pairs had more than 4 differences".format(c[1])
        print "-" * 15
        print "\n".join("{}\t{}".format(a, b) for a, b in c[2])