# GCP Libraries
from google.cloud import bigquery
from google.colab import auth

# Data Analytics
import numpy as np
from scipy import stats

# if you're using Google Colab, authenticate to gcloud with the following
auth.authenticate_user()

# alternatively, use the gcloud SDK
#!gcloud auth application-default login

# set the google project that will be billed for this notebook's computations
google_project = 'google-project' ## change me

# Create a client to access the data within BigQuery
client = bigquery.Client(google_project)

# set parameters for query
genes = "'RHOT1','MYO7A','ZBTB10','MATK','ST18','RPS23','GCNT1','DROSHA','NUAK1','CCPG1',\
'PDGFD','KLRAP1','MTAP','RNF13','THBS1','MLX','FAP','TIMP3','PRSS1','SLC7A11',\
'OLFML3','RPS20','MCM5','POLE','STEAP4','LRRC8D','WBP1L','ENTPD5','SYNE1','DPT',\
'COPZ2','TRIO','PDPR'"

# run query and put results in data frame
pathways = client.query(("""
  SELECT
    DISTINCT pathway.*

  FROM
    `isb-cgc-bq.reactome_versioned.pathway_v77` as pathway

  INNER JOIN `isb-cgc-bq.reactome_versioned.pe_to_pathway_v77` as pe2pathway
    -- link pathways to physical entities via intermediate table
    ON pathway.stable_id = pe2pathway.pathway_stable_id

  INNER JOIN `isb-cgc-bq.reactome_versioned.physical_entity_v77` AS pe
    -- link pathways to physical entities
    ON pe2pathway.pe_stable_id = pe.stable_id

  WHERE
    -- filter by stronger evidence: "Traceable Author Statement" 
    pe2pathway.evidence_code = 'TAS'

    -- filter by pathways that are related to genes in list
    AND pe.name IN ({genes}) 

  ORDER BY pathway.name ASC
""").format(
  genes=genes
)).result().to_dataframe()

# run query and put results in data frame
lowest_pathways = client.query(("""
  SELECT
    DISTINCT pathway.*

  FROM
    `isb-cgc-bq.reactome_versioned.pathway_v77` as pathway

  INNER JOIN `isb-cgc-bq.reactome_versioned.pe_to_pathway_v77` as pe2pathway
    -- link pathways to physical entities via intermediate table
    ON pathway.stable_id = pe2pathway.pathway_stable_id

  INNER JOIN `isb-cgc-bq.reactome_versioned.physical_entity_v77` AS pe
    -- link pathways to physical entities
    ON pe2pathway.pe_stable_id = pe.stable_id

  WHERE
    -- filter by stronger evidence: "Traceable Author Statement" 
    pe2pathway.evidence_code = 'TAS'

    -- filter by pathways that are related to genes in list
    AND pe.name IN ({genes})

    -- filter to include just lowest level pathways
    AND pathway.lowest_level = TRUE

  ORDER BY pathway.name ASC
""").format(
  genes=genes
)).result().to_dataframe()

# set query parameters
lowest_level = True # only show pathways at the lowest level

gene_list_query = """
  -- Table that contains a list of all distinct genes that map to Reactome
  -- physical entities 
  SELECT
    DISTINCT pe.uniprot_id

  FROM
    `isb-cgc-bq.reactome_versioned.physical_entity_v77` AS pe

  WHERE
    -- filter by pathways that are related to genes in list
    pe.name IN ({genes})
""".format(
  genes=genes
).strip("\n")

gene_pp_query = """
  -- Table that maps pathways to the total number of interesting genes within
  -- that pathway
  SELECT
    COUNT(DISTINCT gene_list_query.uniprot_id) as num_genes,
    pathway.stable_id,
    pathway.name

  FROM
    gene_list_query

  INNER JOIN `isb-cgc-bq.reactome_versioned.physical_entity_v77` AS pe
    -- filter for interactions with genes that match a reactome
    -- physical entity
    ON gene_list_query.uniprot_id = pe.uniprot_id

  INNER JOIN `isb-cgc-bq.reactome_versioned.pe_to_pathway_v77` AS pe2pathway
    -- link physical entities to pathways via intermediate table
    ON pe.stable_id = pe2pathway.pe_stable_id

  INNER JOIN `isb-cgc-bq.reactome_versioned.pathway_v77` AS pathway
    -- link physical entities to pathways
    ON pe2pathway.pathway_stable_id = pathway.stable_id

  WHERE
    -- filter by stronger evidence: "Traceable Author Statement" 
    pe2pathway.evidence_code = 'TAS'

  GROUP BY pathway.stable_id, pathway.name
  ORDER BY num_genes DESC
""".strip("\n")

