%%bash

wget -q https://www.ebi.ac.uk/gwas/api/search/downloads/full -O tmp.tsv
cat tmp.tsv | \
    awk 'BEGIN {FS="\t";OFS="\t"} {chrom=$12; gsub(chrom,"chr"chrom,$12)}{print $0}' | \
    sed s/,//g > gwas.tsv
rm tmp.tsv

myBucket = "gs://fc-cad72548-2d6b-41ce-82aa-975cb7e8b764"

"""
Explanation: Finding the most representative GWAS associated with cell-specific enhancers
(Execution on Google Cloud File System)
In this tutorial we are going to use a GWAS dataset (accessible from this link) together with the whole ENCODE BroadPeak dataset to find which mutations (and their associated traits) are most represented in enhancer regions which are present in a limited set of cells.
As first thing let's download the data.
End of explanation
"""


!gsutil cp ./gwas.tsv $myBucket/

"""
Explanation: In order to run the query on HDFS, we have to put the file there. We will use the bucket for this Terra Notebook.
End of explanation
"""


import gmql as gl
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

"""
Explanation: Library imports
End of explanation
"""


gl.set_master("yarn")

"""
Explanation: Setting the master of the cluster
In this example, the data reside in the HDFS of the spark cluster. Let's say that the cluster is managed by the YARN resource manager. We have therefore to tell PyGMQL to use it.
End of explanation
"""


gmql_repository = "gs://geco_repository/"
gwas_path = myBucket + "/gwas.tsv"

"""
Explanation: Loading of the GWAS dataset
In this example, we have loaded the GMQL repository on the Google Cloud Storage. It is convenient to store in a variable the path of the repository.
End of explanation
"""


gwas = gl.load_from_file(gwas_path, 
                         parser=gl.parsers.RegionParser(
                             chrPos=11, 
                             startPos=12, 
                             stopPos=12, 
                             otherPos=[(7, "trait", 'string')]))

"""
Explanation: The GWAS data comes from a single TSV file. Therefore we can import it using the load_from_file function.
Notice that we have to specify a parser to properly load our data. Therefore it is wise to take a look at the schema of the downloaded file.
We are mainly interested in the mutation position (11-th and 12-th column) and the associated trait (7-th).
End of explanation
"""


gwas.head().regs

"""
Explanation: Inspecting the dataset
We can load a tiny part of the dataset to make sense of the data types and schema. You can inspect the dataset using the head function. This function returns a GDataframe object, which enables the access to regions (regs) and metadata (meta)
End of explanation
"""


gwas.schema

"""
Explanation: We can also simply look at the schema
End of explanation
"""


gwas_data = gwas.materialize().regs

"""
Explanation: Plotting the traits
We want to get an idea of the trait distribution. In order to do that we have to load the data in memory. Thereofre we can call the materialize function and take the regions.
End of explanation
"""


plt.figure(figsize=(20,5))
sns.countplot(data=gwas_data[gwas_data.trait.isin(
        gwas_data.trait.value_counts().iloc[:30].index)], x='trait')
plt.xticks(rotation=90)
plt.title("Top represented GWAS traits", fontsize=20)
plt.show()

"""
Explanation: We now plot the number of regions for each of the top 30 represented traits.
End of explanation
"""


broad = gl.load_from_path(gmql_repository + "HG19_ENCODE_BROAD/")

broad.schema

"""
Explanation: Loading of the ENCODE BroadPeak dataset
We now load the ENCODE BroadPeak dataset.
If the data come already in the GDM format, they can be loaded using the load_from_path function. A GDM dataset is stored as a folder having the following structure:
/path/to/dataset/:
    - sample1.gdm
    - sample1.gdm.meta
    - sample2.gdm
    - sample2.gdm.meta
    - ...
    - schema.xml
The first dataset we load is the one from the GWAS study.
End of explanation
"""


acetyl = broad[broad['experiment_target'] == 'H3K27ac-human']

"""
Explanation: Getting the enhancers
We identify enhancers thanks to the presence of H3K27ac peaks. We therefore select all acetylation peaks from ENCODE thanks to the experiment_target metadata attribute.
End of explanation
"""


