# %% [code]
# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

import matplotlib.pyplot as plt

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))

        
def testmap_data(x):
    return x.split('\n', 1)[1].replace('\n', '');

def testmap_group(x):
    return x.split('\n', 1)[0].split('\t',1)[1];

def testmap_id(x):
    return x.split('\n', 1)[0].split('\t',1)[0];

def trainmap_id(x):
    return x.split('\n', 1)[0].split(' ',1)[0];

def trainmap_header(x):
    return x.split('\n', 1)[0];

def trainmap_group(x):
    if not 'OX=' in x:
        return "";
    return x.split('OX=', 1)[1].split(' ',1)[0];

def gobasic_id_ext(x):
    return (x.split('\n', 2)[1]).split(': ',1)[1];


def seqData_init(m):
    zids = np.array(list(map(testmap_id, m)));
    zgroups = np.array(list(map(testmap_group, m)));
    zdats = np.array(list(map(testmap_data, m)), dtype=object);
    return zids, zgroups, zdats;


def seqData_init2(m):
    zids = np.array(list(map(trainmap_id, m)));
    zhead = list(map(trainmap_header, m));
    zgroups = np.array(list(map(trainmap_group, zhead)));
    zdats = np.array(list(map(testmap_data, m)), dtype=object);
    return zids, zhead, zgroups, zdats; #zgroups,


        
f = open("/kaggle/input/cafa-5-protein-function-prediction/Train/train_sequences.fasta");
w = f.read();
f.close();

f5 = open("/kaggle/input/cafa-5-protein-function-prediction/Test (Targets)/testsuperset.fasta");
zstr = f5.read();
f5.close();

f = open("/kaggle/input/cafa-5-protein-function-prediction/Train/go-basic.obo");
xstr = f.read();
f.close();

#w2 = w.split('\n>')

f2 = np.array(pd.read_csv("/kaggle/input/cafa-5-protein-function-prediction/Train/train_terms.tsv", sep='\t'));

f4 = np.array(pd.read_csv("/kaggle/input/cafa-5-protein-function-prediction/Test (Targets)/testsuperset-taxon-list.tsv", sep='\t', encoding='iso-8859-1'));

w2 = w.split('\n>');

w2[0] = w2[0][1:];

zstr2 = zstr.split('\n>');

zstr2[0] = zstr2[0][1:];

xstr2 = xstr.split("\n[Term]");

xstr2 = xstr2[1:];


test_id, test_g, test_seq = seqData_init(zstr2);

train_id, train_h, train_g, train_seq = seqData_init2(w2); #


x3 = list(map(gobasic_id_ext, xstr2));

gb_index =  dict(zip(x3, range(len(x3))));

gb_data = xstr2;


tid_index = dict(zip(train_id, range(train_id.shape[0])));


xmatch = 0;

xm_list = np.zeros((test_id.shape[0]));

for i in range(test_id.shape[0]):
    if test_id[i] in tid_index:
        xmatch += 1;
        xm_list[i] = 1;

print(xmatch, "id'd test set proteins");

print("train set:", train_seq.shape, ", test set:", test_seq.shape, ", go terms:", len(gb_data));

print(f2.shape, "protein:go table");

pg_ids = pd.unique(f2[:,0]);
pg_gos = pd.unique(f2[:,1]);

print("listed protein ids:", pg_ids.shape, ", listed go terms:", pg_gos.shape)

pg_index = dict(zip(pg_ids, range(pg_ids.shape[0])));

pg_index2 = dict(zip(pg_gos, range(pg_gos.shape[0])));

pg_mat = np.zeros((pg_ids.shape[0], pg_gos.shape[0]), np.int8);

xmatch = 0;

for i in range(test_id.shape[0]):
    if test_id[i] in pg_index:
        xmatch += 1;

print(xmatch, "go'd test set proteins");

xmatch2=0;

for i in range(train_id.shape[0]):
    if train_id[i] in pg_index:
        xmatch2 += 1;

print(xmatch2, "go'd train set proteins");

test_human = test_g == '9606';

print(np.sum(test_human & (xm_list==0)), " unlabeled h test");
print(np.sum(test_human & (xm_list==1)), " labeled h test");

print(np.sum((test_human==0) & (xm_list==0)), " unlabeled not-h test");
print(np.sum((test_human==0) & (xm_list==1)), " labeled not-h test");

for i in range(f2.shape[0]):

    if i%250000 == 0:
        print("iter ",i);
        
    pg_mat[pg_index[f2[i, 0]], pg_index2[f2[i,1]]] = 1;
    

print("completed", pg_mat.shape);

pgs1 = np.sum(pg_mat, axis=0);
pgs2 = np.sum(pg_mat, axis=1);

plt.figure();
plt.plot(np.minimum(np.sort(pgs1)[::500], 100));
plt.plot(np.minimum(np.sort(pgs2)[::500], 100));


plt.figure();
plt.plot(np.maximum(np.sort(pgs1)[::500], 100));
plt.plot(np.maximum(np.sort(pgs2)[::500], 100));

#pd.DataFrame(pg_mat).to_csv("protein_go.csv");

# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All" 
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session