import argparse
import json
import re
from collections import defaultdict

import BioExt
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from BioExt.uds import _align_par


class HyPhySummary(object):

    def __init__(self, arguments, summary_json=None, annotation_json=None):
        self.arguments = arguments
        self.ref_map = ''
        self.summary_json = {}
        self.annotation_json = {}
        self.include_in_annotation = {}
        self.test_map = {}
        self.site_reports = {}
        self.labels = {}
        self.ref_seq_map = []
        self.cfel = {}
        self.relax = {}
        self.busted = {}
        self.slac = {}
        self.fel = {}
        self.meme = {}
        self.meme_full = {}
        self.prime = {}
        self.fade = {}
        self.bgm = {}
        self.ref_genes = [
            ['genome', ' ']
        ]
        # 0-based gene coordinates
        # from, to, name, within gene codon offset (to deal with ORF1a/b overlap)
        self.gene_coordinates = [[265, 13482, 'ORF1a', 0],
                                 [13467, 21554, 'ORF1b', -1],
                                 [21562, 25383, 'S', 0],
                                 [25392, 26219, 'ORF3a', 0],
                                 [26244, 26471, 'E', 0],
                                 [26522, 27190, 'M', 0],
                                 [27201, 27386, 'ORF6', 0],
                                 [27393, 27758, 'ORF7a', 0],
                                 [27893, 28258, 'ORF8', 0],
                                 [28273, 29532, 'N', 0],
                                 [29557, 29673, 'ORF10', 0]]
        self.score_matrix_ = BioExt.scorematrices.DNA95.load()
        if summary_json is not None:
            self.summary_json = summary_json
        if annotation_json is not None:
            self.annotation_json = annotation_json
        try:
            if arguments.summary_input:
                with open(arguments.summary_input) as fh:
                    self.summary_json = json.load(fh)
        except Exception:
            pass
        try:
            if arguments.annotation_input:
                with open(arguments.annotation_input) as fh:
                    self.annotation_json = json.load(fh)
        except Exception:
            pass
        self.gene = self.arguments.gene
        self._get_incoming_labels()
        self._get_map()

    def summary(self):
        try:
            if self.summary_json is not None:
                if self.gene not in self.summary_json:
                    self.summary_json[self.gene] = {}
            if self.cfel is not None:
                self.cfel_summary()
            if self.relax is not None:
                self.relax_summary()
            if self.busted is not None:
                self.busted_summary()
            if self.slac is not None:
                self.slac_summary()
            if self.fel is not None:
                self.fel_summary()
            if self.meme is not None:
                self.meme_summary()
            if self.meme_full is not None:
                self.meme_full_summary()
            if self.prime is not None:
                self.prime_summary()
            if self.fade is not None:
                self.fade_summary()
            if self.bgm is not None:
                self.bgm_summary()
        except Exception:
            if self.gene:
                del self.summary_json[self.gene]
                if self.gene in self.annotation_json:
                    del self.annotation_json[self.gene]
            raise
        return self.annotation_json, self.summary_json

    def cfel_summary(self):
        self.cfel = self._load_json(self.arguments.cfel)
        if self.cfel is None:
            return
        node_tags = {}
        _ = self._newick_parser(self.cfel['input']['trees']['0'], False, node_tags, self.cfel)['json']
        if self.summary_json is not None:
            omegas = {}
            T = {}
            for k in [[k.split('*')[1], v[0][0]] for k, v in self.cfel['fits']['Global MG94xREV']['Rate Distributions'].items()]:
                if k[0] != 'background':
                    self.test_map[k[0]] = 'Test'
                else:
                    self.test_map[k[0]] = 'Reference'
                omegas[k[0]] = k[1]
                T[k[0]] = 0.
            for branch, nt in (self.cfel['tested']['0']).items():
                if branch not in node_tags:
                    node_tags[branch] = []
                info = self.cfel['branch attributes']['0'][branch]
                if nt != '':
                    T[nt] += info['Global MG94xREV']
                node_tags[branch].append(info['Global MG94xREV'])
            self.summary_json[self.gene]['rates'] = {'mean-omega': omegas, 'T': T}
        beta_indices = []
        p_indices = []
        subs = []
        for i, tag in enumerate(self.cfel['MLE']['headers']):
            if tag[0].find('beta') == 0:
                beta_indices.append([i, re.split(r'\(|\)', tag[0])[1]])
            elif tag[0].find('P-value') == 0:
                p_indices.append([i, re.split(r'\(|\)|for ', tag[0])[1]])
            elif tag[0].find('subs') == 0:
                subs.append([i, re.split(r'\(|\)', tag[0])[1]])
        for i, row in enumerate(self.cfel['MLE']['content']['0']):
            if self.annotation_json is not None and len(self.ref_map):  # if this is specified, write everything out
                gs = self._get_genomic_annotation(i)
                if gs[0] >= 0:
                    self.include_in_annotation[i] = gs[0]
                    self.annotation_json[gs[0]] = {
                        'G': gs[1],
                        'S': gs[2],
                        'index': i,
                        'bCFEL': {
                            'p': row[4],
                            'a': row[0],
                            'b': self._make_report_dict(row, beta_indices),
                            'pi': self._make_report_dict(row, p_indices),
                            'pp': row[-2],
                            's': self._make_report_dict(row, subs),
                            'q': row[-3]
                        }
                    }
            if row[-4] <= self.arguments.pvalue:
                self.site_reports[i] = {'cfel': row}

    def relax_summary(self):
        self.relax = self._load_json(self.arguments.relax)
        if self.relax is None:
            return
        if self.summary_json is not None:
            relax_d = {}
        for r, rr in self.summary_json[self.gene]['rates']['mean-omega'].items():
            relax_d[r] = []
            for ignored, rd in self.relax['fits']['RELAX alternative']['Rate Distributions'][self.test_map[r]].items():
                relax_d[r].append(rd)
        self.summary_json[self.gene]['rates']['relax'] = relax_d
        self.summary_json[self.gene]['relax'] = {
            'p': self.relax['test results']['p-value'],
            'K': self.relax['test results']['relaxation or intensification parameter']
        }

    def busted_summary(self):
        self.busted = self._load_json(self.arguments.busted)
        if self.busted is None:
            return
        if self.summary_json is not None:
            self.summary_json[self.gene]['rates']['busted'] = self.busted['fits']['Unconstrained model']['Rate Distributions']
            self.summary_json[self.gene]['busted'] = {'p': self.busted['test results']['p-value'], }

    def slac_summary(self):
        self.slac = self._load_json(self.arguments.slac)
        if self.slac is None:
            return

        def def_value():
            return defaultdict(int)
        compressed_subs = {}
        node_tags = {}
        the_tree = self._newick_parser(self.slac['input']['trees']['0'], False, node_tags, self.slac)['json']
        root_node = None
        if self.summary_json is not None:
            for branch, info in self.slac['branch attributes']['0'].items():
                if branch in node_tags:
                    node_tags[branch].append(info['Global MG94xREV'])
                else:
                    root_node = branch
            self.summary_json[self.gene]['tree'] = self.slac['input']['trees']['0']
            self.summary_json[self.gene]['tree_tags'] = node_tags
        if len(self.include_in_annotation):
            for i in self.include_in_annotation:
                report = self.annotation_json[self.include_in_annotation[i]]
                counts_codon_site = {}
                counts_aa_site = {}
                gs = self._get_genomic_annotation(i)
                if gs[0] >= 0:
                    self.labels[root_node] = self.slac['branch attributes']['0'][root_node]['codon'][0][i]
                    try:
                        self._traverse_tree_in_order(the_tree, self.slac['branch attributes']['0'], i, None, root_node)
                    except Exception:
                        raise
                    compressed_subs[gs[0]] = self.labels
                for k in set([k[0] for k in node_tags.values()]):
                    if len(k):
                        counts_codon_site[k] = defaultdict(int)
                        counts_aa_site[k] = defaultdict(int)
                for branch, tag in node_tags.items():
                    if len(tag[0]) > 0 and tag[1] is False:
                        codon = self.slac['branch attributes']['0'][branch]['codon'][0][i]
                        aa = self.slac['branch attributes']['0'][branch]['amino-acid'][0][i]
                        counts_codon_site[tag[0]][codon] += 1
                        counts_aa_site[tag[0]][aa] += 1
                report['cdn'] = counts_codon_site
                report['aa'] = counts_aa_site
            self.summary_json[self.gene]['subs'] = compressed_subs

    def fel_summary(self):
        self.fel = self._load_json(self.arguments.fel)
        if None in [self.fel, self.cfel]:
            return
        for i, row in enumerate(self.fel['MLE']['content']['0']):
            if i in self.include_in_annotation:
                self.annotation_json[self.include_in_annotation[i]]['bFEL'] = {'a': row[0], 'b': row[1], 'p': row[4]}
            if i in self.site_reports or row[4] <= self.arguments.pvalue and row[1] > row[0]:
                if i in self.site_reports:
                    self.site_reports[i]['fel'] = row
                else:
                    self.site_reports[i] = {'fel': row,
                                            'cfel': self.cfel['MLE']['content']['0'][i]}

    def meme_summary(self):
        self.meme = self._load_json(self.arguments.meme)
        if None in [self.fel, self.cfel, self.meme]:
            return
        for i, row in enumerate(self.meme['MLE']['content']['0']):
            if i in self.include_in_annotation:
                self.annotation_json[self.include_in_annotation[i]]['bMEME'] = {
                    'p': row[6],
                    'a': row[0],
                    'b+': row[3],
                    'w+': row[4],
                    'b-': row[1],
                    'w-': row[2],
                    'br': row[7]
                }
            if i in self.site_reports or row[6] <= self.arguments.pvalue:
                if i in self.site_reports:
                    self.site_reports[i]['meme'] = row
                else:
                    self.site_reports[i] = {'meme': row,
                                            'fel': self.fel['MLE']['content']['0'][i],
                                            'cfel': self.cfel['MLE']['content']['0'][i]}
        # annotate branches with EBF support
        for n, info in self.meme['branch attributes']['0'].items():

            if n in self.summary_json[self.gene]['tree_tags']:
                sig_sites = []
                for tag, ebf in info.items():
                    bits = tag.split(' ')
                    if len(bits) >= 4 and ebf >= 100:
                        sig_sites.append(self.include_in_annotation[int(bits[2]) - 1])
                self.summary_json[self.gene]['tree_tags'][n].append(sig_sites)

    def meme_full_summary(self):
        self.meme_full = self._load_json(self.arguments.meme_full)
        if None in [self.fel, self.cfel, self.meme, self.meme_full]:
            return
        for i, row in enumerate(self.meme_full['MLE']['content']['0']):
            if i in self.include_in_annotation:
                self.annotation_json[self.include_in_annotation[i]]['lMEME'] = {
                    'p': row[6],
                    'a': row[0],
                    'b+': row[3],
                    'w+': row[4],
                    'b-': row[1],
                    'w-': row[2],
                    'br': row[7]
                }
            if i in self.site_reports or row[6] <= self.arguments.pvalue:
                if i in self.site_reports:
                    self.site_reports[i]['full-meme'] = row
                else:
                    self.site_reports[i] = {'full-meme': row,
                                            'meme': self.meme['MLE']['content']['0'][i],
                                            'fel': self.fel['MLE']['content']['0'][i],
                                            'cfel': self.cfel['MLE']['content']['0'][i]}
        # annotate branches with EBF support
        for n, info in self.meme_full['branch attributes']['0'].items():
            if n in self.summary_json[self.gene]['tree_tags']:
                sig_sites = []
                for tag, ebf in info.items():
                    bits = tag.split(' ')
                    if len(bits) >= 4 and ebf >= 100:
                        sig_sites.append(self.include_in_annotation[int(bits[2]) - 1])
                self.summary_json[self.gene]['tree_tags'][n].append(sig_sites)

    def prime_summary(self):
        self.prime = self._load_json(self.arguments.prime)
        if self.prime is None:
            return
        if self.summary_json is not None:
            h = self.prime['MLE']['headers']
            self.summary_json[self.gene]['prime-properties'] = [h[k][1].replace('Importance for ', '') for k in range(6, len(h), 3)]
        if len(self.include_in_annotation):
            for i in self.include_in_annotation:
                report = self.annotation_json[self.include_in_annotation[i]]
                prime_info = self.prime['MLE']['content']['0'][i]
                if prime_info:
                    report['prime'] = {
                        'p': [prime_info[k] for k in ([5, ] + list(range(7, len(prime_info), 3)))],
                        'lambda': [prime_info[k] for k in range(6, len(prime_info), 3)]
                    }
                else:
                    report['prime'] = None  # invariable

    def fade_summary(self):
        self.fade = self._load_json(self.arguments.fade)
        if self.fade is None:
            return
        if len(self.include_in_annotation):
            for i in self.include_in_annotation:
                report = self.annotation_json[self.include_in_annotation[i]]
                report['fade'] = {}
                for residue, info in self.fade['MLE']['content'].items():
                    # Skip anything past the end of this section.
                    if i >= len(info['0']):
                        continue
                    if len(residue) == 1:
                        report['fade'][residue] = {'rate': info['0'][i][1], 'BF': info['0'][i][-1]}

    def bgm_summary(self):
        self.bgm = self._load_json(self.arguments.bgm)
        if self.bgm is None:
            return
        if self.summary_json is not None:
            try:
                self.summary_json[self.gene]['bgm'] = self.bgm['MLE']['content']
            except KeyError:
                self.summary_json[self.gene]['bgm'] = []

    def _load_json(self, filename):
        if filename is None:
            return None
        try:
            with open(filename, 'r') as fh:
                return json.load(fh)
        except Exception:
            raise

    def _get_map(self):
        for seq_record in SeqIO.parse(self.arguments.combined, 'fasta'):
            seq_id = seq_record.description
            ref_seq_re = re.compile(self.arguments.name)
            if ref_seq_re.search(seq_id):
                ref_seq = str(seq_record.seq).upper()
                self.aligned_str = None

                def output_record(x):
                    listified_input = list(x)
                    if len(listified_input) == 1:
                        self.aligned_str = listified_input[0]

                def ignore_record(x):
                    pass

                for s in self.ref_genes:
                    _align_par(SeqRecord(Seq(s[1]), id=s[0]), [SeqRecord(Seq(ref_seq), id='ref')],
                               self.score_matrix_, False, False, 0.8, ignore_record, output_record)
                    if (self.aligned_str is not None):
                        break

                self.ref_map = self.aligned_str.seq.strip('-')
                c = 0
                i = 0
                map_to_genome = []
                while i < len(self.ref_map):
                    if self.ref_map[i:i + 3] != '---':
                        map_to_genome.append(i)
                    i += 3
                i = 0
                c = 0
                while i < len(ref_seq):
                    if ref_seq[i:i + 3] != '---':
                        self.ref_seq_map.append(map_to_genome[c // 3] + self.aligned_str.annotations['position'])
                        c += 3
                    else:
                        self.ref_seq_map.append(-1)
                    i += 3
        if self.summary_json is not None:
            if self.gene not in self.summary_json:
                self.summary_json[self.gene] = dict()
            self.summary_json[self.gene]['map'] = self.ref_seq_map

    def _make_report_dict(self, row, indices):
        result = {}
        for i, t in indices:
            result[t] = row[i]
        return result

    def _get_genomic_annotation(self, site):
        genomic_site_coord = -1
        gene_name = ''
        gene_site = -1
        if len(self.ref_seq_map):
            genomic_site_coord = self.ref_seq_map[site]
            if genomic_site_coord < 0:
                gene_site = 'Not in SC2 (deletion)'
            else:
                gene_name = None
                for k in self.gene_coordinates:
                    if k[0] <= genomic_site_coord and k[1] > genomic_site_coord:
                        genomic_site = ((genomic_site_coord + k[3]) - k[0]) // 3
                        gene_name = k[2]
                        gene_site = genomic_site + 1
                        break
                if gene_name is None:
                    gene_name = 'Not mapped'
        else:
            gene_name = 'N/A'
        return (genomic_site_coord, gene_name, gene_site)

    def _traverse_tree_in_order(self, node, slac_data, i, parent_tag, root):
        node_tag = None
        if node is None:
            return
        try:
            nn = root if node['name'] == 'root' else node['name']
        except Exception:
            raise
        if nn in slac_data:
            node_tag = slac_data[nn]['codon'][0][i]
            if (parent_tag != node_tag):
                self.labels[nn] = node_tag
                self.labels[node['name']] = node_tag
        if 'children' in node:
            for c in node['children']:
                if c is not None:
                    if 'name' in c:
                        self._traverse_tree_in_order(c, slac_data, i, node_tag, root)

    def _match_node_names(self, qry_node, ref_node, mapping):
        if 'children' in qry_node and 'children' in ref_node:
            mapping[ref_node['name']] = qry_node['name']
            if len(qry_node['children']) != len(ref_node['children']):
                raise Exception('Internal topology mismatch')
            for i, n in enumerate(ref_node['children']):
                self._match_node_names(qry_node['children'][i], n, mapping)
        elif 'children' in qry_node:
            raise Exception('Topology mismatch')
        elif 'children' in ref_node:
            raise Exception('Topology mismatch')
        else:
            if qry_node['name'] != ref_node['name']:
                raise Exception('Leaf name mismatch')

    def _get_incoming_labels(self):
        json_data = self._load_json(self.arguments.labels)
        self.incoming_labels = json_data

    def _newick_parser(self, nwk_str, bootstrap_values, track_tags, json_map):
        clade_stack = []
        automaton_state = 0
        current_node_name = ''
        current_node_attribute = ''
        current_node_annotation = ''
        quote_delimiter = None
        name_quotes = {"'": 1, '"': 1}

        def add_new_tree_level():
            new_level = {'name': None}
            the_parent = clade_stack[len(clade_stack) - 1]
            if ('children' not in the_parent):
                the_parent['children'] = []
            clade_stack.append(new_level)
            the_parent['children'].append(clade_stack[len(clade_stack) - 1])
            clade_stack[len(clade_stack) - 1]['original_child_order'] = len(the_parent['children'])

        def finish_node_definition():
            nonlocal current_node_name
            nonlocal current_node_annotation
            nonlocal current_node_attribute
            this_node = clade_stack.pop()
            if (bootstrap_values and 'children' in this_node):
                this_node['bootstrap_values'] = current_node_name
            else:
                this_node['name'] = current_node_name
            this_node['attribute'] = current_node_attribute
            this_node['annotation'] = current_node_annotation
            try:
                if 'children' not in this_node:
                    node_tag = self.arguments.default_tag
                    if json_map:
                        tn = json_map['branch attributes']['0'][this_node['name']]
                    else:
                        tn = this_node
                    nn = tn['original name'] if 'original name' in tn else tn['name']
                    for k, v in self.incoming_labels.items():
                        if nn.find(k) >= 0:
                            node_tag = v
                            break
                else:
                    counts = {}
                    node_tag = ''
                    for n in this_node['children']:
                        counts[n['tag']] = 1 + (counts[n['tag']] if n['tag'] in counts else 0)
                    if len(counts) == 1:
                        node_tag = list(counts.keys())[0]
                    this_node['tag'] = node_tag
            except Exception:
                raise
            if track_tags is not None:
                track_tags[this_node['name']] = [this_node['tag'], 'children' in this_node]
            current_node_name = ''
            current_node_attribute = ''
            current_node_annotation = ''

        def generate_error(location):
            unexpected = nwk_str[location]
            before = nwk_str[location - 20:location + 1]
            after = nwk_str[location + 1:location + 20]
            return {
                'json': None,
                'error': 'Unexpected %s in %s [ERROR HERE] %s' % (unexpected, before, after)
            }
        tree_json = {'name': 'root'}
        clade_stack.append(tree_json)
        space = re.compile(r'\s')
        for char_index in range(len(nwk_str)):
            try:
                current_char = nwk_str[char_index]
                if automaton_state == 0:
                    # look for the first opening parenthesis
                    if (current_char == '('):
                        add_new_tree_level()
                        automaton_state = 1
                elif automaton_state == 1 or automaton_state == 3:
                    # case 1: // name
                    # case 3: { // branch length
                    # reading name
                    if (current_char == ':'):
                        automaton_state = 3
                    elif current_char == ',' or current_char == ')':
                        try:
                            finish_node_definition()
                            automaton_state = 1
                            if (current_char == ','):
                                add_new_tree_level()
                        except Exception:
                            return generate_error(char_index)
                    elif (current_char == '('):
                        if len(current_node_name) > 0:
                            return generate_error(char_index)
                        else:
                            add_new_tree_level()
                    elif (current_char in name_quotes):
                        if automaton_state == 1 and len(current_node_name) == 0 and len(current_node_attribute) == 0 and len(current_node_annotation) == 0:
                            automaton_state = 2
                            quote_delimiter = current_char
                            continue
                        return generate_error(char_index)
                    else:
                        if (current_char == '['):
                            if len(current_node_annotation):
                                return generate_error(char_index)
                            else:
                                automaton_state = 4
                        else:
                            if (automaton_state == 3):
                                current_node_attribute += current_char
                            else:
                                if (space.search(current_char)):
                                    continue
                                if (current_char == ';'):
                                    char_index = len(nwk_str)
                                    break
                                current_node_name += current_char
                elif automaton_state == 2:
                    # inside a quoted expression
                    if (current_char == quote_delimiter):
                        if (char_index < len(nwk_str - 1)):
                            if (nwk_str[char_index + 1] == quote_delimiter):
                                char_index += 1
                                current_node_name += quote_delimiter
                                continue
                            quote_delimiter = 0
                            automaton_state = 1
                            continue
                    else:
                        current_node_name += current_char
                elif automaton_state == 4:
                    # inside a comment / attribute
                    if (current_char == ']'):
                        automaton_state = 3
                    else:
                        if (current_char == '['):
                            return generate_error(char_index)
                        current_node_annotation += current_char
            except Exception:
                return generate_error(char_index)

        if (len(clade_stack) != 1):
            return generate_error(len(nwk_str) - 1)

        if (len(current_node_name)):
            tree_json['name'] = current_node_name

        return {
            'json': tree_json,
            'error': None
        }


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Summarize selection analysis results.')
    parser.add_argument('--combined', help='Combined reference and query alignment from TN-93', required=False, type=str)
    parser.add_argument('--pvalue', help='p-value to use', required=False, type=float, default=0.05)
    parser.add_argument('--gene', help='Name of the gene or sequence being analyzed', required=False, type=str)
    parser.add_argument('--labels', help='JSON file with labels', required=False, type=str)
    parser.add_argument('--annotation-output', help='Write a JSON file with site annotations', required=True, type=str)
    parser.add_argument('--summary-output', help='Write a JSON file here segment annotations', required=True, type=str)
    parser.add_argument('--annotation-inputs', help='Comma-separated list of site annotation files to merge', required=False, type=str)
    parser.add_argument('--summary-inputs', help='Comma-separated list of segment annotation files to merge', required=False, type=str)
    parser.add_argument('--default-tag', help='Default name for sequences that have no explicit label', required=False, type=str, default='Reference')
    parser.add_argument('--name', help='The sequence ID to highlight', required=False, default='MN908947')
    parser.add_argument('--mode', help='Operation mode, generate a summary or merge multiple summaries', type=str, choices=['summary', 'merge'], default='summary')
    parser.add_argument('--relax', help='Path to RELAX.json file', required=False, type=str)
    parser.add_argument('--busted', help='Path to BUSTED.json file', required=False, type=str)
    parser.add_argument('--slac', help='Path to SLAC.json file', required=False, type=str)
    parser.add_argument('--fel', help='Path to FEL.json file', required=False, type=str)
    parser.add_argument('--cfel', help='Path to CFEL.json file', required=False, type=str)
    parser.add_argument('--meme', help='Path to MEME.json file', required=False, type=str)
    parser.add_argument('--meme-full', help='Path to MEME-full.json file', dest='meme_full', required=False, type=str)
    parser.add_argument('--prime', help='Path to PRIME.json file', required=False, type=str)
    parser.add_argument('--fade', help='Path to FADE.json file', required=False, type=str)
    parser.add_argument('--bgm', help='Path to BGM.json file', required=False, type=str)
    arguments = parser.parse_args()
    if arguments.mode == 'summary':
        analyzer = HyPhySummary(arguments)
        annotation_json, summary_json = analyzer.summary()
        if annotation_json is not None:
            with open(arguments.annotation_output, 'w') as fh:
                json.dump(annotation_json, fh, indent=1)
        if summary_json is not None:
            with open(arguments.summary_output, 'w') as fh:
                json.dump(summary_json, fh, indent=1)
    else:
        summary = {}
        annotation = {}
        for filename in arguments.annotation_inputs.split(','):
            with open(filename, 'r') as fh:
                annotation.update(json.load(fh))
        for filename in arguments.summary_inputs.split(','):
            with open(filename, 'r') as fh:
                summary.update(json.load(fh))
        with open(arguments.annotation_output, 'w') as fh:
            json.dump(annotation, fh, indent=1)
        with open(arguments.summary_output, 'w') as fh:
            json.dump(summary, fh, indent=1)
    exit(0)