not_gene_list_query = """
  -- Table that contains a list of all genes that are NOT in the interest list
  -- This query depends on the previous "gene_list_query" sub-query.
  SELECT
    DISTINCT pe.uniprot_id

  FROM `isb-cgc-bq.reactome_versioned.physical_entity_v77` AS pe

  WHERE
    pe.uniprot_id NOT IN (
      SELECT uniprot_id FROM gene_list_query
    )
""".strip("\n")

not_gene_pp_query = """
  -- Table that maps pathways to the number of proteins that are NOT drug
  -- targets in that pathway.
  SELECT
    COUNT(DISTINCT not_gene_list_query.uniprot_id) AS num_not_genes,
    pathway.stable_id,
    pathway.name

  FROM not_gene_list_query

  INNER JOIN `isb-cgc-bq.reactome_versioned.physical_entity_v77` AS pe
    ON not_gene_list_query.uniprot_id = pe.uniprot_id

  INNER JOIN `isb-cgc-bq.reactome_versioned.pe_to_pathway_v77` AS pe2pathway
    ON pe.stable_id = pe2pathway.pe_stable_id
  
  INNER JOIN `isb-cgc-bq.reactome_versioned.pathway_v77` AS pathway
    ON pe2pathway.pathway_stable_id = pathway.stable_id

  WHERE
    -- filter by stronger evidence: "Traceable Author Statement" 
    pe2pathway.evidence_code = 'TAS'

  GROUP BY pathway.stable_id, pathway.name
  ORDER BY num_not_genes DESC
""".strip("\n")

gene_count_query = """
  -- Table that contains the counts of # of genes that are/are not targets
  SELECT
    gene_count,
    not_gene_count,
    gene_count + not_gene_count AS total_count

  FROM 
    (SELECT COUNT(*) AS gene_count FROM gene_list_query),
    (SELECT COUNT(*) AS not_gene_count FROM not_gene_list_query)
""".strip("\n")

observed_query = """
  -- Table with observed values per pathway in the contingency matrix
  SELECT
    gene_pp_query.num_genes AS in_gene_in_pathway,
    not_gene_pp_query.num_not_genes AS not_gene_in_pathway,
    gene_count_query.gene_count - gene_pp_query.num_genes AS in_gene_not_pathway,
    gene_count_query.not_gene_count - not_gene_pp_query.num_not_genes AS not_gene_not_pathway,
    gene_pp_query.stable_id,
    gene_pp_query.name

  FROM 
    gene_pp_query,
    gene_count_query

  INNER JOIN not_gene_pp_query
    ON gene_pp_query.stable_id = not_gene_pp_query.stable_id
""".strip("\n")

sum_query = """
  -- Table with summed observed values per pathway in the contingency matrix
  SELECT
    observed_query.in_gene_in_pathway + observed_query.not_gene_in_pathway AS pathway_total,
    observed_query.in_gene_not_pathway + observed_query.not_gene_not_pathway AS not_pathway_total,
    observed_query.in_gene_in_pathway + observed_query.in_gene_not_pathway AS gene_total,
    observed_query.not_gene_in_pathway + observed_query.not_gene_not_pathway AS not_gene_total,
    observed_query.stable_id,
    observed_query.name

  FROM
    observed_query
""".strip("\n")

expected_query = """  
  -- Table with the expected values per pathway in the contingency matrix
  SELECT 
    sum_query.gene_total * sum_query.pathway_total / gene_count_query.total_count AS exp_in_gene_in_pathway,
    sum_query.not_gene_total * sum_query.pathway_total / gene_count_query.total_count AS exp_not_gene_in_pathway,
    sum_query.gene_total * sum_query.not_pathway_total / gene_count_query.total_count AS exp_in_gene_not_pathway,
    sum_query.not_gene_total * sum_query.not_pathway_total / gene_count_query.total_count AS exp_not_gene_not_pathway,
    sum_query.stable_id,
    sum_query.name

  FROM 
    sum_query, gene_count_query
""".strip("\n")

