#
# Filters a GFF file using a list of attribute values. Attribute values must
# be in the first column of the file; subsequent columns are ignored.
# Usage:
# python gff_filter_by_attribute_values.py <gff_file> <attribute_name> <ids_file> <output_file>
#
from __future__ import print_function

import sys


def parse_gff_attributes(attr_str):
    """
    Parses a GFF/GTF attribute string and returns a dictionary of name-value
    pairs. The general format for a GFF3 attributes string is
        name1=value1;name2=value2
    The general format for a GTF attribute string is
        name1 "value1" ; name2 "value2"
    The general format for a GFF attribute string is a single string that
    denotes the interval's group; in this case, method returns a dictionary
    with a single key-value pair, and key name is 'group'
    """
    attributes_list = attr_str.split(";")
    attributes = {}
    for name_value_pair in attributes_list:
        # Try splitting by space and, if necessary, by '=' sign.
        pair = name_value_pair.strip().split(" ")
        if len(pair) == 1:
            pair = name_value_pair.strip().split("=")
        if len(pair) == 1:
            # Could not split for some reason -- raise exception?
            continue
        if pair == "":
            continue
        name = pair[0].strip()
        if name == "":
            continue
        # Need to strip double quote from values
        value = pair[1].strip(' "')
        attributes[name] = value

    if len(attributes) == 0:
        # Could not split attributes string, so entire string must be
        # 'group' attribute. This is the case for strictly GFF files.
        attributes["group"] = attr_str
    return attributes


def gff_filter(gff_file, attribute_name, ids_file, output_file):
    # Put ids in dict for quick lookup.
    ids_dict = {}
    for line in open(ids_file):
        ids_dict[line.split("\t")[0].strip()] = True

    # Filter GFF file using ids.
    with open(output_file, "w") as output, open(gff_file) as ingff:
        for line in ingff:
            if not line.rstrip() or line.startswith("#"):
                output.write(line)
                continue
            fields = line.split("\t")
            assert len(fields) == 9, f"Line should have exactly 9 fields: {line}"
            attributes = parse_gff_attributes(fields[8])
            if attribute_name in attributes and attributes[attribute_name] in ids_dict:
                output.write(line)


if __name__ == "__main__":
    # Handle args.
    if len(sys.argv) != 5:
        print("usage: python %s <gff_file> <attribute_name> <ids_file> <output_file>" % sys.argv[0], file=sys.stderr)
        sys.exit(-1)
    gff_file, attribute_name, ids_file, output_file = sys.argv[1:]
    gff_filter(gff_file, attribute_name, ids_file, output_file)
