import torch

from utilities import roc, prc

import heterograph_construction

import karateclub

import scipy.stats as st

import numpy as np


def predict(heads, tails, eid):
    pred = []
    for head, tail in zip(eid[0], eid[1]):
        pred.append(torch.sigmoid(torch.sum(heads[head] * tails[tail], dim=-1)))
    return pred


def metrics(constructor, phenotypes, drugs):
    _, eidRDB = constructor.DISNETHeterographNetworkX(withoutRepoDB=True)
    randomedges = randomEids()

    predictions = predict(phenotypes, drugs, eidRDB)

    predictionsN = predict(phenotypes, drugs, randomedges)

    preds = torch.tensor(predictions).cpu().detach()
    predsN = torch.tensor(predictionsN).cpu().detach()

    labels1 = torch.ones(len(preds))
    labels2 = torch.zeros(len(predsN))

    pure_predictions = [item.item() for sublist in [preds, predsN] for item in sublist]
    labels = torch.tensor([item for sublist in [labels1, labels2] for item in sublist])

    _, _, rocV = roc(labels, pure_predictions)
    _, _, prcV = prc(labels, pure_predictions)

    return rocV, prcV


def randomEids(size=5013):
    tensor1 = torch.randint(0, 30729, (size,))
    tensor2 = torch.randint(0, 3944, (size,))
    return tensor1, tensor2


def baselineModel(modelName, full):
    # When transforming to homogeneous phenotypes go from 0 to 30728 and drugs go 30729 from to 34672
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    constructor = heterograph_construction.DISNETConstructor(device=device)

    g, _ = constructor.DISNETHeterographNetworkX(full=full)

    # Train model and generate embedding
    if modelName == 'Role2Vec':
        model = karateclub.Role2Vec()

    elif modelName == 'DeepWalk':
        model = karateclub.DeepWalk()

    elif modelName == 'Node2Vec':
        model = karateclub.Node2Vec()

    elif modelName == 'NetMF':
        model = karateclub.NetMF()

    model.fit(g)
    embedding = model.get_embedding()

    phenotypes = torch.from_numpy(embedding[0:30729])
    drugs = torch.from_numpy(embedding[30729:34673])

    return phenotypes, drugs, constructor


if __name__ == '__main__':
    k = 20
    rocL, prcL = np.array([]), np.array([])
    for i in range(0, k):
        print(i)
        phenotypesEmbed, drugsEmbed, constructor = baselineModel('DeepWalk', True)
        rocV, prcV = metrics(constructor, phenotypesEmbed, drugsEmbed)

        rocL = np.append(rocL, rocV)
        prcL = np.append(prcL, prcV)

    rocM = sum(rocL) / k
    prcM = sum(prcL) / k

    r = st.t.interval(confidence=0.95, df=k - 1, loc=np.mean(rocL), scale=st.sem(rocL))
    p = st.t.interval(confidence=0.95, df=k - 1, loc=np.mean(prcL), scale=st.sem(prcL))

    print("AUCROC:  ", rocM, "+-", rocM-r[0])
    print("AUCPR:   ", prcM, "+-", prcM-p[0])
