from indra.preassembler import grounding_mapper as gm
import sys
import logging
from indra.statements import Complex
import pickle
from indra.preassembler import Preassembler
from indra.preassembler.hierarchy_manager import hierarchies
from indra.databases import biogrid_client as bg
import csv

logger = logging.getLogger('analyze_complexes')

# Filter out complexes from the statement list
def load_file(stmts_file):
    logger.info("Loading results...")
    with open(stmts_file) as f:
        results = pickle.load(f)
    return results


def get_biogrid_interactors(gene_list):
    pass

if __name__ == '__main__':

    # Load the statements
    if len(sys.argv) < 2:
        print "Usage: %s reach_stmts_file" % sys.argv[0]
        sys.exit()
    results = load_file(sys.argv[1])

    all_stmts = [stmt for paper_stmts in results.values()
                      for stmt in paper_stmts]

    # Map grounding
    logger.info('Mapping grounding...')
    gmap = gm.GroundingMapper(gm.default_grounding_map)
    map_stmts = gmap.map_agents(all_stmts)
    map_stmts = gmap.rename_agents(map_stmts)

    # Combine duplicates
    logger.info('Removing duplicates...')
    pa = Preassembler(hierarchies, map_stmts)
    pa.combine_duplicates()

    # Get complexes
    complexes = [s for s in pa.unique_stmts if isinstance(s, Complex)]
    protein_complexes = [s for s in complexes
                           if all([True if 'HGNC' in ag.db_refs.keys()
                                        else False
                                        for ag in s.agent_list()])]

    genes = list(set([ag.db_refs['HGNC'] for stmt in protein_complexes
                                         for ag in stmt.members]))

    # Get complexes from BioGrid and combine duplicates
    num_genes_per_query = 50
    start_indices = range(0, len(genes), num_genes_per_query)
    end_indices = [i + num_genes_per_query
                   if i + num_genes_per_query < len(genes) else len(genes)
                   for i in start_indices]
    bg_complexes = []
    for i in range(len(start_indices)):
        logger.info("Querying biogrid for %s" %
                    str(genes[start_indices[i]:end_indices[i]]))
        bg_complexes += (bg.get_statements(
                                genes[start_indices[i]:end_indices[i]]))

    # Filter out Biogrid statements not involving genes in the gene list
    # (this will make duplicate removal more efficient
    bg_filt = []
    for stmt in bg_complexes:
        if stmt.members[0].name in genes and \
           stmt.members[1].name in genes:
            bg_filt.append(stmt)
    # Might as well free up some memory
    del bg_complexes

    logger.info("Combining duplicates with biogrid...")
    pa = Preassembler(hierarchies, bg_filt + protein_complexes)
    pa.combine_duplicates()

    reach_only = []
    bg_only = []
    reach_and_bg = []
    for stmt in pa.unique_stmts:
        evidence_source_list = set([])
        for e in stmt.evidence:
            evidence_source_list.add(e.source_api)
        if 'reach' in evidence_source_list and \
           'biogrid' in evidence_source_list:
            reach_and_bg.append(stmt)
        elif 'reach' in evidence_source_list and \
             'biogrid' not in evidence_source_list:
            reach_only.append(stmt)
        elif 'reach' not in evidence_source_list and \
             'biogrid' in evidence_source_list:
            bg_only.append(stmt)

    rows = []
    for stmt in reach_only:
        rows.append([stmt.members[0].name, stmt.members[1].name,
                     len(stmt.evidence)])
    with open('unmatched_complexes.tsv', 'w') as f:
        csvwriter = csv.writer(f, delimiter='\t')
        csvwriter.writerows(rows)


@@ -288,16 +288,72 @@
 port csv
+%0Afrom indra.databases import hgnc_client, uniprot_client
 %0A%0Alogger
@@ -597,58 +597,8 @@
 s%0A%0A%0A
-def get_biogrid_interactors(gene_list):%0A    pass%0A%0A
 if _
@@ -1332,16 +1332,41 @@
 mplex)%5D%0A
+    # Get HGNC grounding%0A
     prot
@@ -1598,20 +1598,75 @@
 )%5D%0A%0A    
-gene
+logger.info('Mapping gene IDs to gene symbols')%0A    gene_id
 s = list
@@ -1786,16 +1786,79 @@
 mbers%5D))
+%0A    genes = %5Bhgnc_client.get_hgnc_name(id) for id in gene_ids%5D
 %0A%0A    # 

Update analyze_complexes to work with HGNC IDs