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

#
#  Copyright (c) 2017 - EMBL-EBI
#
#  File author(s): Dénes Türei (denes@ebi.ac.uk)
#
#  Website: http://www.ebi.ac.uk/~denes
#

import pypath
import re

int1 = \
'''MEK1_pS217_S221->MAPK_pT202_Y204
GSK3A_pS21_S9->GSK3A_pS9
GSK3B_pS21_S9->GSK3B_pS9
RPS6KB1_pT389->RPS6_pS235_S236
EGFR_pY1068->Src_pY416
HER2_pY1248->EGFR_pY1068
Src_pY416->EGFR_pY1068
RPS6_pS240_S244->RPS6_pS235_S236
AKT_pS473->GSK3A_pS21_S9
AKT_pS473->GSK3B_pS21_S9
MEK1_pS217_S221->BAD_pS112
EGFR_pY1068->HER2_pY1248
NRG1->HER3_pY1298
HGF->MET_pY1235
AKT_pS473->AKT_pT308
AKT_pT308->AKT_pS473
MAPK_pT202_Y204->MEK1_pS217_S221
GSK3A_pS9->GSK3A_pS21_S9
GSK3B_pS9->GSK3B_pS21_S9
RPS6_pS235_S236->RPS6KB1_pT389
GSK3A_pS21_S9->PRAS40_pT246
GSK3B_pS21_S9->PRAS40_pT246
RPS6KB1_pT389->BAD_pS112
RPS6KB1_pT389->RPS6_pS240_S244'''

int2 = \
'''Serum->MAPK14_pT180_Y182
HGF->RAF1_pS338
EGF->Src_pY416
EGF->EGFR_pY1068
Rictor_pT1135->mTOR_pS2448
EGF->EGFR_pY1173
EGF->EIF4EBP1_pS65
EGF->GSK3A_pS21_S9
EGF->GSK3B_pS21_S9
PBS1->RPS6_pS240_S244
BAD_pS112->JUN_pS73
GSK3A_pS9->RPS6KB1_pT389
GSK3B_pS9->RPS6KB1_pT389
Src_pY416->EGFR_pY1068
RPS6KB1_pT389->NDRG1_pT346
NDRG1_pT346->GSK3A_pS9
NDRG1_pT346->GSK3B_pS9
AKT_pS473->AKT_pT308
AKT_pT308->AKT_pS473
GSK3A_pS9->GSK3A_pS21_S9
GSK3A_pS9->GSK3A_pS21_S9
NRG1->HER3_pY1298
HGF->MET_pY1235
RPS6KA1_pT359_S363->BAD_pS112'''

def map_int(interactions):
    
    return \
        list(
            map(
                lambda i:
                    list(
                        map(
                            lambda p:
                                p.split('_'),
                            i
                        )
                    ),
                map(
                    lambda i:
                        i.split('->'),
                    interactions.split('\n')
                )
            )
        )

def map_mods(mods):
    
    return \
        list(
            map(
                lambda m:
                    (m[0], int(m[1:])),
                map(
                    lambda m:
                        m.replace('p', ''),
                    filter(
                        lambda m:
                            m != 'alpha' and m != 'beta',
                        mods
                    )
                )
            )
        )

def to_uniprot(interactions):
    
    result = []
    
    unmapped = set([])
    
    m = pypath.mapping.Mapper()
    
    for i in interactions:
        
        u0 = m.map_name(i[0][0], 'genesymbol', 'uniprot')
        u1 = m.map_name(i[1][0], 'genesymbol', 'uniprot')
        
        if not len(u0):
            
            unmapped.add(i[0][0])
        
        if not len(u1):
            
            unmapped.add(i[1][0])
        
        if len(u0) and len(u1):
            
            mods0 = map_mods(i[0][1:])
            mods1 = map_mods(i[1][1:])
            
            result.append([[u0[0], i[0][0]] + mods0,
                           [u1[0], i[1][0]] + mods1])
    
    return result, unmapped

def lookup_interactions(uint, pa):
    
    result = []
    
    for i in uint:
        
        e = pa.edge_exists(i[0][0], i[1][0])
        
        res = [i[0][0], i[0][1], i[1][0], i[1][1]]
        
        if type(e) is not int:
            
            res.extend(['', '', ''])
            result.append(res)
            print('Could not find edge between %s and %s' % (i[0][1], i[1][1]))
            continue
        
        e = pa.graph.es[e]
        
        dirr = (i[0][0], i[1][0])
        d = e['dirs']
        ptms = set(i[1][2:])
        
        
        
        if d.get_dir(dirr):
            
            res.append('directed')
            if dirr in e['refs_by_dir']:
                refs = list(map(lambda r: r.pmid, e['refs_by_dir'][dirr]))
                res.append(';'.join(refs))
            
            else:
                refs = list(map(lambda r: r.pmid, e['refs_by_dir']['undirected']))
                res.append(';'.join(refs))
        
        elif d.get_dir('undirected'):
            
            undir_refs = list(map(lambda r: r.pmid, e['refs_by_dir']['undirected']))
            res.append('undirected')
            res.append(';'.join(undir_refs))
        
        ptms_in_dbs = set([])
        
        for ptm in e['ptm']:
            
            if ptm.ptm.protein == i[1][0]:
                
                this_ptm = (ptm.ptm.residue.name, ptm.ptm.residue.number)
                
                if this_ptm in ptms:
                    
                    ptms_in_dbs.add(this_ptm)
        
        res.append(';'.join(list(map(lambda p: '%s%u' % p, ptms_in_dbs))))
        
        result.append(res)
    
    return result

def table(fname, ll):
    
    with open(fname, 'w') as fp:
        fp.write('\n'.join(map(lambda l: '\t'.join(l), ll)))

uint1, unmapped1 = to_uniprot(map_int(int1))
uint2, unmapped2 = to_uniprot(map_int(int2))

pa = pypath.PyPath(ncbi_tax_id = 9606, loops = True)

pa.load_omnipath()
pa.load_ptms()

res1 = lookup_interactions(uint1, pa)
res2 = lookup_interactions(uint2, pa)

table('mean.tab', res1)
table('var.tab', res2)