#!/usr/bin/env python3
import collections
import os
import sys
import glob
import argparse
import itertools
import gzip


def decodeKvPairs(kv):
    keyValues = {}
    for part in kv.split(';'):
        kv = part.strip().split()
        if len(kv) == 2:
            key = kv[0]
            value = kv[1].replace('"', '')
            keyValues[key] = value
    return keyValues


def exonGTF_to_intronGTF(exon_path, id):
    # (gene,chrom,strand)-> [range(start,end), range(start,end)]
    geneToExonRanges = collections.defaultdict(list)
    prevChrom = None
    capture = list(set(['gene_id',
                        'gene_version',
                        'transcript_version',
                        'gene_name',
                        'gene_source',
                        'gene_biotype',
                        'transcript_name',
                        'transcript_source',
                        'transcript_biotype',
                        'transcript_support_level']) - set([id]))
    print(capture)

    id_to_features = {}

    with gzip.open(exon_path, 'rt') if exon_path.endswith('.gz') else open(exon_path) as f:
        for line in f:
            if line.startswith('#'):
                yield line
                continue
            chromosome, source, feature_type, feature_start, feature_end, at5, strand, at7, kv = line.split(
                None, 8)

            if feature_type != 'exon':
                continue

            meta = decodeKvPairs(kv)
            id_to_features[meta[id]] = {
                c: meta.get(c) for c in capture if c in meta
            }
            geneToExonRanges[(
                meta[id], chromosome, strand, source
            )
            ].append(range(int(feature_start), int(feature_end) + 1))

            if prevChrom is not None and prevChrom != chromosome:
                for intr in generate_introns(
                        geneToExonRanges, id, id_to_features):
                    yield intr
                geneToExonRanges = collections.defaultdict(list)

            prevChrom = chromosome

    if prevChrom is not None:
        for intr in generate_introns(geneToExonRanges, id, id_to_features):
            yield intr
        geneToExonRanges = collections.defaultdict(list)


def generate_introns(geneToExonRanges, id, id_to_features):
    for g in geneToExonRanges:
        if len(geneToExonRanges[g]) == 0:
            continue
        geneStart = min([min(r) for r in geneToExonRanges[g]])
        geneEnd = max([max(r) for r in geneToExonRanges[g]])
        gene_identifier_or_other, chromosome, strand, source = g

        introns = []
        in_intron = None
        current_intron_start = None
        rest_data = id_to_features.get(gene_identifier_or_other, {})

        if len(rest_data) > 0:
            rest = '; '.join(
                [f'{key} {value}' for key, value in rest_data.items()])
        else:
            rest = ''

        for i in sorted(

            [geneStart, geneEnd] +
            [i.start for i in geneToExonRanges[g]] +
                [i.stop for i in geneToExonRanges[g]]):
            if any([i in r for r in geneToExonRanges[g]]):
                if current_intron_start is not None and in_intron is True:
                    introns.append([current_intron_start, i])

                    yield '\t'.join([str(x) for x in [
                        chromosome,
                        source,
                        'intron',
                        current_intron_start,
                        i - 1,
                        '.',
                        strand,
                        '.', f'{id} "{gene_identifier_or_other}";{rest}']]) + '\n'

                in_intron = False
            else:
                if in_intron is False:
                    in_intron = True
                    current_intron_start = i


if __name__ == '__main__':
    argparser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Convert exonic GTF file to intronic GTF file")
    argparser.add_argument('gffin')
    argparser.add_argument('-id', help='identifier', default='transcript_id')
    argparser.add_argument(
        '-o',
        help="Output GTF file",
        type=str,
        required=True)
    args = argparser.parse_args()

    with gzip.open(args.o, 'wt') if args.o.endswith('.gz') else open(args.o, 'w') as o:
        for intron in exonGTF_to_intronGTF(args.gffin, args.id):
            if intron is not None:
                o.write(intron)