import sys;
import numpy as np;
import random;

if (len(sys.argv) != 8):
	print("Usage: python ./get-nodepair-feats.py [shared list] [human emb] [yeast emb] [k] [out] [go frequency] [rarity thresh]");
	sys.exit();

sharedf = sys.argv[1];
humnf = sys.argv[2];
yeanf = sys.argv[3];
k = int(sys.argv[4]);
# infeat = sys.argv[5];
outf = sys.argv[5];
meaningf = sys.argv[6];
thresh = int(sys.argv[7]);

infeat = "gdvdiff"

outstem = outf.split("atleast")[0];

def get_featmat(file):

	featmat = {};
	with open(file, "r") as fi:
		for line in fi:
			larr = line.split();
			prot = larr[0][1:];
			feat = list(map(float, larr[1:]));
			featmat[prot] = feat;
	return featmat;

# given the feature vector of node1 and the feature vector of node2, performs element-wise average to get feature of a node pair
def get_avg(pair, featmat1, featmat2):
	fv1 = featmat1[pair[0]];
	fv2 = featmat2[pair[1]];
	out = [];
	# recommend to use numpy functions instead of a loop
	for i in range(len(fv1)):
		out.append((fv1[i]+fv2[i]) / 2.0);
	return out;

# gets the feature vector of a node pair
def get_feat(pair, featmat1, featmat2):
	out = get_avg(pair, featmat1, featmat2)
	return out;

# get the list of GO terms given a rarity threshold
def get_meaningful_terms(mf, thresh):
	mterms = [];
	with open(mf, "r") as fi: 
		for line in fi:
			larr = line.split();
			term = larr[0];
			freq = int(larr[1]);
			if freq <= thresh:
				mterms.append(term);
	return mterms


human_featmat = get_featmat(humnf);
yeast_featmat = get_featmat(yeanf);

pos_pairs = [];
neg_pairs = [];
ct = 0;

mterms = get_meaningful_terms(meaningf, thresh);
atleastkfile = outf.split("/")[-1].split(".")[0];

with open(sharedf, "r") as fi:
	for line in fi:
		
		# progress checker
		# remove break if not testing
		if ct % 1000000 == 0:
			print(ct)
			# break;
		larr = line.split();
		hprot = larr[0];
		yprot = larr[1];
		
		# need this to check that terms exist
		if len(larr) >= k + 2:
			sharedterms = larr[2:];
			# check if pair shares at least k terms
			if len(set(sharedterms)) >= k:				
				tmpc = 0;
				for t in list(set(sharedterms)): # check how many terms satisfy rarity threshold
					if t in mterms:
						tmpc += 1;
				if tmpc == len(list(set(sharedterms))): # only add to positive pairs if all satisfy rarity threshold
					pos_pairs.append((hprot, yprot));

		# if NO pairs are shared, then add to negative pairs
		elif len(larr) == 2:
			neg_pairs.append((hprot, yprot));
		ct += 1;

# 10 different sampling instances
seeds = range(0, 10);
for s in seeds:
	random.seed(s);
	neg_samp_pairs = random.sample(neg_pairs, len(pos_pairs));

	with open(outf + str(s), "w") as fo:
		for p in pos_pairs:
			feat = get_feat(p, human_featmat, yeast_featmat);
			fo.write(p[0] + "\t" + p[1] + "," + ",".join(list(map(str, feat))));
			fo.write(",1\n");
		for n in neg_samp_pairs:
			feat = get_feat(n, human_featmat, yeast_featmat);
			fo.write(n[0] + "\t" + n[1] + "," + ",".join(list(map(str, feat))));
			fo.write(",0\n");


