import sys;
import numpy as np;
import random;
import time;
if (len(sys.argv) != 7):
	print("Usage: python ./get-feats.py [shared-list] [human.ncount] [yeast.ncount] [k] [feat] [out]");
	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[6];

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

# helper functions

# read GDVs, create map so that node name returns GDVs
def get_gdvs(file):
	gdvs = {};
	with open(file, "r") as fi:
		for line in fi:
			larr = line.split();
			prot = larr[0];
			feat = list(map(int, larr[1:]));
			gdvs[prot] = feat;

	return gdvs;

# given two GDVs, return absolute difference
def get_gdvdiff(pair, gdm1, gdm2):
	
	gdv1 = gdm1[pair[0]];
	gdv2 = gdm2[pair[1]];
	gp = [];
	for i in range(len(gdv1)):
		gp.append(np.absolute(gdv1[i]-gdv2[i]));
	return gp;

# given two GDVs, compute the weighted difference according to the GDV-similarity calucaltion
def get_gdvsimdiff(pair, gdm1, gdm2):
	weight_factor = [1, 2, 2, 2, 3, 4, 3, 3, 4, 3, 4, 4, 4, 4, 3, 4, 6, 5, 4, 5, 6, 6, 4, 4, 4, 5, 7, 4, 6, 6, 7, 4, 6, 6, 6, 5, 6, 7, 7, 5, 7, 6, 7, 6, 5, 5, 6, 8, 7, 6, 6, 8, 6, 9, 5, 6, 4, 6, 6, 7, 8, 6, 6, 8, 7, 6, 7, 7, 8, 5, 6, 6, 4]
	gdv1 = gdm1[pair[0]];
	gdv2 = gdm2[pair[1]];

	num_orbits = len(gdv1);
	weight = [];
	sum_coef = 0;
	for i in range(0, num_orbits):
	    weight.append(1 - np.log10(weight_factor[i]) / np.log10(num_orbits));
	    sum_coef += weight[i];

	gdvsim = [];
	for i in range(len(gdv1)):
		di = weight[i] * np.absolute(np.log10(gdv1[i] + 1) - np.log10(gdv2[i] + 1)) / np.log10(np.maximum(gdv1[i], gdv2[i]) + 2);
		gdvsim.append(di);
	return gdvsim;

# return the feature
def get_feat(pair, gdm1, gdm2, feat):
	if feat == "gdvdiff":
		out = get_gdvdiff(pair, gdm1, gdm2);
	elif feat == "gdvsim":
		out = get_gdvsimdiff(pair, gdm1, gdm2);
	else:
		out = [];
	return out;

humangdvs = get_gdvs(humnf);
yeastgdvs = get_gdvs(yeanf);

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


atleastkfile = outf.split("/")[-1].split(".")[0];
genlistetimestart = time.time();
# get the list of positive and negative pairs, but also time how long it takes
with open(outstem + "human-yeast-prot-go-list-" + atleastkfile + ".txt.time", "w") as fo:
	hd = {};
	yd = {};
	with open(sharedf, "r") as fi:
		for line in fi:
			larr = line.split();
			hprot = larr[0];
			yprot = larr[1];
			if len(larr) >= k + 2:
				sharedterms = larr[2:];
				if len(set(sharedterms)) >= k:
					pos_pairs.append((hprot, yprot));
					hd[hprot] = list(set(sharedterms));
					yd[yprot] = list(set(sharedterms));
			elif len(larr) == 2:
				neg_pairs.append((hprot, yprot));

	fo.write(str(time.time() - genlistetimestart));

# randomly sample negative pairs 10 times, write the features out to a file with their corresponding labels
seeds = range(0, 10);
for s in seeds:
	starttime = time.time();
	random.seed(s);
	print(len(pos_pairs))
	neg_samp_pairs = random.sample(neg_pairs, len(pos_pairs));
	print(len(neg_samp_pairs));

	with open(outf + str(s), "w") as fo:
		for p in pos_pairs:
			feat = get_feat(p, humangdvs, yeastgdvs, infeat);
			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, humangdvs, yeastgdvs, infeat);
			fo.write(n[0] + "\t" + n[1] + "," + ",".join(list(map(str, feat))));
			fo.write(",0\n");

	with open(outf + str(s) + ".time", "w") as fo:
		fo.write(str(time.time() - starttime));

