import argparse
import json
import sys
from math import log as ln
from scipy.special import gammaln, betaln

def gain_ln_probability(graph):
    total_genes = 0
    total = 0
    for module in graph:
        total_genes += len(module['gains'])
        total += gammaln(len(module['gains']) + 1)

    return total + gammaln(len(graph)) - gammaln(len(graph) + total_genes)


def loss_ln_probability(graph):
    active_genes = []
    total = 0
    for module in graph:
        inherited_genes = set()
        for dependency in module['dependencies']:
            inherited_genes.update(active_genes[dependency])

        lost_gene_count = len(module['losses'])
        present_gene_count = len(inherited_genes) - lost_gene_count

        total += betaln(lost_gene_count + 1, present_gene_count + 1)

        module_active_genes = set(inherited_genes)
        module_active_genes.update(module['gains'])
        module_active_genes.difference_update(module['losses'])

        active_genes.append(module_active_genes)

    return total

def graph_prior_ln_probability(graph, is_dependency_graph):
    if not is_dependency_graph:
        return 0
    else:
        species_count = 0
        dependencies_count = 0
        for module in graph:
            if module['species'] is not None:
                species_count += 1

            dependencies_count += len(module['dependencies'])

        return betaln(dependencies_count + 1, (len(graph) * (len(graph) - 1)) / 2 - dependencies_count  + 1) - ln(pow(2, species_count) - 1)


def evaluate(graph, is_dependency_graph):
    for index, module in enumerate(graph):
        for dependency in module['dependencies']:
            assert dependency < index
    return (gain_ln_probability(graph) + loss_ln_probability(graph) + graph_prior_ln_probability(graph, is_dependency_graph)) / ln(2)

def main():
    parser = argparse.ArgumentParser(description="Dependency Graph Evaluator")
    parser.add_argument("type", type=str, choices=['tree', 'dependency'], help='type of graph')
    parser.add_argument("filename", type=argparse.FileType('r'), nargs='+', help='json file containing a dependency graph')

    args = parser.parse_args()

    for input_file in args.filename:
        try:
            graph = json.load(input_file)
        except json.decoder.JSONDecodeError as error:
            print(error, file=sys.stderr)
            sys.exit(1)

        score = evaluate(graph, args.type == 'dependency')

        print(input_file.name, score, sep='\t')

if __name__ == '__main__':
    main()