#!/usr/bin/env Rscript

# Copyright Saez Lab 2020
#
# Authors: Denes Turei and Alberto Valdeolivas
# Contact: turei.denes@gmail.com
#
# License: MIT License https://opensource.org/licenses/MIT
#
# Combines the OmniPath PPI and transcriptional regulatory network
# with the human-SARS CoV-2 interactions from Gordon et al. 2020

require(dplyr)
require(tibble)
require(purrr)
require(openxlsx)
require(igraph)
require(qgraph)
require(tidyr)

if(!'OmnipathR' %in% installed.packages()[,"Package"]){
    require(devtools)
    install_github('saezlab/OmnipathR')
}

require(OmnipathR)

# URLs for supplementary tables S2 and S3 from Gordon et al. 2020
# https://www.biorxiv.org/content/10.1101/2020.03.22.002386v2
url_s2 <-
    paste0(
        'https://www.biorxiv.org/content/biorxiv/early/',
        '2020/03/23/2020.03.22.002386/DC4/embed/media-4.xlsx?download=true'
    )
url_s3 <-
    paste0(
        'https://www.biorxiv.org/content/biorxiv/early/',
        '2020/03/23/2020.03.22.002386/DC5/embed/media-5.xlsx?download=true'
    )

# the OmniPath PPI interaction network
ia_omnipath <- import_Omnipath_Interactions() %>% as_tibble()

# in the example we don't use these but if you need higher coverage on PPI
# add these with `bind_rows` below
ia_ligrec <- import_LigrecExtra_Interactions() %>% as_tibble()
ia_pwextra <- import_PathwayExtra_Interactions() %>% as_tibble()
ia_kinaseextra <- import_KinaseExtra_Interactions() %>% as_tibble()

# transcriptional regulation
# if you want more interactions, pass the argument
# `confidence_level = c('A', 'B', 'C', 'D')`
ia_transcriptional <- import_TFregulons_Interactions() %>% as_tibble()

# post-transcriptional regulation
# here we don't use it but if you are interested
# add these by `bind_rows` below
ia_post_transcriptional <- import_miRNAtarget_Interactions() %>% as_tibble()

# downloading the host-pathogen interactions from the paper
human_sarscov2_raw <- read.xlsx(url_s2, startRow = 2) %>% as_tibble()
# downloading the compound-host-pathogen interactions from the paper
# here we don't use it but if you are interested combine with the other
# interactions a similar way as the host-pathogen ones
human_sarscov2_compounds_raw <- read.xlsx(url_s3) %>% as_tibble()

# combining all the interactions to one data frame
ia_human_sarscov2 <-
    human_sarscov2_raw %>%
    select(
        source = Bait,
        target = Preys,
        source_genesymbol = Bait,
        target_genesymbol = PreyGene,
        MIST,
        Saint_BFDR,
        AvgSpec,
        FoldChange
    ) %>%
    # manually adding the known S--ACE2 and S--BSG interactions
    add_row(
        source = 'SARS-CoV2 Spike',
        target = 'Q9BYF1',
        source_genesymbol = 'Spike',
        target_genesymbol = 'ACE2'
    ) %>%
    add_row(
        source = 'SARS-CoV2 Spike',
        target = 'P35613',
        source_genesymbol = 'Spike',
        target_genesymbol = 'BSG'
    )

