#!/usr/bin/env python3

import fire

# Workflow 2, Module 1E: Gene interactions
from BioLink.biolink_client import BioLinkWrapper
from pprint import pprint

import pandas as pd

from translator_modules.core import Config
from translator_modules.core import Payload


class GeneInteractions(object):

    def __init__(self):
        self.blw = BioLinkWrapper(Config().get_biolink_api_endpoint())
        self.meta = {
            'input_type': {
                'complexity': 'set',
                'id_type': 'HGNC',
                'data_type': 'gene',
            },
            'output_type': {
                'complexity': 'set',
                'id_type': 'HGNC',
                'data_type': 'gene',
            },

            'source': 'Monarch Biolink',
            'predicate': ['blm:interacts with']
        }

    def metadata(self):
        print("""Mod1E Interaction Network metadata:""")
        pprint(self.meta)

    @staticmethod
    # RMB: July 5, 2019 - gene_records is a Pandas DataFrame
    def load_gene_set(gene_records):
        annotated_gene_set = []
        for gene in gene_records.to_dict(orient='records'):
            annotated_gene_set.append({
                'input_id': gene['hit_id'],
                'sim_input_curie': gene['hit_id'],
                'input_symbol': gene['hit_symbol']
            })
        return annotated_gene_set

    def get_interactions(self, input_gene_set, threshold):

        annotated_input_gene_set = self.load_gene_set(input_gene_set)
        lower_bound = int(threshold)

        results = []
        for gene in annotated_input_gene_set:
            interactions = self.blw.gene_interactions(gene_curie=gene['sim_input_curie'])
            for assoc in interactions['associations']:
                interaction = \
                    self.blw.parse_association(
                        input_id=gene['sim_input_curie'],
                        input_label=gene['input_symbol'],
                        association=assoc
                    )
                results.append({
                    'input_id': interaction['input_id'],
                    'input_symbol': interaction['input_symbol'],
                    'hit_symbol': interaction['hit_symbol'],
                    'hit_id': interaction['hit_id'],
                    'score': 0,
                })

        # Process the results
        results = pd.DataFrame().from_records(results)
        counts = results['hit_symbol'].value_counts().rename_axis('unique_values').to_frame('counts').reset_index()
        high_counts = counts[counts['counts'] > lower_bound]['unique_values'].tolist()
        results = pd.DataFrame(results[results['hit_symbol'].isin(high_counts)])

        return results


class GeneInteractionSet(Payload):

    def __init__(self, threshold=12, input_gene_set_file=None):

        super(GeneInteractionSet, self).__init__(GeneInteractions())

        input_gene_set = None
        if input_gene_set_file:
            with open(input_gene_set_file) as stream:
                # assuming it's JSON and it's a record list
                input_gene_set = pd.read_json(stream, orient='records')

        self.results = self.mod.get_interactions(input_gene_set, threshold)


if __name__ == '__main__':
    fire.Fire(GeneInteractionSet)