peaked = acetyl.reg_project(new_field_dict={
        'peak': acetyl.right/2 + acetyl.left/2})

"""
Explanation: We get the peak region of the Chip-Seq using the reg_project function. The peak position (peak) is given by the center of the region.
$$
    peak = \frac{right + left}{2}
$$
End of explanation
"""


enlarge = peaked.reg_project(new_field_dict={
    'left': peaked.peak - 1500, 
    'right': peaked.peak + 1500})

"""
Explanation: Once we have the peak, we extend the search region to $\pm 1500 bp$. We use again reg_project
End of explanation
"""


enhancers_by_cell_line = enlarge.normal_cover(1, "ANY", 
                                              groupBy=['biosample_term_name'])

"""
Explanation: Grouping by cell line and aggregating the signals
We are interested in enhancers which are cell specific. Therefore it is important to group our data by cell line. In addition to this we merge the signals coming from different tracks for the same cell line. We can do both of these actions using the normal_cover function.
End of explanation
"""


max_overlapping = 2
cell_specific_enhancers = enhancers_by_cell_line.normal_cover(1, max_overlapping)

cell_specific_enhancers.schema

cell_specific_enhancers_by_cell_line = enhancers_by_cell_line.join(
                                    cell_specific_enhancers, 
                                    [gl.DLE(0)], 'left', 
                                    refName="en", expName="csen")

"""
Explanation: To select only the cell-specific enhancers we can now apply again normal_cover and constraining the maximum number of overlaps between the regions to be a selected threshold.
In this case we select a threshold of 2.
End of explanation
"""


gwas.schema

enhancer_gwas = cell_specific_enhancers_by_cell_line.map(
                        gwas, refName="csen", expName="gwas", 
                        new_reg_fields={'traits': gl.BAG('trait')})
enhancer_gwas = enhancer_gwas.reg_project(
                    ["count_csen_gwas", "traits"],
                    new_field_dict={'cell_line': enhancer_gwas['csen.en.biosample_term_name','string']})

"""
Explanation: Mapping mutations to cell specific enhancers
We now map the mutations in the GWAS dataset on the enhancer regions. We store the list of traits associated to each enhancer using the gl.BAG expression.
End of explanation
"""


enhancer_gwas = enhancer_gwas.materialize()

"""
Explanation: Materializing the result
We now can call the materialize function to execute the full query. The result will be collected in a GDataframe object.
End of explanation
"""


enhancer_gwas.regs['traits'] = enhancer_gwas.regs.traits\
                                .map(lambda x: x.split(",") if pd.notnull(x) else x)

"""
Explanation: The traits column of the resulting region is the list of traits associated with the cell specific enhancer. The data comes in the form of a string of trait names.
We convert the string to a list.
End of explanation
"""


cell_trait = pd.DataFrame.from_records([(k, v) for k, vs in enhancer_gwas.regs[enhancer_gwas.regs.count_csen_gwas > 0]\
                                        .groupby("cell_line").traits.sum().to_dict().items() for v in vs], 
                                       columns=['cell_line', 'trait'])

cross = pd.crosstab(cell_trait.cell_line, cell_trait.trait)

"""
Explanation: Analysis
The final part of the analysis regards the matching of cell lines and traits. We want to understand if a cell line (which is represented by its specific enhancers) has some particular mutation trait associated.
The analysis is performed in Pandas using the result region attributes traits and cell_line.
We build an association matrix between cell lines and traits by firstly converting the result to a list of (cell_line, trait), converting it to a Pandas DataFrame, and finally using the crosstab Pandas function to extract the matrix.
End of explanation
"""


plt.figure(figsize=(50, 15))
sns.heatmap(cross[cross.sum(0).sort_values(ascending=False).iloc[:100].index], cmap='Reds', vmax=70, linewidths=1, annot=True, cbar=False)
plt.xticks(fontsize=20)
plt.yticks(fontsize=20)

plt.xlabel("Trait", fontsize=30)
plt.ylabel("Cell line", fontsize=30)
plt.show()

"""
Explanation: We finally plot the result as an heatmap.
End of explanation
"""


