# Path: experiments/exp_setup.py
# SDCIT_RESULT_DIR = SDCIT_WORKING_DIR + '/results'
# 
# PARALLEL_JOBS = multiprocessing.cpu_count() // 2
# 
# Path: experiments/exp_utils.py
# def read_chaotic(independent, gamma, trial, N, dir_at=SDCIT_DATA_DIR + '/'):
#     X, Y, Z = read_chaotic_data(independent, gamma, trial, N, dir_at)
#     kx, ky, kz = rbf_kernel_median(X, Y, Z)
#     Dz = K2D(kz)
#     return kx, ky, kz, Dz
# 
# def read_postnonlinear_noise(independent, noise, trial, N, dir_at=SDCIT_DATA_DIR + '/'):
#     X, Y, Z = read_postnonlinear_noise_data(independent, noise, trial, N, dir_at)
#     kx, ky, kz = rbf_kernel_median(X, Y, Z)
# 
#     dist_mat_file = os.path.expanduser(dir_at + 'dist_{}_{}_{}_{}_postnonlinear.mat'.format(noise, trial, independent, N))
#     mat_load = scipy.io.loadmat(dist_mat_file, squeeze_me=True, struct_as_record=False)
#     Dz = np.array(mat_load['D'])
# 
#     return cythonize(kx, ky, kz, Dz)
# 
# def chaotic_configs():
#     return list(itertools.product([0, 1], [200, 400], ['0.0', '0.1', '0.2', '0.3', '0.4', '0.5']))
# 
# def postnonlinear_noise_configs():
#     return list(itertools.product(range(5), [0, 1], [200, 400])) + \
#            list(itertools.product([9, 19, 49], [0, 1], [400]))  # high-dimensional
# 
# Path: sdcit/sdcit_mod.py
# def SDCIT(Kx: np.ndarray, Ky: np.ndarray, Kz: np.ndarray, Dz=None, size_of_null_sample=1000, with_null=False, seed=None, adjust=True, to_shuffle=True):
#     """SDCIT (Lee and Honavar, 2017)
# 
#     Parameters
#     ----------
#     Kx : np.ndarray
#         N by N kernel matrix of X
#     Ky : np.ndarray
#         N by N kernel matrix of Y
#     Kz : np.ndarray
#         N by N kernel matrix of Z
#     Dz : np.ndarray
#         N by N pairwise distance matrix of Z
#     size_of_null_sample : int
#         The number of samples in a null distribution
#     with_null : bool
#         If true, resulting null distribution is also returned
#     seed : int
#         Random seed
#     adjust : bool
#         whether to adjust null distribution and test statistics based on 'permutation error' information
#     to_shuffle : bool
#         shuffle the order of given data at the beginning, which minimize possible issues with getting a bad permutation
# 
#     References
#     ----------
#         Lee, S., Honavar, V. (2017). Self-Discrepancy Conditional Independence Test.
#         In Proceedings of the Thirty-third Conference on Uncertainty in Artificial Intelligence. Corvallis, Oregon: AUAI Press.
#     """
#     if seed is not None:
#         np.random.seed(seed)
# 
#     if Dz is None:
#         Dz = K2D(Kz)
# 
#     if to_shuffle:
#         Kx, Ky, Kz, Dz = shuffling(seed, Kx, Ky, Kz, Dz)  # categorical Z may yield an ordered 'block' matrix and it may harm permutation.
# 
#     Kxz = Kx * Kz
# 
#     test_statistic, error_statistic, mask, _ = MMSD(Ky, Kz, Kxz, Dz)
#     mask, Pidx = mask_and_perm(penalized_distance(Dz, mask))
# 
#     # avoid permutation between already permuted pairs.
#     mmsd_distr_under_null, error_distr_under_null = emp_MMSD(Kxz, Ky[np.ix_(Pidx, Pidx)], Kz, penalized_distance(Dz, mask), size_of_null_sample)
# 
#     if adjust:
#         fix_null, fix_test_statistic = adjust_errors(error_distr_under_null, mmsd_distr_under_null, error_statistic, test_statistic)
#         fix_null = fix_null - fix_null.mean()
#     else:
#         fix_null = mmsd_distr_under_null - mmsd_distr_under_null.mean()
#         fix_test_statistic = test_statistic
# 
#     if with_null:
#         return fix_test_statistic, p_value_of(fix_test_statistic, fix_null), fix_null
#     else:
#         return fix_test_statistic, p_value_of(fix_test_statistic, fix_null)


