import sys;
import numpy as np;
import random;
import random;
from sklearn.linear_model import LinearRegression, LogisticRegression, Lasso;
from subprocess import call;
from os import listdir;

if (len(sys.argv) != 6):
    print("Usage: python ./prediction.py [data] [outf] [seed] [perc] [thresh]");
    sys.exit();

dataf = sys.argv[1];
outdir = sys.argv[2];
s = int(sys.argv[3]);
percent = float(sys.argv[4]);
thresh = int(sys.argv[5]);

# these are just used to get the list of proteins
humanf = "./data/embeddings/human.emb";
yeastf = "./data/embeddings/yeast.emb";

def getprots(file):
    prots = [];
    with open(humanf, "r") as fi:
        for line in fi:
            larr = line.split();
            prots.append(larr[0][1:]); # remove the initial 'h' or 's'
    return prots;

humanprots = getprots(humanf);
yeastprots = getprots(yeastf);

def get_data_from_file_rand(file, perc, seed):
    random.seed(seed + 0);

    X = [];
    y = [];
    Xp = [];
    yp = [];
    Xn = [];
    yn = [];

    with open(file, "r") as fi:
        for line in fi:
 
            # the first column contains the two proteins, tab delimited
            # the remaining columns are comma delimited
            larr = line.split(",")[0].split("\t") + line.split(",")[1:];
            larr[-1] = larr[-1].split()[0]; # remove the ending \n
            larr = larr[0:2] + list(map(float, larr[2:])); # remake into array

            feats = larr[0:-1];
            lab = np.array(larr[-1]);

            # positive and negative classes
            if larr[-1] == 1:
                Xp.append(feats);
                yp.append(lab);
            else:
                Xn.append(feats);
                yn.append(lab);
            
            X.append(feats);
            y.append(lab);

    # amount of training data to use
    ssize = int(perc * len(Xn));

    # sample, but create equal sized classes
    # pop removes the element from Xp/Xn
    Xtr = [Xp.pop(random.randrange(len(Xp))) for _ in range(ssize)] + \
          [Xn.pop(random.randrange(len(Xn))) for _ in range(ssize)]

    # doesn't matter if the positive labels are rearranged among themselves
    ytr = [yp.pop(random.randrange(len(yp))) for _ in range(ssize)] + \
          [yn.pop(random.randrange(len(yn))) for _ in range(ssize)]

    return Xtr, np.array(ytr), Xp + Xn, np.array(yp + yn);

# remove the first two elements, which are the proteins
def format_feats(X):
    Xout = [];
    for arr in X:
        Xout.append(np.array(arr[2:]));
    return np.array(Xout);

def log_reg(file, perc, seed):
    
    # need to keep _Xtr and _Xte to recover the proteins for the alignment
    _Xtr, ytr, _Xte, yte = get_data_from_file_rand(file, perc, seed);

    Xtr = format_feats(_Xtr);
    Xte = format_feats(_Xte);

    # data requires reformatting for some reason
    newXtr = [];
    for i in range(len(Xtr)):
        a = [];
        for j in range(len(Xtr[0])):
            a.append(Xtr[i][j]);
        newXtr.append(a);
    Xtr = newXtr;
    newXte = [];
    for i in range(len(Xte)):
        a = [];
        for j in range(len(Xte[0])):
            a.append(Xte[i][j]);
        newXte.append(a);
    Xte = newXte;

    # logistic regression classifier
    LogR = LogisticRegression();
    LogR.fit(Xtr, ytr);

    ypred = LogR.predict(Xte);

    return _Xte, ypred, _Xtr;



# file naming setup
outalnf = outdir + "-" + str(percent * 100) + ".txt";
outresfile = outalnf + ".out" + str(s);
outresfile = outresfile.split("/")[-1];
filedir = "/".join(outdir.split("/")[0:-1])
outalnf = outdir + "-" + str(percent * 100) + ".txt" + str(s)

_Xte, ypred, _Xtr = log_reg(dataf + str(s), percent, s);

# write the alignment (any pair predicted as functionally related)
with open(outalnf, "w") as fo:

    # the order of predictions and _Xte is the same, which is why this works
    for i, pred in enumerate(ypred):
        if pred == 1:
            fo.write(_Xte[i][0] + "\t" + _Xte[i][1] + "\n");

# ground truth data
gofile = "data/atleast" + dataf.split("atleast")[-1][0];
tmpgof = gofile + "-EXP.txt";

if "EXP" in dataf:
    gofile += "-EXP" + str(thresh) + ".txt";
else:
    gofile += ".txt";

newgof = gofile;

outalnf = outdir + "-" + str(percent * 100) + ".txt";
configf = outdir + "-" + str(percent * 100) + ".config" + str(s);

# format of the file for the protein functional prediction framework
with open(configf, "w") as fo:
    fo.write("Aligned node pairs: " + "/".join(outalnf.split("/")[1:]) + str(s) + "\n");
    fo.write("Network1: data/networks/human.gw\n");
    fo.write("Network2: data/networks/yeast.gw\n");
    fo.write("GO data for Network1: " + gofile + "\n");
    fo.write("GO data for Network2: " + gofile + "\n");
    fo.write("Measures: P-PF,R-PF,F-PF\n");

call(["python", "./scripts/main.py", configf, outalnf + ".out" + str(s)])
