#!/usr/bin/env python3

"""
Combine structural annotations made from different gff. 

In case of collapsing annotations, RNA annotation are prefered to CDS annotations.
Faa file from prodigal is processed to filter out overlapping sequences. 

:Example:
merge_annotations.py -c prodigal.gff -r barrnap.gff -t trnascan_se.gff --contig_seq contigs.fasta --faa_file prodigal.faa
"""

# Metadata
__author__ = 'Mainguy Jean - Plateforme bioinformatique Toulouse'
__copyright__ = 'Copyright (C) 2022 INRAE'
__license__ = 'GNU General Public License'
__version__ = '0.1'
__email__ = 'support.bioinfo.genotoul@inra.fr'
__status__ = 'dev'


from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter, FileType
import logging
import gzip
import csv
from collections import defaultdict
import pyfastx



def parse_arguments():
    """Parse script arguments."""
    parser = ArgumentParser(description="...",
                            formatter_class=ArgumentDefaultsHelpFormatter)

    parser.add_argument(
        '-c', '--cds', help='gff file with CDS annotations.', required=True)

    parser.add_argument(
        '-t', '--trna', help='gff file with tRNA annotations.', required=True)

    parser.add_argument(
        '-r', '--rrna', help='gff file with rRNA annotations.', required=True)

    parser.add_argument(
        '--contig_seq', help='fasta file of contigs. Needed to extract gene sequence', required=True)

    parser.add_argument(
        '--faa_file', help='fasta file of protein sequences generated by prodigal.', required=True)

    parser.add_argument(
        '--gff_output', help='final gff file with all annotations.', default="all_annotation.gff")

    parser.add_argument(
        '--ffn_output', help='final ffn file with all annotations.', default="all_annotation.ffn")

    parser.add_argument(
        '--faa_output', help='final faa file with all CDS annotations in amino acid.', default="all_annotation.faa")

    parser.add_argument(
        '--report_output', help='Prokka report like to be able to show annotations in multiqc.', default="annotation_report.txt")

    parser.add_argument("-v", "--verbose", help="increase output verbosity",
                        action="store_true")

    args = parser.parse_args()
    return args


def read_file_and_ignore_hashtag(file_path):

    proper_open = gzip.open if file_path.endswith('.gz') else open
    with proper_open(file_path, 'rt') as fl:
        for line in fl:
            if line.startswith('#'):
                continue
            yield line


def group_annotations_by_contigs(*gff_annotations):

    contig2annotations = defaultdict(list)
    for annotations in gff_annotations:
        for annotation in annotations:
            contig2annotations[annotation['seqname']].append(annotation)

    return contig2annotations


def parse_gff_file(gff_file, feature_to_keep=False):

    gff_headers = ("seqname", "_3", "feature", "start",
                   "end", "_2", "strand", "_1", "attribute")

    gff_annotations = csv.DictReader(read_file_and_ignore_hashtag(gff_file),
                                     delimiter='\t',
                                     fieldnames=gff_headers)

    if feature_to_keep:
        logging.info(f'Keeping only {feature_to_keep} feature from {gff_file}')
        gff_annotations = (
            annotation for annotation in gff_annotations if annotation['feature'] == feature_to_keep)

    return gff_annotations


def remove_overlapping_cds(cds_file, contig2rnas):

    contig_annotations = []
    current_contig = None

    for cds in parse_gff_file(cds_file, 'CDS'):

        reading_next_contig = cds['seqname'] != current_contig

        if contig_annotations and reading_next_contig:

            yield current_contig, contig_annotations
            contig_annotations = []

        current_contig = cds['seqname']

        is_overlapping = False
        for rna in contig2rnas[current_contig]:
            if (int(rna['end']) < int(cds['start']) or int(rna['start']) > int(cds['end'])):
                continue

            else:  # overlap -> remove cds
                is_overlapping = True
                overlap = f"[{max(rna['start'], cds['start'])},{min(rna['end'], cds['end'])}]"
                rna_info = f"{rna['feature']} [{rna['start']}, {rna['end']}]"
                cds_info = f"CDS [{cds['start']}, {cds['end']}]"
                logging.info(
                    f"overlap: {cds_info} overlapping with {rna_info} at {overlap} on contig={current_contig}")

        if not is_overlapping:
            contig_annotations.append(cds)

    yield current_contig, contig_annotations


def merging_cds_and_rna(cds_per_contig, contig2rnas):

    contig_processed = []
    for contig, cds_features in cds_per_contig:
        contig_processed.append(contig)
        rna_features = contig2rnas[contig]
        yield contig, rna_features + cds_features

    # check that all rrna contigs have been processed if not processe them
    for contig, rnas in contig2rnas.items():
        if contig not in contig_processed:
            logging.info(f'{contig} has only rna annotation')
            yield contig, rnas