interactions <-
    as_tibble(
        bind_rows(
            ia_omnipath %>% mutate(type = 'ppi'),
            ia_pwextra %>% mutate(type = 'ppi'),
            ia_kinaseextra %>% mutate(type = 'ppi'),
            ia_ligrec %>% mutate(type = 'ppi'),
            ia_transcriptional %>% mutate(type = 'transcriptional'),
            ia_human_sarscov2 %>% mutate(
                source_genesymbol = sub('SARS-CoV2 ', '', source_genesymbol),
                type = 'host_pathogen',
                # the direction or effect of these interactions is unknown
                is_directed = 0,
                is_stimulation = 0,
                is_inhibition = 0,
                consensus_direction = 0,
                consensus_stimulation = 0,
                consensus_inhibition = 0
            )
        )
    ) %>%
    group_by(
        source_genesymbol,
        target_genesymbol,
        type,
        is_directed,
        is_inhibition,
        is_stimulation
    ) %>%
    summarize_all(first) %>%
    ungroup() %>%
    mutate(
        curation_effort = ifelse(
            type == 'host_pathogen',
            # for our purpose these interactions are of high importance
            # hence we assign a high "weight"
            1000,
            ifelse(
                is.na(curation_effort),
                # no curation effort means no literature evidence available
                # here we assign a small weight < 1
                .5,
                curation_effort
            )
        )
    ) %>%
    # remove any ambiguity between gene symbols and uniprot ids
    group_by(source) %>%
    mutate(source_genesymbol = first(source_genesymbol)) %>%
    ungroup() %>%
    group_by(target) %>%
    mutate(target_genesymbol = first(target_genesymbol)) %>%
    ungroup()

# Obtaining pneumocyte protein expression data from the Human Protein Atlas
# using the annotations database of OmniPath
# From the annotations database you can access a broad variety of data about
# proteins: function, localization, structure, expression, etc, overall
# from more than 40 resources. See a table of resources, properties and
# values at http://omnipathdb.org/annotations_summary or use the
# `get_annotation_databases` method from the OmnipathR package.
human_proteins <-
    c(
        interactions %>% filter(type != 'host_pathogen') %>% pull(source),
        interactions %>% pull(target)
    ) %>%
    unique()

pneumocyte_expression <-
    import_Omnipath_annotations(
        filter_databases = c('HPA_tissue'),
        select_genes = human_proteins
    ) %>%
    spread(label, value) %>%
    as_tibble() %>%
    filter(organ == 'lung', tissue == 'pneumocytes')

# creating an igraph network from the interactions data frame
net_human_sarscov2 <-
    interaction_graph(interactions = interactions) %>%
    igraph::simplify(remove.multiple = FALSE, remove.loops = TRUE)

# labeling the proteins by organism
sarscov2_proteins <- ia_human_sarscov2 %>% pull(source) %>% unique()

V(net_human_sarscov2)$organism <-
    ifelse(
        V(net_human_sarscov2)$up_ids %in% sarscov2_proteins,
        'SARS-CoV2',
        'human'
    )

# labeling direct viral targets and transcription factors
V(net_human_sarscov2)$viral_target <-
    V(net_human_sarscov2)$up_ids %in% (
        ia_human_sarscov2 %>% pull(target) %>% unique()
    )

V(net_human_sarscov2)$tf <-
    V(net_human_sarscov2)$up_ids %in% (
        ia_transcriptional %>% pull(source) %>% unique()
    )

# node indices of the SARS-CoV2 proteins
nodes_sarscov2 <- (V(net_human_sarscov2)$organism == 'SARS-CoV2') %>% which()

# adding pneumocyte expression levels as a vertex attribute
expression_by_protein <-setNames(
    pneumocyte_expression %>% pull(level),
    pneumocyte_expression %>% pull(uniprot)
)

V(net_human_sarscov2)$expression <-
    V(net_human_sarscov2)$up_ids %>%
    map(
        function(uniprot){
            `if`(
                uniprot %in% names(expression_by_protein),
                expression_by_protein[[uniprot]],
                'Not available'
            )
        }
    ) %>%
    unlist()

# creating subgraphs of the neighborhood of each virus protein
neighborhoods_human_sarscov2_1 <-
    make_ego_graph(
        net_human_sarscov2,
        order = 1,
        nodes = nodes_sarscov2,
        mode = 'all'
    )

nodes_in_neighborhoods <-
    neighborhoods_human_sarscov2_1 %>%
    map(function(g){V(g)$name}) %>%
    unlist() %>%
    unique()

