#!/usr/bin/env python3.2

from __future__ import division, print_function

import sys

from os import getcwd
from os.path import exists, join

from Bio import AlignIO

from hy454 import graph_logo


def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]

    try:
        labels = None
        for i in range(len(argv)-2, -1, -1):
            arg = argv[i]
            if arg[0] == '-':
                if arg[1] == 'l':
                    labels = argv[i+1].split(',')
                    del argv[i]
                    del argv[i]
                else:
                    assert(0)
        assert(len(argv) == 3)
        assert(exists(argv[0]))

        # now handle column specification
        colspecs = argv[1].split(',')
        columns = []
        for colspec in colspecs:
            if ':' in colspec:
                xyz = [int(v) for v in colspec.split(':')]
                xyz[1] += 1
                if len(xyz) in (2, 3):
                    columns.extend(range(*xyz))
                else:
                    raise ValueError()
            else:
                columns.append(int(colspec))
        assert(True if labels is None else len(labels) == len(columns))
    except:
        print('usage: seqlogo [-l LABELS] MSAFILE COLSPEC OUTPUTFILE', file=sys.stderr)
        print('  COLSPEC must be a comma-delimited list of 1-indices or colon-delimited range specifications start:stop[:skip]', file=sys.stderr)
        sys.exit(-1)

    filename = join(getcwd(), argv[2])

    with open(argv[0]) as fh:
        msa = AlignIO.read(fh, 'fasta')

    graph_logo(msa, columns, filename, format='pdf', labels=labels)

    return 0


if __name__ == '__main__':
    sys.exit(main())
