import sys;
import numpy as np;
import random;
import random;
import time;

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-with-prot-pairs.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]);
rarethresh = sys.argv[5];

print(dataf + " " + str(s) + " " + str(percent));

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

    X = [];
    y = [];
    Xp = [];
    yp = [];
    Xn = [];
    yn = [];
    with open(file, "r") as fi:
        for line in fi:
            # the first element of the feature vector is the protein pair
            larr = [line.split(",")[0]] + list(map(float, line.split(",")[1:]));
            
            feats = larr[0:-1];
            lab = np.array(larr[-1]);
            
            if larr[-1] == 1:
                Xp.append(feats);
                yp.append(lab);
            else:
                Xn.append(feats);
                yn.append(lab);
            
            X.append(feats);
            y.append(lab);
    
    ssize = int(perc * len(Xn));

    # randomly sample from feature vectors, but make sure labels are the corresponding ones
    Xtr = [Xp.pop(random.randrange(len(Xp))) for _ in range(ssize)] + \
          [Xn.pop(random.randrange(len(Xn))) for _ in range(ssize)]
    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);

def format_feats(X):
    Xout = [];
    for arr in X:
        Xout.append(np.array(arr[1:]));
    return np.array(Xout);

def log_reg(file, perc, seed):
    _Xtr, ytr, _Xte, yte = get_data_from_file_rand(file, perc, seed);

    # ran into some issues with data formatting, so hacked some code to ensure it was correct.
    Xtr = format_feats(_Xtr);
    Xte = format_feats(_Xte);

    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;

    LogR = LogisticRegression();
    LogR.fit(Xtr, ytr);

    ypred = LogR.predict(Xte);
    return _Xte, ypred;

trainingtimestart = time.time();

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

if outresfile in listdir(filedir):
    print("skipping " + outresfile + " because it already exists");
    sys.exit();

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

trainingtimeend = time.time();
with open(outdir + "-training-time" + str(percent * 100) + ".time", "w") as fo:
    fo.write(str(trainingtimeend - trainingtimestart));

predtimestart = time.time();

# any pair that is predicted as 1 is added to the alignment
outalnf = outdir + "-" + str(percent * 100) + ".txt" + str(s)
with open(outalnf, "w") as fo:
    for i, pred in enumerate(ypred):
        if pred == 1:
            fo.write(_Xte[i][0] + "\n");
predtimeend = time.time();
with open(outdir + "-pred-time" + str(percent * 100) + ".time", "w") as fo:
    fo.write(str(predtimeend- predtimestart));

# setup for protein function prediction. See https://www3.nd.edu/~cone/LNA_GNA/download.html
gofile = "data/real-world/atleast" + dataf.split("atleast")[-1][0];
if "EXP" in dataf:
    gofile += "-EXP-" + rarethresh + ".txt";


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

with open(configf, "w") as fo:
    fo.write("Aligned node pairs: " + "/".join(outalnf.split("/")[1:]) + str(s) + "\n");
    fo.write("Network1: ./data/real-world/human.gw\n");
    fo.write("Network2: ./data/real-world/yeast.gw\n");
    fo.write("GO data for Network1: " + gofile + "\n");
    fo.write("GO data for Network2: " + gofile + "\n");
    fo.write("Measures: NCV,GS3,NCV-GS3,GC,P-PF,R-PF,F-PF\n")

pftimestart = time.time();
call(["python", "./code/main.py", configf, outalnf + ".out" + str(s)])
pftimeend = time.time();
with open(outalnf + ".time", "w") as fo:
    fo.write(str(pftimeend - pftimestart));