from HECI import *
from NNCL import *
from os import listdir
import pandas as pd

np.set_printoptions(precision=4)
# pairs with more than 2 variables in tuebingen
not_bivariate_tuebingen = [47, 52,53,54, 55, 70, 71, 105, 107]


# testing HEC
def testDataset(path, tuebingen, mode, addvar=False):
    if tuebingen:
        skipped = []
        causal_direction_pred = np.zeros(108)
        causal_direction_gt = np.zeros(108)
        weights = np.zeros(108)
        f = open(path+"/pairmeta.txt")
        for i in range(108):
            line = f.readline()
            information = line.split(" ")[1]
            if information == "1":
                causal_direction_gt[i] = 1
            else:
                causal_direction_gt[i] = 0
            weights[i] = line.split()[5]
        f.close()
        rescaled = 0
        for i in range(1, 109):
            if i in not_bivariate_tuebingen:
                skipped.append(i-1)
                continue
            datapath = path+"/pair{:04d}.txt".format(i)
            data = np.loadtxt(datapath)
            xsrc = data[:, 0]
            ysrc = data[:, 1]
            _,score1,score2 = HECI(xsrc,ysrc,standardize=addvar)

            if score1 < score2:
                causal_direction_pred[i-1] = 1
            else:
                causal_direction_pred[i-1] = 0

        correct = causal_direction_gt == causal_direction_pred
        correct = np.delete(correct,skipped)
        weights = np.delete(weights,skipped)
        
        weighted_accuracy = sum(correct * weights)/sum(weights)
        print("_________________________________________")
        print("Dataset:", path.split("/")[-1])
        print("HEC Accuracy:", weighted_accuracy)
        return weighted_accuracy


    else:
        causal_direction_gt = np.loadtxt(path + "/pairs_gt.txt")
        weights = np.ones((len(causal_direction_gt),))
        num_pairs = len(causal_direction_gt)
        causal_direction_pred = np.zeros((num_pairs,))

        for i in range(num_pairs):
            
            if mode == 1:
                datapath = path+"/pair_{:d}.txt".format(i+1)
                data = np.loadtxt(datapath,delimiter=",",usecols=(1,2),skiprows=1)
                xsrc = data[:, 0]
                ysrc = data[:, 1]
            elif mode == 2:
                datapath = path+"/pair_{:d}.txt".format(i+1)
                data = np.loadtxt(datapath)
                xsrc = data[:, 0]
                ysrc = data[:, 1]
            _ , score1, score2 = HECI(xsrc,ysrc,standardize=addvar)

            if score1 < score2:
                causal_direction_pred[i] = 1
            else:
                causal_direction_pred[i] = 0

        correct = causal_direction_gt == causal_direction_pred
        print("_________________________________________")
        print("Dataset:", path.split("/")[-1])
        print("HECI Accuracy:", np.mean(correct))
        return np.mean(correct)
        
def testNNCL(path,tuebingen,mode):
    if tuebingen:
        causal_direction_pred = np.zeros(108)
        causal_direction_gt = np.zeros(108)
        skipped = []
        weights = np.zeros(108)
        f = open(path+"/pairmeta.txt")
        for i in range(108):
            line = f.readline()
            information = line.split(" ")[1]
            if information == "1":
                causal_direction_gt[i] = 1
            else:
                causal_direction_gt[i] = 0
            weights[i] = line.split()[5]
        f.close()
        for i in range(1, 109):
            if i in not_bivariate_tuebingen:
                skipped.append(i-1)
                continue
            datapath = path+"/pair{:04d}.txt".format(i)
            data = np.loadtxt(datapath)
            xsrc = data[:, 0]
            ysrc = data[:, 1]
            direction,_,_ = non_invertible_causal(xsrc, ysrc)
            causal_direction_pred[i-1] = direction

        correct = causal_direction_gt == causal_direction_pred   
        correct = np.delete(correct,skipped)
        weights = np.delete(weights,skipped)
        
        weighted_accuracy = sum(correct * weights)/sum(weights)
        print("_________________________________________")
        print("Dataset:", path.split("/")[-1])
        print("NNCL Accuracy:", weighted_accuracy)
        return weighted_accuracy   
    else:
        causal_direction_gt = np.loadtxt(path + "/pairs_gt.txt")
        weights = np.ones((len(causal_direction_gt),))
        num_pairs = len(causal_direction_gt)
        causal_direction_pred = np.zeros((num_pairs,))

        for i in range(num_pairs):
            if mode == 1:
                datapath = path+"/pair_{:d}.txt".format(i+1)
                data = np.loadtxt(datapath,delimiter=",",usecols=(1,2),skiprows=1)
                xsrc = data[:, 0]
                ysrc = data[:, 1]
            elif mode == 2:
                datapath = path+"/pair_{:d}.txt".format(i+1)
                data = np.loadtxt(datapath)
                xsrc = data[:, 0]
                ysrc = data[:, 1]
            direction,_,_= non_invertible_causal(xsrc, ysrc)
            causal_direction_pred[i] = direction

        correct = causal_direction_pred == causal_direction_gt
        print("_________________________________________")
        print("Dataset:", path.split("/")[-1])
        print("NNCL Accuracy:", np.mean(correct))
        return np.mean(correct)
    
    