# reducing the network to the neighborhoods
net_human_sarscov2_1 <-
    net_human_sarscov2 %>%
    delete_vertices(
        which(
            !V(net_human_sarscov2)$name %in% nodes_in_neighborhoods
        )
    )

# reducing the network to the viral targets having PPI among
# themselves
net_human_sarscov2_1 <-
    net_human_sarscov2_1 %>%
    delete_vertices(
        which(!(
            V(net_human_sarscov2_1)$viral_target |
            V(net_human_sarscov2_1)$organism == 'SARS-CoV2'
        ))
    )

nodes_having_ppi <-
    net_human_sarscov2_1 %>%
    get.edges(
        E(net_human_sarscov2_1)[
            E(net_human_sarscov2_1)$type %in% c('ppi', 'transcriptional')
        ]
    ) %>%
    c() %>%
    unique()

sarscov2_nodes <-
    (V(net_human_sarscov2_1)$organism == 'SARS-CoV2') %>%
    which()

nodes_to_keep <- c(nodes_having_ppi, sarscov2_nodes) %>% unique()

net_human_sarscov2_1 <-
    net_human_sarscov2_1 %>%
    induced_subgraph(nodes_to_keep)

net_human_sarscov2_1 <-
    net_human_sarscov2_1 %>%
    delete_vertices(degree(.) == 0)

# creating a figure of the direct neighborhood network
fr_layout <-
    net_human_sarscov2_1 %>%
    get.edgelist(names = FALSE) %>%
    qgraph.layout.fruchtermanreingold(
        vcount = vcount(net_human_sarscov2_1),
        area = vcount(net_human_sarscov2_1) ** 2.3,
        repulse.rad = vcount(net_human_sarscov2_1) ** 2.1,
        niter = 500
    )

png(
    filename = 'SARS-CoV2_OmniPath_neighborhood.png',
    width = 7,
    height = 7,
    units = 'in',
    res = 600
)

# igraph is not able to plot arrows with and without head at the same time
# plotting first only the undirected edges, then the directed
plot(
    net_human_sarscov2_1,
    layout = fr_layout,
    vertex.size = 0,
    vertex.label = NA,
    edge.arrow.size = 0,
    edge.width = .2,
    edge.color = ifelse(
        E(net_human_sarscov2_1)$is_directed,
        NA,
        '#646567AA'
    )
)

plot(
    net_human_sarscov2_1,
    layout = fr_layout,
    vertex.size = 5,
    vertex.label.cex = .66,
    vertex.label.dist = 1,
    vertex.label.degree = -pi * .75,
    vertex.color = ifelse(
        V(net_human_sarscov2_1)$organism == 'SARS-CoV2',
        '#FCCC06',
        ifelse(
            V(net_human_sarscov2_1)$viral_target,
            '#97BE73',
            '#B6B7B9'
        )
    ),
    vertex.label.family = 'DINPro',
    vertex.label.color = '#454447',
    vertex.frame.width = 0,
    vertex.frame.color = NA,
    edge.color = ifelse(
        E(net_human_sarscov2_1)$type == 'host_pathogen',
        '#FCCC06',
        ifelse(
            E(net_human_sarscov2_1)$is_inhibition,
            '#E25C49AA',
            ifelse(
                E(net_human_sarscov2_1)$is_stimulation,
                '#49969AAA',
                ifelse(
                    E(this_community)$is_directed,
                    '#646567AA',
                    NA
                )
            )
        )
    ),
    edge.arrow.size = .2,
    edge.width = .2,
    add = TRUE
)

dev.off()

# creating the same plot with showing the pneumocyte expression levels
png(
    filename = 'SARS-CoV2_OmniPath_neighborhood_pneumocytes.png',
    width = 7,
    height = 7,
    units = 'in',
    res = 600
)