from os.path import exists
from joblib import Parallel
from joblib import delayed
from tqdm import tqdm, trange
from experiments.exp_setup import SDCIT_RESULT_DIR, PARALLEL_JOBS
from experiments.exp_utils import read_chaotic, read_postnonlinear_noise, chaotic_configs, postnonlinear_noise_configs
from sdcit.sdcit_mod import SDCIT
from sdcit.utils import *




def test_chaotic(independent, gamma, trial, N):
    np.random.seed(trial)
    mmsd, pval = SDCIT(*read_chaotic(independent, gamma, trial, N), seed=trial, to_shuffle=False)
    return independent, gamma, trial, N, mmsd, pval


def test_postnonlinear(independent, noise, trial, N):
    np.random.seed(trial)
    mmsd, pval = SDCIT(*read_postnonlinear_noise(independent, noise, trial, N), seed=trial, to_shuffle=False)
    return independent, noise, trial, N, mmsd, pval


def main():
    with Parallel(PARALLEL_JOBS) as parallel:
        if not exists(SDCIT_RESULT_DIR + '/sdcit_chaotic.csv'):
            for independent, N, gamma in tqdm(chaotic_configs()):
                outs = parallel(delayed(test_chaotic)(independent, gamma, trial, N) for trial in trange(300))
                with open(SDCIT_RESULT_DIR + '/sdcit_chaotic.csv', 'a') as f:
                    for out in outs:
                        print(*out, sep=',', file=f, flush=True)
        else:
            print('skipping SDCIT on chaotic time series data')


