#!/usr/bin/env python3
"""
Copyright 2019 Ryan Wick (rrwick@gmail.com)
https://github.com/rrwick/Compare-annotations/
This program is free software: you can redistribute it and/or modify it under the terms of the GNU
General Public License as published by the Free Software Foundation, either version 3 of the
License, or (at your option) any later version.
This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without
even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
General Public License for more details.
You should have received a copy of the GNU General Public License along with this program.  If not,
see <https://www.gnu.org/licenses/>.
"""

import argparse
from Bio import SeqIO
from Bio import pairwise2
from Bio.pairwise2 import format_alignment
import itertools
import sys


def addArr(arrA, arrB):
    res = []
    for x in range(0, min(len(arrA), len(arrB))):
      res.append(arrA[x] + arrB[x])
    return res

def get_arguments():
    parser = argparse.ArgumentParser(description='Compare GenBank annotations',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('annotation_1', type=str,
                        help='First annotated genome in Genbank format')
    parser.add_argument('annotation_2', type=str,
                        help='Second annotated genome in Genbank format')


    parser.add_argument('--match_identity_threshold', type=float, default=0.7,
                        help='Two genes must have at least this identity to be considerd the same (0.0 to 1.0)')
    parser.add_argument('--allowed_skipped_genes', type=int, default=10,
                        help='This many missing genes are allowed when aligning the annotations')
    parser.add_argument("--addNotes", action="store_true", help="Add Note fields")

    parser.add_argument(
        "-sumOut", type=argparse.FileType("w"), help="Summary out file"
    )
    args = parser.parse_args()
    return args


def main():
    args = get_arguments()

    # Load in the CDS features from the two assemblies.
    old = SeqIO.parse(args.annotation_1, 'genbank')
    new = SeqIO.parse(args.annotation_2, 'genbank')
    old_record = next(old)
    new_record = next(new)
    old_features, new_features = [], []
    for f in old_record.features:
        if f.type == 'CDS':
            old_features.append(f)
    for f in new_record.features:
        if f.type == 'CDS':
            new_features.append(f)

    args.sumOut.write('Features in First Genbank\'s assembly:\t' + str(len(old_features)) + "\n")
    args.sumOut.write('Features in Second Genbank\'s assembly:\t' + str(len(new_features)) + "\n\n")

    # Align the features to each other.
    offsets = sorted(list(itertools.product(range(args.allowed_skipped_genes + 1),
                                            range(args.allowed_skipped_genes + 1))),
                     key=lambda x: x[0]+x[1])
    old_i, new_i = 0, 0
    exactRec = 0
    inexactRec = [0, 0 ,0]
    hypoRec = [0, 0, 0]
    newCount = 0
    oldCount = 0
    if args.addNotes:
      print("First Record CDS Product\tSimilarity\tSecond Record CDS Product\tPercent Identity\tLength Difference\tFirst Gbk's CDS Location\tSecond Gbk's CDS Location\tHypothetical Status\tFirst Record's Notes\tSecond Record's Notes\n")
    else:
      print("First Record CDS Product\tSimilarity\tSecond Record CDS Product\tPercent Identity\tLength Difference\tFirst Gbk's CDS Location\tSecond Gbk's CDS Location\tHypothetical Status\n")
    while True:
        if old_i >= len(old_features) and new_i >= len(new_features):
            break

        for old_offset, new_offset in offsets:
            try:
                old_feature = old_features[old_i + old_offset]
            except IndexError:
                old_feature = None
            try:
                new_feature = new_features[new_i + new_offset]
            except IndexError:
                new_feature = None
            try:
                match, identity, length_diff = compare_features(old_feature, new_feature, old_record, new_record, args.match_identity_threshold)
            except TypeError:
                break
            if match:
                for j in range(old_offset):
                    print_in_old_not_new(old_features[old_i + j])
                    oldCount += 1
                for j in range(new_offset):
                    print_in_new_not_old(new_features[new_i + j])
                    newCount += 1
                if identity == 1.0:
                  exactRec += 1
                res1, res2 = print_match(old_features[old_i + old_offset], new_features[new_i + new_offset], identity, length_diff, args.addNotes)
                inexactRec = addArr(inexactRec, res1)
                hypoRec = addArr(hypoRec, res2)
                old_i += old_offset
                new_i += new_offset
                break
        else:
            sys.stderr.write("Exceeded allowed number of skipped genes (" + str(args.allowed_skipped_genes) + "), unable to maintain alignment and continue comparison.\n")
            exit(2)

        if old_feature is None and new_feature is None:
            break

        old_i += 1
        new_i += 1

    args.sumOut.write('Exact Match:\t' + str(exactRec) + "\n\n")

    args.sumOut.write('Inexact Match:\t' + str(inexactRec[0] + inexactRec[1] + inexactRec[2]) + "\n")
    args.sumOut.write('  Same length:\t' + str(inexactRec[0]) + "\n")
    args.sumOut.write('  Second Gbk Seq longer:\t' + str(inexactRec[2]) + "\n")
    args.sumOut.write('  First Gbk Seq longer:\t' + str(inexactRec[1]) + "\n\n")
     
    args.sumOut.write('In Second Gbk but not in first:\t' + str(newCount) + "\n")
    args.sumOut.write('In First Gbk but not in second:\t' + str(oldCount) + "\n\n")

    args.sumOut.write('Hypothetical Annotation Change:\t' + str(hypoRec[1] + hypoRec[2]) + "\n")
    args.sumOut.write('Hypothetical:\t' + str(hypoRec[0] + hypoRec[2]) + "\n")
    

def print_match(f1, f2, identity, length_diff, outNotes):
    #print('', flush=True)
    line = f1.qualifiers['product'][0] + "\t"
    matchArr = [0, 0, 0]
    hypoArr = [0, 0, 0]
    if identity == 1.0:
#        print('Exact match')
        line += 'Exact match\t' + f2.qualifiers['product'][0] + "\t100.0\tSame Length\t"
    else:
#        print('Inexact match (' + '%.2f' % (identity * 100.0) + '% ID, ', end='')
        line += 'Inexact match\t' + f2.qualifiers['product'][0] + "\t%.2f\t" % (identity * 100.0)
        if length_diff == 0:
#            print('same length)')
            line +="Same Length\t"
            matchArr[0] += 1
        elif length_diff > 0:
#            print('old seq longer)')
            line +="First Gbk Seq Longer\t"
            matchArr[1] += 1
        elif length_diff < 0:
#            print('new seq longer)')
            line +="Second Gbk Seq Longer\t"
            matchArr[2] += 1
#    print('  old: ', end='')
#    print_feature_one_line(f1)
    line += print_feature_one_line(f1) + "\t"
#    print('  new: ', end='')
#    print_feature_one_line(f2)
    line += print_feature_one_line(f2) + "\t"
    p1 = f1.qualifiers['product'][0].lower()
    p2 = f2.qualifiers['product'][0].lower()
    if 'hypothetical' in p1 and 'hypothetical' in p2:
#        print('  still hypothetical')
        line += "Hypothetical\t"
        hypoArr[0] += 1
    elif 'hypothetical' in p1 and 'hypothetical' not in p2:
#        print('  no longer hypothetical')
        line += "No Longer Hypothetical\t"
        hypoArr[1] += 1
    elif 'hypothetical' not in p1 and 'hypothetical' in p2:
#        print('  became hypothetical')
        line += "Became Hypothetical\t"
        hypoArr[2] += 1
    else:
        line += "'Hypothetical' not in second nor first Gbk's product tag"
    
    if outNotes:
      line += "\t"
      if "note" in f1.qualifiers.keys():
        for x in f1.qualifiers["note"]:
          line += x
        line += "\t"
      else:
        line += "N/A\t"
      if "note" in f2.qualifiers.keys():
        for x in f2.qualifiers["note"]:
          line += x
      else:
        line += "N/A"

    print(line)
    return matchArr, hypoArr

def print_in_old_not_new(f): # rename file outputs
    line = f.qualifiers['product'][0] + "\tIn First Gbk but not Second\tN/A\t0.00\t" + str(f.location.end - f.location.start) + "\t" + print_feature_one_line(f) + "\tN/A\tN/A"
#    print('')
#    print('In old but not in new:')
#    print('  ', end='')
#    print_feature_one_line(f)
    print(line)


def print_in_new_not_old(f): # rename file outputs
    line = "N/A\tIn Second Gbk but not First\t" + f.qualifiers['product'][0] + "\t0.00\t" + str(f.location.end - f.location.start) + "\tN/A\t" + print_feature_one_line(f) + "\tN/A"
    #print('')
    #print('In new but not in old:')
    #print('  ', end='')
    #print_feature_one_line(f)
    print(line)


def print_feature_one_line(f):
    #f_str = f.qualifiers['product'][0]
    f_str = ""
    strand = '+' if f.location.strand == 1 else '-'
    f_str += '(' + str(f.location.start) + '-' + str(f.location.end) + ' ' + strand + ', '
    f_str += str(f.location.end - f.location.start) + ' bp)'
    return(f_str)


def compare_features(f1, f2, r1, r2, match_identity_threshold):
    if f1 is None or f2 is None:
        return False

    s1 = f1.extract(r1).seq
    s2 = f2.extract(r2).seq
    score = pairwise2.align.globalms(s1, s2, 1, 0, 0, 0, score_only=True)
    identity = score / max(len(s1), len(s2))
    match = identity >= match_identity_threshold
    length_diff = len(s1) - len(s2)
    return match, identity, length_diff


if __name__ == '__main__':
    main()