plot(
    net_human_sarscov2_1,
    layout = fr_layout,
    vertex.size = 0,
    vertex.label = NA,
    edge.arrow.size = 0,
    edge.width = .2,
    edge.color = ifelse(
        E(net_human_sarscov2_1)$is_directed,
        NA,
        '#646567AA'
    )
)

plot(
    net_human_sarscov2_1,
    layout = fr_layout,
    vertex.size = 5,
    vertex.label.cex = .66,
    vertex.label.dist = 1,
    vertex.label.degree = -pi * .75,
    vertex.color = ifelse(
        V(net_human_sarscov2_1)$organism == 'SARS-CoV2',
        '#FCCC06',
        ifelse(
            V(net_human_sarscov2_1)$expression == 'High',
            '#ff1414',
            ifelse(
                V(net_human_sarscov2_1)$expression == 'Medium',
                '#ff7676',
                ifelse(
                    V(net_human_sarscov2_1)$expression == 'Low',
                    '#ffc4c4',
                    'whitesmoke'
                )
            )
        )
    ),
    vertex.label.family = 'DINPro',
    vertex.label.color = '#454447',
    vertex.frame.width = 0,
    vertex.frame.color = NA,
    edge.color = ifelse(
        E(net_human_sarscov2_1)$type == 'host_pathogen',
        '#FCCC06',
        ifelse(
            E(net_human_sarscov2_1)$is_inhibition,
            '#E25C49AA',
            ifelse(
                E(net_human_sarscov2_1)$is_stimulation,
                '#49969AAA',
                ifelse(
                    E(this_community)$is_directed,
                    '#646567AA',
                    NA
                )
            )
        )
    ),
    edge.arrow.size = .2,
    edge.width = .2,
    add = TRUE
)

dev.off()


# creating a function for recursive clustering, i.e. to call the methods
# repeatedly until also large clusters are chopped into pieces of the
# desired maximum size
cluster_recursive <- function(
        graph,
        membership = NULL,
        max_size = 100,
        method = cluster_louvain,
        verbose = FALSE,
        ...
    ){
    
    if(is.null(membership)){
        
        membership <- rep('1', vcount(graph))
        
    }
    
    cl_sizes <- table(membership)
    
    for(cl_idx in names(cl_sizes)[which(cl_sizes > max_size)]){
        
        this_cluster_idx <- which(membership == cl_idx)
        
        if(verbose){
            message(
                sprintf(
                    'Cluster %s has a size of %d, clustering further.',
                    cl_idx,
                    length(this_cluster_idx)
                )
            )
        }
        
        this_cluster_graph <-
            graph %>%
            induced_subgraph(this_cluster_idx)
        
        subclusters <- method(this_cluster_graph, ...)
        
        membership[this_cluster_idx] <-
            subclusters$membership %>%
            map(function(m){sprintf('%s.%d', cl_idx, m)}) %>%
            unlist()
        
    }
    
    if(max(table(membership)) <= max_size){
        
        return(membership)
        
    }else{
        
        return(
            cluster_recursive(
                graph,
                membership,
                max_size = max_size,
                method = method,
                ...
            )
        )
        
    }
    
}


# using curation effort as weight for a better clustering
E(net_human_sarscov2)$weight <- E(net_human_sarscov2)$curation_effort

# some clustering methods work only with undirected and simple graphs
net_human_sarscov2_undir <-
    net_human_sarscov2 %>%
    as.undirected() %>%
    simplify()

# clustering the network by the louvain algorithm
human_sarscov2_louvain <-
    net_human_sarscov2_undir %>%
    cluster_recursive(method = cluster_louvain, max_size = 150)

names(human_sarscov2_louvain) <- V(net_human_sarscov2_undir)$name

# creating a vertex attribute from cluster memberships
V(net_human_sarscov2)$cluster <-
    V(net_human_sarscov2)$name %>%
    map(
        function(name){
            `if`(
                name %in% names(human_sarscov2_louvain),
                human_sarscov2_louvain[[name]],
                NA
            )
        }
    ) %>%
    unlist()