# Path: experiments/exp_setup.py
# SDCIT_RESULT_DIR = SDCIT_WORKING_DIR + '/results'
# 
# PARALLEL_JOBS = multiprocessing.cpu_count() // 2
# 
# Path: experiments/exp_utils.py
# def read_chaotic(independent, gamma, trial, N, dir_at=SDCIT_DATA_DIR + '/'):
#     X, Y, Z = read_chaotic_data(independent, gamma, trial, N, dir_at)
#     kx, ky, kz = rbf_kernel_median(X, Y, Z)
#     Dz = K2D(kz)
#     return kx, ky, kz, Dz
# 
# def read_postnonlinear_noise(independent, noise, trial, N, dir_at=SDCIT_DATA_DIR + '/'):
#     X, Y, Z = read_postnonlinear_noise_data(independent, noise, trial, N, dir_at)
#     kx, ky, kz = rbf_kernel_median(X, Y, Z)
# 
#     dist_mat_file = os.path.expanduser(dir_at + 'dist_{}_{}_{}_{}_postnonlinear.mat'.format(noise, trial, independent, N))
#     mat_load = scipy.io.loadmat(dist_mat_file, squeeze_me=True, struct_as_record=False)
#     Dz = np.array(mat_load['D'])
# 
#     return cythonize(kx, ky, kz, Dz)
# 
# def chaotic_configs():
#     return list(itertools.product([0, 1], [200, 400], ['0.0', '0.1', '0.2', '0.3', '0.4', '0.5']))
# 
# def postnonlinear_noise_configs():
#     return list(itertools.product(range(5), [0, 1], [200, 400])) + \
#            list(itertools.product([9, 19, 49], [0, 1], [400]))  # high-dimensional
# 
# Path: sdcit/sdcit_mod.py
# def SDCIT(Kx: np.ndarray, Ky: np.ndarray, Kz: np.ndarray, Dz=None, size_of_null_sample=1000, with_null=False, seed=None, adjust=True, to_shuffle=True):
#     """SDCIT (Lee and Honavar, 2017)
# 
#     Parameters
#     ----------
#     Kx : np.ndarray
#         N by N kernel matrix of X
#     Ky : np.ndarray
#         N by N kernel matrix of Y
#     Kz : np.ndarray
#         N by N kernel matrix of Z
#     Dz : np.ndarray
#         N by N pairwise distance matrix of Z
#     size_of_null_sample : int
#         The number of samples in a null distribution
#     with_null : bool
#         If true, resulting null distribution is also returned
#     seed : int
#         Random seed
#     adjust : bool
#         whether to adjust null distribution and test statistics based on 'permutation error' information
#     to_shuffle : bool
#         shuffle the order of given data at the beginning, which minimize possible issues with getting a bad permutation
# 
#     References
#     ----------
#         Lee, S., Honavar, V. (2017). Self-Discrepancy Conditional Independence Test.
#         In Proceedings of the Thirty-third Conference on Uncertainty in Artificial Intelligence. Corvallis, Oregon: AUAI Press.
#     """
#     if seed is not None:
#         np.random.seed(seed)
# 
#     if Dz is None:
#         Dz = K2D(Kz)
# 
#     if to_shuffle:
#         Kx, Ky, Kz, Dz = shuffling(seed, Kx, Ky, Kz, Dz)  # categorical Z may yield an ordered 'block' matrix and it may harm permutation.
# 
#     Kxz = Kx * Kz
# 
#     test_statistic, error_statistic, mask, _ = MMSD(Ky, Kz, Kxz, Dz)
#     mask, Pidx = mask_and_perm(penalized_distance(Dz, mask))
# 
#     # avoid permutation between already permuted pairs.
#     mmsd_distr_under_null, error_distr_under_null = emp_MMSD(Kxz, Ky[np.ix_(Pidx, Pidx)], Kz, penalized_distance(Dz, mask), size_of_null_sample)
# 
#     if adjust:
#         fix_null, fix_test_statistic = adjust_errors(error_distr_under_null, mmsd_distr_under_null, error_statistic, test_statistic)
#         fix_null = fix_null - fix_null.mean()
#     else:
#         fix_null = mmsd_distr_under_null - mmsd_distr_under_null.mean()
#         fix_test_statistic = test_statistic
# 
#     if with_null:
#         return fix_test_statistic, p_value_of(fix_test_statistic, fix_null), fix_null
#     else:
#         return fix_test_statistic, p_value_of(fix_test_statistic, fix_null)


# Path: experiments/run_SDCIT.py
from os.path import exists
from joblib import Parallel
from joblib import delayed
from tqdm import tqdm, trange
from experiments.exp_setup import SDCIT_RESULT_DIR, PARALLEL_JOBS
from experiments.exp_utils import read_chaotic, read_postnonlinear_noise, chaotic_configs, postnonlinear_noise_configs
from sdcit.sdcit_mod import SDCIT
from sdcit.utils import *





def test_chaotic(independent, gamma, trial, N):
    np.random.seed(trial)
    mmsd, pval = SDCIT(*read_chaotic(independent, gamma, trial, N), seed=trial, to_shuffle=False)
    return independent, gamma, trial, N, mmsd, pval


def test_postnonlinear(independent, noise, trial, N):
    np.random.seed(trial)
    mmsd, pval = SDCIT(*read_postnonlinear_noise(independent, noise, trial, N), seed=trial, to_shuffle=False)
    return independent, noise, trial, N, mmsd, pval


def main():
    with Parallel(PARALLEL_JOBS) as parallel:
        if not exists(SDCIT_RESULT_DIR + '/sdcit_chaotic.csv'):
            for independent, N, gamma in tqdm(chaotic_configs()):
                outs = parallel(delayed(test_chaotic)(independent, gamma, trial, N) for trial in trange(300))
                with open(SDCIT_RESULT_DIR + '/sdcit_chaotic.csv', 'a') as f:
                    for out in outs:
                        print(*out, sep=',', file=f, flush=True)
        else:
            print('skipping SDCIT on chaotic time series data')


        if not exists(SDCIT_RESULT_DIR + '/sdcit_postnonlinear.csv'):