import argparse

from datetime import datetime


def Parser():
    the_parser = argparse.ArgumentParser()
    the_parser.add_argument(
        '--input', action="store", type=str, help="input miRBase GFF3 file")
    the_parser.add_argument(
        '--output', action="store", type=str,
        help="output GFF3 file with converted mature mir coordinates")
    args = the_parser.parse_args()
    return args


def get_gff_header(gff_input_file):
    string_list = []
    for line in open(gff_input_file, "r"):
        if line[0] == '#':
            string_list.append(line)
    string_list.append('# generated by mature_mir_gff_translation.py %s\n#\n' %
                       str(datetime.now()))
    return ''.join(string_list)


def load_gff_in_dict(gff_input_file):
    '''
    Reads the gff3 file and return a dictionary of dictionaries
    with keys equal to standard gff3 fields (9)
    Note that the key of the primary dictionary is the ID
    '''
    gff_dict = {}
    for line in open(gff_input_file, "r"):
        if line[0] == "#":
            continue
        gff_fields = line[:-1].split("\t")
        ID = gff_fields[8].split("ID=")[1].split(";")[0]
        gff_dict[ID] = {}
        gff_dict[ID]["seqid"] = gff_fields[0]
        gff_dict[ID]["source"] = gff_fields[1]
        gff_dict[ID]["type"] = gff_fields[2]
        gff_dict[ID]["start"] = gff_fields[3]
        gff_dict[ID]["end"] = gff_fields[4]
        gff_dict[ID]["score"] = gff_fields[5]
        gff_dict[ID]["strand"] = gff_fields[6]
        gff_dict[ID]["phase"] = gff_fields[7]
        gff_dict[ID]["attributes"] = gff_fields[8]
        if "erives_from" in gff_dict[ID]["attributes"]:
            parent_primary_transcript = gff_dict[ID]["attributes"].split(
                "erives_from=")[1]
            parent_primary_transcript = gff_dict[parent_primary_transcript][
                                            "attributes"].split("Name=")[1]
            gff_dict[ID]["attributes"] = "%s;Parent_mir_Name=%s" % (
                gff_dict[ID]["attributes"], parent_primary_transcript)
    return gff_dict


def genome_to_mir_gff(gff_dict, output, header):
    '''
    Converts seqid field from chromosome to item Name
    Then converts coordinates relative to "miRNA_primary_transcript"
    Note that GFF files are 1-based coordinates
    '''
    for key in gff_dict:
        name = gff_dict[key]["attributes"].split("Name=")[1].split(";")[0]
        gff_dict[key]["seqid"] = name
        if "erives_from=" in gff_dict[key]["attributes"]:
            parent_ID = gff_dict[key]["attributes"].split(
                                        "erives_from=")[1].split(";")[0]
            gff_dict[key]["start"] = str(int(gff_dict[key]["start"])-int(
                                       gff_dict[parent_ID]["start"])+1)
            gff_dict[key]["end"] = str(int(gff_dict[key]["end"])-int(
                                     gff_dict[parent_ID]["start"])+1)
    hairpins = {}
    matures = {}
    # treats miRNA_primary_transcript coordinates
    # in a second loop to avoid errors in conversion
    for key in gff_dict:
        if gff_dict[key]["type"] == "miRNA_primary_transcript":
            gff_dict[key]["end"] = str(int(gff_dict[key]["end"])-int(
                                       gff_dict[key]["start"]) + 1)
            gff_dict[key]["start"] = '1'
            # now, do a dict[ID]=Name but only for miRNA_primary_transcript
            hairpins[key] = gff_dict[key]["attributes"].split(
                                                        "Name=")[1].split(
                                                        ";")[0]
        else:
            matures[key] = gff_dict[key]["attributes"].split(
                                                        "Name=")[1].split(
                                                         ";")[0]
    with open(output, "w") as output:
        output.write(header)
        for ID in sorted(hairpins, key=hairpins.get):
            output.write("\t".join([gff_dict[ID]["seqid"],
                         gff_dict[ID]["source"], gff_dict[ID]["type"],
                         gff_dict[ID]["start"], gff_dict[ID]["end"],
                         gff_dict[ID]["score"], gff_dict[ID]["strand"],
                         gff_dict[ID]["phase"], gff_dict[ID]["attributes"]]))
            output.write("\n")
            for id in sorted(matures, key=matures.get, reverse=True):
                if ID in gff_dict[id]["attributes"]:
                    output.write("\t".join([gff_dict[id]["seqid"],
                                 gff_dict[id]["source"], gff_dict[id]["type"],
                                 gff_dict[id]["start"], gff_dict[id]["end"],
                                 gff_dict[id]["score"],
                                 gff_dict[id]["strand"],
                                 gff_dict[id]["phase"],
                                 gff_dict[id]["attributes"]]))
                    output.write("\n")


def main(infile, outfile):
    gff_dict = load_gff_in_dict(infile)
    genome_to_mir_gff(gff_dict, outfile, get_gff_header(infile))


if __name__ == "__main__":
    args = Parser()
    main(args.input, args.output)