import os
import re
import gzip
import numpy as np
from itertools import combinations

def arguments():
    """adding arguments for the script"""
    import argparse
    parser = argparse.ArgumentParser(description="""
    Given a directory that contains pair counts as generated by
    "motif_pairs_counts.py", this script creates a single file
    that is the symmetric double-sitecount matrix for all pairs of
    motifs within the directory.
    """)
    parser.add_argument('-d', '--dir', dest='dirname', action='store',
                        type=str, required=True,
                        help="""The input directory containing pair counts files""")
    parser.add_argument('-o', '--out', dest='outfile', action='store',
                        type=str, required=False,
                        help="""Optional output file name, if not given the results
                        appear on console""")

    args = parser.parse_args()
    return args


def main():
    args = arguments()
    motifs = np.array([f.replace('.gz', '').split("_with_")
                       for f in os.listdir(args.dirname)])
    motifs = np.sort(np.unique(motifs.flatten()))
    print len(motifs)
    
    index = dict([(motif, i) for i, motif in enumerate(motifs)])
    mat = np.zeros((len(motifs), len(motifs)), dtype=np.float32)
    for pair in combinations(motifs, 2):
        print pair
        fnames = [f for f in os.listdir(args.dirname) if re.search(pair[0], f)]
        try:
            fname = os.path.join(args.dirname, [f for f in fnames if re.search(pair[1], f)].pop())
        except IndexError:
            print pair
        del fnames
        double_counts = np.sum([float(l.split()[-1]) for l in
                              gzip.open(fname, 'r')])
        print fname, double_counts
        mat[index[pair[0]], index[pair[1]]] = \
            mat[index[pair[1]], index[pair[0]]] = double_counts
    if args.outfile:
        with open(args.outfile, 'w') as outf:
            outf.write('\t'.join(motifs) + '\n')
            for i, row in enumerate(mat):
                outf.write('\t'.join([
                    motifs[i],
                    '\t'.join(map(str, np.round(row, 6))) + '\n',
                ]))
    else:
        print '\t'.join(motifs)
        for i, row in enumerate(mat):
            print '\t'.join([
                motifs[i],
                '\t'.join(map(str, np.round(row, 6)))
                ])

    return 0


if __name__ == '__main__':
    main()
