#!/usr/bin/env python
# coding: utf-8

# # Building a gene network from the IDR and STRINGdb
# This notebook exemplifies how to build Figure 1 of the paper
# using the IDR API.

# ### Dependencies
#  * [Matplotlib](http://matplotlib.org/)
#  * [NumPy](http://www.numpy.org/)
#  * [Pandas](http://pandas.pydata.org/)
#  * [NetworkX](https://networkx.github.io/)
#  * [Py2Cytoscape](https://pypi.python.org/pypi/py2cytoscape)
#  * [Requests](http://docs.python-requests.org/)
#  
# The cell below will install dependencies if you choose to run the notebook in [Google Colab](https://colab.research.google.com/notebooks/intro.ipynb#recent=true).

# In[1]:


get_ipython().run_line_magic('pip', 'install idr-py')


# In[1]:


from IPython import get_ipython
import warnings
import omero
from idr import connection

import requests
from pandas import DataFrame
from pandas import read_csv
from pandas import concat
from io import StringIO

import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from bokeh.models import ColumnDataSource
from bokeh.plotting import figure, output_notebook, show
from bokeh.models import HoverTool
from bokeh.palettes import brewer

output_notebook()
get_ipython().run_line_magic('matplotlib', 'inline')


# ### Method definitions

# In[2]:


def getBulkAnnotationAsDf(screenID, conn):
    sc = conn.getObject('Screen', screenID)
    for ann in sc.listAnnotations():
        if isinstance(ann, omero.gateway.FileAnnotationWrapper):
            if (ann.getFile().getName() == 'bulk_annotations'):
                if (ann.getFile().getSize() > 147625090):
                    print("that's a big file...")
                    return None
                ofId = ann.getFile().getId()
                break

    print(ofId)
    original_file = omero.model.OriginalFileI(ofId, False)

    table = conn.c.sf.sharedResources().openTable(original_file)
    try:
        rowCount = table.getNumberOfRows()

        column_names = [col.name for col in table.getHeaders()]

        black_list = []
        column_indices = []
        for column_name in column_names:
            if column_name in black_list:
                continue
            column_indices.append(column_names.index(column_name))

        table_data = table.slice(column_indices, None)
    finally:
        table.close()

    data = []
    for index in range(rowCount):
        row_values = [column.values[index] for column in table_data.columns]
        data.append(row_values)

    dfAnn = DataFrame(data)
    dfAnn.columns = column_names
    return dfAnn


# In[3]:


def getGenesFromPhenotype(df, phTerm):
    # all gene from bulk_annotation 'df' annotated with CMPO term 'phTerm'
    colElong = []
    for col in df.columns:
        if 'Term Accession' in col:
            if phTerm in df[col].unique():
                colElong.append(col)

    dfElong = concat([df[df[col] != ''] for col in colElong])
    return dfElong['Gene Identifier'].unique()


# ### Connect to the IDR server

# In[4]:


conn = connection('idr.openmicroscopy.org')


# ## Querying the IDR
# We will download the annotations for the three screens under study
# as panda DataFrames, and sub-select the genes from each which are
# annotated with the phenotype we are looking for, CMPO_0000077 a.k.a.
# 'elongated cell phenotype'.
# The next step is to build, from that list, a list of IDs we can
# query the STRING database with. The translation table was built offline
# using biomart and pombase.

# In[5]:


# CMPO term to look for
ph_term = 'CMPO_0000077'

# ids of screens:
# sc_id = 3  # Graml et al.
# sc_id = 206  # Rohn et al., B
# sc_id = 1202  # Fuchs et al., B

screens = [3, 206, 1202]

genes = []
for sc_id in screens:
    print('loading ' + str(sc_id))

    # loading bulk_annotations of screens as dataframes
    df = getBulkAnnotationAsDf(sc_id, conn)
    # unique genes with CPMO term
    cur = getGenesFromPhenotype(df, ph_term)
    print('got ' + str(len(cur)) + ' genes')
    genes.extend(cur)


# ### Disconnect when done loading data

# In[6]:


conn.close()


# ### Translation

# In[7]:


# The table was built offline using biomart
dfTrans = read_csv('https://raw.githubusercontent.com/IDR/idr-notebooks/master/includes/TableOfGenesWithElongatedCellPhenotype.csv')

# extract IDs
key = 'Human Ortholog Ensembl 84'
genesE84 = concat([dfTrans[dfTrans['Screen GeneID'] == x][key] for x in genes])
genesE84 = genesE84[genesE84 != '(null)']

print(genes[:10])
print(genesE84.head(10))


# ### REST query of STRINGdb
# We use the STRINGdb REST API to get all the known interactions
# of all of our genes.

# In[8]:


# Building STRINGdb REST api query
url = 'http://string-db.org/api/psi-mi-tab/interaction_partners?species=9606&limit=5&identifiers='
# genesE84 = genesE84[:10]
for g in genesE84:
    url = url + g + '%0d'
res = requests.get(url)


# In[9]:


df_original = read_csv(StringIO(res.text), sep='\t', header=None)
df_full = df_original.replace(to_replace ='string:', value = '', regex = True)
df = df_full[:int(len(df_full.index)/5)]
df_full.head(10)


# ### Some displays

# #### A. NetworkX

# In[11]:


warnings.filterwarnings('ignore')
g = nx.from_pandas_edgelist(df, 2, 3)
plt.figure(figsize=(12, 12))
nx.draw_spring(g, with_labels=True)


# #### B. Bokeh

# In[12]:


# Use spring layout from networkx for display
pts = nx.spring_layout(g)

TOOLS = "tap,pan,wheel_zoom,reset"
p = figure(plot_width=800, plot_height=800, title="My chart", tools=TOOLS)

# color certain genes of interest
# comment out this section if not needed
colors = np.ones(len(g.nodes()))
colors[[x == 'ASH2L' for x in pts.keys()]] = 4
colors[[x == 'SETD1A' for x in pts.keys()]] = 2
colors[[x == 'SETD1B' for x in pts.keys()]] = 2
cm = brewer['Set1'][5]
colors = [cm[int(x)] for x in colors]
names = list(pts.keys())
cntr = 0
for x in pts.values():
    cir = p.circle(x=x[0], y=x[1], radius=.01,
                   color=colors[cntr], name=names[cntr])
    cntr = cntr + 1

# uncomment next line if not
# coloring genes of interest
# cir=p.circle('x','y',source=sourceNode,radius=.01)
sourceEdges = ColumnDataSource(
    data=dict(
        x1=[pts[x[0]][0] for x in g.edges()],
        y1=[pts[x[0]][1] for x in g.edges()],
        x2=[pts[x[1]][0] for x in g.edges()],
        y2=[pts[x[1]][1] for x in g.edges()]
    ))
seg = p.segment('x1', 'y1', 'x2', 'y2', source=sourceEdges)
hover = HoverTool(
        tooltips=[
            ("Name", "@names"),
        ]
    )
p.add_tools(hover)

show(p)


# ### License (BSD 2-Clause)¶ 
# Copyright (C) 2016-2021 University of Dundee. All Rights Reserved.
# 
# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
# 
# Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.