sarscov2_proteins <- interactions %>%
    filter(type == 'host_pathogen') %>%
    pull(source) %>%
    unique()

human_target_proteins <- interactions %>%
    filter(type == 'host_pathogen') %>%
    pull(target) %>%
    unique()

# shortcut for the vertex sequence of the large network
vs <- V(net_human_sarscov2)

# plotting each cluster containing viral targets
dir.create('clusters', showWarnings = FALSE)
list.files('clusters') %>%
walk(function(f){file.remove(file.path('clusters', f))})

cairo_pdf(
    'SARS-CoV2_OmniPath_clusters_pneumocytes.pdf',
    width = 7,
    height = 7,
    onefile = TRUE
)

for(cl_id in unique(vs$cluster)){
    
    has_target_proteins <-
        vs$up_ids[which(vs$cluster == cl_id)] %>%
        intersect(human_target_proteins) %>%
        length() %>%
        `>`(0)
    
    if(has_target_proteins){
        
        this_community <-
            net_human_sarscov2 %>%
            induced_subgraph(
                which(
                    # the members of the cluster and all virus proteins
                    vs$cluster == cl_id |
                    vs$up_ids %in% sarscov2_proteins
                )
            )
        
        # removing nodes with no connections
        this_community <-
            this_community %>%
            delete_vertices(which(degree(this_community) == 0))
        
        if(vcount(this_community) == 0) next
        
        fr_layout <-
            this_community %>%
            get.edgelist(names = FALSE) %>%
            qgraph.layout.fruchtermanreingold(
                vcount = vcount(this_community),
                area = vcount(this_community) ** 2.3,
                repulse.rad = vcount(this_community) ** 2.1,
                niter = 500
            )
        
        pngpath <- file.path(
            'clusters',
            sprintf(
                'SARS-CoV2_OmniPath_cluster-%s_pneumocytes.png',
                cl_id
            )
        )
        
        message(
            sprintf(
                'Community %s with %d members and %d connections',
                cl_id,
                vcount(this_community),
                ecount(this_community)
            )
        )
        
        # in case you want png
        # png(
        #     filename = pngpath,
        #     width = 7,
        #     height = 7,
        #     units = 'in',
        #     res = 600
        # )
        
        plot(
            this_community,
            layout = fr_layout,
            vertex.size = 0,
            vertex.label = NA,
            edge.arrow.size = 0,
            edge.width = .2,
            edge.color = ifelse(
                E(this_community)$is_directed,
                NA,
                '#646567AA'
            )
        )

        plot(
            this_community,
            layout = fr_layout,
            vertex.size = 5,
            vertex.label.cex = .66,
            vertex.label.dist = 1,
            vertex.label.degree = -pi * .75,
            vertex.color = ifelse(
                V(this_community)$organism == 'SARS-CoV2',
                '#FCCC06',
                ifelse(
                    V(this_community)$expression == 'High',
                    '#ff1414',
                    ifelse(
                        V(this_community)$expression == 'Medium',
                        '#ff7676',
                        ifelse(
                            V(this_community)$expression == 'Low',
                            '#ffc4c4',
                            'whitesmoke'
                        )
                    )
                )
            ),
            vertex.label.family = 'DINPro',
            vertex.label.color = '#454447',
            vertex.frame.width = 0,
            vertex.frame.color = NA,
            edge.color = ifelse(
                E(this_community)$type == 'host_pathogen',
                '#FCCC06',
                ifelse(
                    E(this_community)$is_inhibition,
                    '#E25C49AA',
                    ifelse(
                        E(this_community)$is_stimulation,
                        '#49969AAA',
                        ifelse(
                            E(this_community)$is_directed,
                            '#646567AA',
                            NA
                        )
                    )
                )
            ),
            edge.arrow.size = .2,
            edge.width = .2,
            add = TRUE
        )

        #dev.off()
        
    }
    
}

dev.off()