chi_squared_query = """
  -- Table with the chi-squared statistic for each pathway
  SELECT
    -- Chi squared statistic with Yates' correction
    POW(ABS(observed_query.in_gene_in_pathway - expected_query.exp_in_gene_in_pathway) - 0.5, 2) / expected_query.exp_in_gene_in_pathway 
    + POW(ABS(observed_query.not_gene_in_pathway - expected_query.exp_not_gene_in_pathway) - 0.5, 2) / expected_query.exp_not_gene_in_pathway
    + POW(ABS(observed_query.in_gene_not_pathway - expected_query.exp_in_gene_not_pathway) - 0.5, 2) / expected_query.exp_in_gene_not_pathway
    + POW(ABS(observed_query.not_gene_not_pathway - expected_query.exp_not_gene_not_pathway) - 0.5, 2) / expected_query.exp_not_gene_not_pathway
    AS chi_squared_stat,
    observed_query.stable_id,
    observed_query.name

  FROM observed_query

  INNER JOIN expected_query
    ON observed_query.stable_id = expected_query.stable_id
""".strip("\n")

lowest_level_filter = """
  INNER JOIN `isb-cgc-bq.reactome_versioned.pathway_v77` AS pathway
    ON chi_squared_query.stable_id = pathway.stable_id

  WHERE pathway.lowest_level = TRUE
""".strip("\n")

final_query = """
  WITH
    gene_list_query AS (
      {gene_list_query}
    ),

    gene_pp_query AS (
      {gene_pp_query}
    ),

    not_gene_list_query AS (
      {not_gene_list_query}
    ),

    not_gene_pp_query AS (
      {not_gene_pp_query}
    ),

    gene_count_query AS (
      {gene_count_query}
    ),

    observed_query AS (
      {observed_query}
    ),

    sum_query AS (
      {sum_query}
    ),

    expected_query AS (
      {expected_query}
    ),
    
    chi_squared_query AS (
      {chi_squared_query}
    )

  SELECT
    observed_query.in_gene_in_pathway,
    observed_query.in_gene_not_pathway,
    observed_query.not_gene_in_pathway,
    observed_query.not_gene_not_pathway,
    chi_squared_query.chi_squared_stat,
    chi_squared_query.stable_id,
    chi_squared_query.name

  FROM chi_squared_query

  INNER JOIN observed_query
    ON chi_squared_query.stable_id = observed_query.stable_id
  {lowest_level_filter}
  ORDER BY chi_squared_stat DESC
""".format(
  # make final query a little easier to read by removing/adding some white space 
  gene_list_query="\n    ".join(gene_list_query.strip().splitlines()),
  gene_pp_query="\n    ".join(gene_pp_query.strip().splitlines()),
  not_gene_list_query="\n    ".join(not_gene_list_query.strip().splitlines()),
  not_gene_pp_query="\n    ".join(not_gene_pp_query.strip().splitlines()),
  gene_count_query="\n    ".join(gene_count_query.strip().splitlines()),
  observed_query="\n    ".join(observed_query.strip().splitlines()),
  sum_query="\n    ".join(sum_query.strip().splitlines()),
  expected_query="\n    ".join(expected_query.strip().splitlines()),
  chi_squared_query="\n    ".join(chi_squared_query.strip().splitlines()),
  lowest_level_filter=(
    "\n  "+"\n".join(lowest_level_filter.strip().splitlines())+"\n" if lowest_level else ""
  )
).strip("\n")

# run query and put results in data frame
chi_squared_pathways = client.query(final_query).result().to_dataframe()

# run query and put results in data frame
num_pathways_result = client.query(("""
  SELECT
    COUNT (*) AS num_pathways
  FROM
    `isb-cgc-bq.reactome_versioned.pathway_v77` as pathway
  {lowest_level_filter}
""").format(
  lowest_level_filter=("WHERE lowest_level = TRUE" if lowest_level else "")
)).result().to_dataframe()

# extract observed values from bigquery result
observed = chi_squared_pathways[[
  'in_gene_in_pathway',
  'in_gene_not_pathway',
  'not_gene_in_pathway',
  'not_gene_not_pathway'
]]

# calculate the chi-squared statistic using the scipy stats package  
chi2_stat = []
chi2_pvalue = []
for index, row in observed.iterrows():
  stat, pvalue, _, _ = stats.chi2_contingency(
    np.reshape(np.matrix(row), (2,2)), correction=True
  )
  chi2_stat.append(stat)
  chi2_pvalue.append(pvalue)

# add columns to the original data frame
chi_squared_pathways['scipy_stat'] = chi2_stat
chi_squared_pathways['scipy_p_value'] = chi2_pvalue