#!/usr/bin/env python
# -*- coding: utf-8 -*-
#

import sys
import os
import gzip
import collections
from pprint import pprint
import argparse
import vcf

def main():

    # Args
    args = parse_arguments()
    args.c = set(args.c)

    # Get all list of consequences
    # get_list_consequences(args.i)

    # Parse vcf
    vcf_reader = vcf.Reader(filename=args.i)
    # Get list of CSQ field keys
    csq_keys = vcf_reader.infos["CSQ"].desc.split("Format: ")[1].split("|")
    
    # Go through each record
    outdata = {}
    for record in vcf_reader:

        # Get CSQs
        csqs = record.INFO["CSQ"]
        # Make into dicts
        csqs = [dict(zip(csq_keys, fields.split("|"))) for fields in csqs]

        # Find overlap with desired consequences
        for csq in csqs:
            csq["Consequence"] = csq["Consequence"].split("&")
            if len(args.c & set(csq["Consequence"])) > 0:
                key = csq[args.g]
                # Add each of the alternative alleles to the ordered set
                for alt_allele in record.ALT:
                    # Make identifier
                    idx = "{0}:{1}_{2}/{3}".format(record.CHROM,
                                                   record.POS,
                                                   record.REF,
                                                   alt_allele)
                    try:
                        outdata[key].add(idx)
                    except KeyError:
                        outdata[key] = OrderedSet([idx])
                        
    # Output to tsv
    with open(args.o, "w") as out_h:
        for group_id in outdata:
            markers = "\t".join(list(outdata[group_id]))
            out_h.write("{0}\t{1}\n".format(group_id, markers))

    return 0

def get_list_consequences(invcf):

    consequences = set([])
    # Parse vcf
    vcf_reader = vcf.Reader(filename=invcf)
    # Get list of CSQ field keys
    csq_keys = vcf_reader.infos["CSQ"].desc.split("Format: ")[1].split("|")
    
    # Go through each record
    outdata = {}
    for record in vcf_reader:

        # Get CSQs
        csqs = record.INFO["CSQ"]
        # Make into dicts
        csqs = [dict(zip(csq_keys, fields.split("|"))) for fields in csqs]

        # Find overlap with desired consequences
        for csq in csqs:
            csq["Consequence"] = csq["Consequence"].split("&")
            for consequence in csq["Consequence"]:
                consequences.add(consequence)

    # Print and quit
    pprint(sorted(consequences))
    sys.exit()

def parse_arguments():
    """ Parses command line arguments.
    """
    # Create top level parser.
    parser = argparse.ArgumentParser(description="Make group file from VEP annotated VCF.")

    # Add options
    parser.add_argument('--i', metavar="<input>",
        help=('Input VCF'),
        required=True,
        type=str)
    parser.add_argument('--o', metavar="<output>",
        help=('Output make-file'),
        required=True,
        type=str)
    parser.add_argument('--g', metavar="<gene/transcript>",
        help=('Group markers by gene or transcript'),
        required=True,
        choices=['Gene', 'Feature'],
        type=str)
    parser.add_argument('--c', metavar="<consequence>",
        help=('Consequences to include in output'),
        required=True,
        nargs="+",
        type=str)

    # Parse the arguments
    args = parser.parse_args()

    # Parse the arguments
    return args


class OrderedSet(collections.MutableSet):

    def __init__(self, iterable=None):
        self.end = end = [] 
        end += [None, end, end]         # sentinel node for doubly linked list
        self.map = {}                   # key --> [key, prev, next]
        if iterable is not None:
            self |= iterable

    def __len__(self):
        return len(self.map)

    def __contains__(self, key):
        return key in self.map

    def add(self, key):
        if key not in self.map:
            end = self.end
            curr = end[1]
            curr[2] = end[1] = self.map[key] = [key, curr, end]

    def discard(self, key):
        if key in self.map:        
            key, prev, next = self.map.pop(key)
            prev[2] = next
            next[1] = prev

    def __iter__(self):
        end = self.end
        curr = end[2]
        while curr is not end:
            yield curr[0]
            curr = curr[2]

    def __reversed__(self):
        end = self.end
        curr = end[1]
        while curr is not end:
            yield curr[0]
            curr = curr[1]

    def pop(self, last=True):
        if not self:
            raise KeyError('set is empty')
        key = self.end[1][0] if last else self.end[2][0]
        self.discard(key)
        return key

    def __repr__(self):
        if not self:
            return '%s()' % (self.__class__.__name__,)
        return '%s(%r)' % (self.__class__.__name__, list(self))

    def __eq__(self, other):
        if isinstance(other, OrderedSet):
            return len(self) == len(other) and list(self) == list(other)
        return set(self) == set(other)

if __name__ == '__main__':

    main()