def add_new_ID_tag_and_Parent(gff_attributes, new_id, parent_id):

    gff_attributes_no_id = ';'.join([attr for attr in gff_attributes.split(
        ';') if attr.split('=')[0] not in ['ID', 'Parent']])
    return f"ID={new_id};Parent={parent_id};{gff_attributes_no_id}"


def get_tag_value(gff_attribute, tag):

    for attr in gff_attribute.split(';'):
        if attr.split('=')[0] == tag:
            return attr.split('=')[1]
    return ''


def writing_features_to_gff_ffn_faa(annotations_per_contig, out_gff, fna_file, out_ffn, faa_file, out_faa):

    contig_fa = pyfastx.Fasta(fna_file)
    protein_fa = pyfastx.Fasta(faa_file)
    gene_count_all = defaultdict(int)
    with open(out_gff, 'w') as fh_gff, open(out_ffn, 'w') as fh_ffn, open(out_faa, 'w') as fh_faa:

        fh_gff.write("##gff-version  3\n")
        for contig, features in annotations_per_contig:
            gene_count = defaultdict(int)
            logging.info(f"writing {contig} annotations to gff file")
            fh_gff.write(f"##sequence-region {contig}\n")

            for feature in sorted(features, key=lambda x: int(x['start'])):
                gene_count[feature['feature']] += 1
                gene_count_all[feature['feature']] += 1
                new_id = f"{feature['seqname']}.{feature['feature']}_{gene_count[feature['feature']]}"
                parent_id = f"{new_id}_gene"

                start, end = int(feature['start']), int(feature['end'])

                # write faa
                if feature['feature'] == "CDS":
                    gff_id = get_tag_value(feature['attribute'], 'ID')

                    # ID in gff from prodigal is <contig_number>_<cds_number>
                    # seq name in faa from prodigal is <contig_name>_<cds_number>
                    cds_number = gff_id.split('_')[-1]
                    cds_prodigal_id = f"{contig}_{cds_number}"
                    fh_faa.write(
                        f">{new_id} {start}-{end}\n{protein_fa[cds_prodigal_id]}\n")

                # Write gff and ffn
                feature['attribute'] = add_new_ID_tag_and_Parent(
                    feature['attribute'], new_id, parent_id)

                # gff line with CDS/tRNA or rRNA in feature column
                feature_gff_line = '\t'.join(feature.values())

                # gff line with gene in feature column
                feature['feature'] = 'gene'
                feature['attribute'] = f'ID={parent_id};locus_tag={new_id}'
                gene_gff_line = '\t'.join(feature.values())

                fh_gff.write(f'{gene_gff_line}\n{feature_gff_line}\n')

                tag_name = get_tag_value(feature['attribute'], 'Name')

                if feature['strand'] == "+":
                    feature_seq = contig_fa[contig][start-1:end].seq
                else:
                    # minus strand: reverse complement the sequence
                    feature_seq = contig_fa[contig][start-1:end].antisense

                fh_ffn.write(
                    f">{new_id} {start}-{end} {tag_name}\n{feature_seq}\n")

    # building prokka like report for multiqc
    report = {"organism": "NA",
              "contigs": len(contig_fa),
              "bases": contig_fa.size}
    report.update(gene_count_all)

    return report


def main():

    args = parse_arguments()

    if args.verbose:
        logging.basicConfig(
            format="%(levelname)s: %(message)s", level=logging.DEBUG)
        logging.info('Mode verbose ON')

    else:
        logging.basicConfig(format="%(levelname)s: %(message)s")

    trna_file = args.trna
    rrna_file = args.rrna
    cds_file = args.cds

    fna_file = args.contig_seq
    faa_file = args.faa_file

    out_gff = args.gff_output
    out_ffn = args.ffn_output
    out_faa = args.faa_output
    out_report = args.report_output

    logging.info('Parsing rRNA annoations.')
    rrna_annotations = parse_gff_file(rrna_file, 'rRNA')

    logging.info('Parsing tRNA annotations.')
    trna_annotations = parse_gff_file(trna_file, 'tRNA')

    logging.info('Grouping RNA annoations by contig.')
    contig2rnas = group_annotations_by_contigs(
        trna_annotations, rrna_annotations)

    logging.info('Removing CDS annoations overlapping a RNA annotations.')
    unoverlapping_cds_per_contig = remove_overlapping_cds(
        cds_file, contig2rnas)

    logging.info('Merge CDS and RNA annotations by contig.')
    annotation_per_contigs = merging_cds_and_rna(
        unoverlapping_cds_per_contig, contig2rnas)

    logging.info(f'Writting CDS and RNA annotations to gff file: {out_gff}.')
    report = writing_features_to_gff_ffn_faa(
        annotation_per_contigs, out_gff, fna_file, out_ffn, faa_file, out_faa)

    with open(out_report, "w") as fl:
        fl.write(''.join(f"{k}: {v}\n" for k, v in report.items()))


if __name__ == '__main__':
    main()