if __name__ == "__main__":

    accuracies_nncl = []
    accuracies_hec = []
    #benchmark on our own datasets
    linear_path = "Data/Linear/"
    for dataset in listdir(linear_path):
        acc_nncl = testNNCL(linear_path+dataset,False,2)
        acc_hec = testDataset(linear_path+dataset,False,2)
        accuracies_nncl.append(acc_nncl)
        accuracies_hec.append(acc_hec)
    df_hec = pd.DataFrame()
    df_hec["Dataset"] = list(listdir(linear_path))
    df_nncl = df_hec.copy()
    df_hec["HECI"] = accuracies_hec
    df_nncl["NNCL"] = accuracies_nncl
    df_hec.to_csv("Results/linear_heci.tab",sep="\t")
    df_nncl.to_csv("Results/linear_nncl.tab",sep="\t")

    accuracies_nncl = []
    accuracies_hec = []
    #benchmark on our own datasets
    heteroscedastic_benchmark_path = "Data/HeteroscedasticPairs/"
    for dataset in listdir(heteroscedastic_benchmark_path):
        acc_nncl = testNNCL(heteroscedastic_benchmark_path+dataset,False,2)
        acc_hec = testDataset(heteroscedastic_benchmark_path+dataset,False,2)
        accuracies_nncl.append(acc_nncl)
        accuracies_hec.append(acc_hec)
    df_hec = pd.DataFrame()
    df_hec["Dataset"] = list(listdir(heteroscedastic_benchmark_path))
    df_nncl = df_hec.copy()
    df_hec["HECI"] = accuracies_hec
    df_nncl["NNCL"] = accuracies_nncl
    df_hec.to_csv("Results/heteroscedastic_heci.tab",sep="\t")
    df_nncl.to_csv("Results/heteroscedastic_nncl.tab",sep="\t")

    accuracies_nncl = []
    accuracies_hec = []
    # test AN, LS and MN datasets by Tagakovska et al, 2019
    an_benchmark_path = "Data/ANLSMN/"
    for dataset in listdir(an_benchmark_path):
        acc_nncl = testNNCL(an_benchmark_path+dataset,False,1)
        acc_hec = testDataset(an_benchmark_path+dataset,False,1,True)
        accuracies_nncl.append(acc_nncl)
        accuracies_hec.append(acc_hec)
    datasets = list(listdir(an_benchmark_path))
    df_hec = pd.DataFrame()
    df_hec["Dataset"] = list(listdir(an_benchmark_path))
    df_nncl = df_hec.copy()
    df_hec["HECI"] = accuracies_hec
    df_nncl["NNCL"] = accuracies_nncl
    df_hec.to_csv("Results/anlsmn_heci.tab",sep="\t")
    df_nncl.to_csv("Results/anlsmn_nncl.tab",sep="\t")

    accuracies_nncl = []
    accuracies_hec = []
    # test AN-s, LS-s datasets by Tagakovska et al, 2019
    an_benchmark_path = "Data/ANLS-s/"
    for dataset in listdir(an_benchmark_path):
        acc_nncl = testNNCL(an_benchmark_path+dataset,False,1)
        acc_hec = testDataset(an_benchmark_path+dataset,False,1,True)
        accuracies_nncl.append(acc_nncl)
        accuracies_hec.append(acc_hec)
    datasets = list(listdir(an_benchmark_path))
    df_hec = pd.DataFrame()
    df_hec["Dataset"] = list(listdir(an_benchmark_path))
    df_nncl = df_hec.copy()
    df_hec["HECI"] = accuracies_hec
    df_nncl["NNCL"] = accuracies_nncl
    df_hec.to_csv("Results/anls-s_heci.tab",sep="\t")
    df_nncl.to_csv("Results/anls-s_nncl.tab",sep="\t")

    # test Tuebingen real world cause-effect pairs
    accuracies_nncl = []
    accuracies_hec = []
    acc_nncl = testNNCL("Data/Tuebingen",True,1)
    acc_hec = testDataset("Data/Tuebingen",True,1)
    accuracies_nncl.append(acc_nncl)
    accuracies_hec.append(acc_hec)
    df_hec = pd.DataFrame()
    df_hec["Dataset"] = ["Tuebingen"]
    df_nncl = df_hec.copy()
    df_hec["HECI"] = accuracies_hec
    df_nncl["NNCL"] = accuracies_nncl

    df_hec.to_csv("Results/tueb_heci.tab",sep="\t")
    df_nncl.to_csv("Results/tueb_nncl.tab",sep="\t")

