from itertools import combinations

import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import numpy.linalg as la
import pandas as pd
import pingouin as pg
import scipy.io as sio

from scipy.stats import norm

data = sio.loadmat('/home/adhaene/Downloads/Data.mat')['D']
n, p = data.shape

graph = nx.complete_graph(p)

draw = lambda x: nx.draw_circular(x, with_labels=True, node_color='b', font_color='w')
draw(graph)

 # Using `norm` from the scipy stats package
phi_inv = lambda x: norm.ppf( norm.cdf(x) )
alpha = 0.05

 # Defining the threshold function
threshold = lambda x: phi_inv(1 - alpha/2) / np.sqrt(n - x - 3)

def phi(df, X, Y, Z=None):
    """
    Get partial correlation between X and Y controlling on set Z

    :param X: index of variable X
    :param Y: index of variable Y
    :param Z: index/indices of controlling variable(s)
    :return phi: partial correlation (Pearson)
    """
    if type(Z) is int:
        Z = str(Z)
    elif Z is not None:
        Z = [str(e) for e in Z]

    return float(df.partial_corr(x=str(X), y=str(Y), covar=Z)['r'])

 # Used to break out of the loop in a highly specific manner
class RemoveEdge(Exception): pass

def pc(data, verbose=False):
    """
    Run PC Algorithm.

    :param data: raw data N x P (samples x nodes)
    :return A: adjacency matrix
    :return separators: separating sets
    """

    graph = nx.complete_graph(p)
    df = pd.DataFrame(data=data, columns=[str(e) for e in list(range(p))])

    # initialize separators list
    separators = np.empty((p, p), dtype=list)

    # loop over degree of test
    for deg in range(p - 2):

        if verbose >= 1:
                print('\n----- ' + 'Performing {}-degree tests\n'.format(deg))

        for x in graph.nodes:

            if verbose == 2:
                print('\n=== ' + 'Checking node {}\n'.format(x))

            remove = []

            for y in graph.neighbors(x):
                try:

                    z = list(graph.neighbors(x))
                    z.remove(y)

                    if len(z) >= deg:

                        for zz in combinations(z, deg):

                            if verbose == 2:
                                print('{} ind. {} with z: {}'.format(x, y, ", ".join([str(e) for e in zz])))

                            if np.abs(phi(df, x, y, zz)) < threshold(p - 2 - len(zz)):

                                if verbose >= 1:
                                    print('{0} ind. {1} | {2}'.format(x, y, zz))

                                remove.append(y)

                                separators[x, y] = zz

                                raise RemoveEdge()

                except RemoveEdge:
                    pass

            for nei in remove:
                if (x, nei) in graph.edges:
                    graph.remove_edge(x, nei)

    return graph, separators

G, separators = pc(data, verbose=1)

A = nx.to_numpy_matrix(G)
draw(G)

def add_collider(A, X, Y, Z, verbose=False):
    """
    Add collider Z between X and Y to ajacency matrix A.
    We have X independent of Y conditioned on nothing, but
    X dependent of Y conditioned on Z.

    :param A: adjancency matrix
    :param X: index of variable X
    :param Y: index of variable Y
    :param Z: index of variabll Z
    :return B: updated adjancency matrix
    """

    if verbose:
        print('\n Adding collider: {0} -> {2} <- {1} \n'.format(X, Y, Z))

    B = np.asarray(A.copy())

    B[X, Z] = 1
    B[Y, Z] = 1
    B[Z, X] = 0
    B[Z, Y] = 0

    return B

def is_collider(df, X, Y, Z, p, separators, verbose=False):
    """
    Check in data if X -> Z <- Y

    :param df: raw data as pandas DataFrame
    :param X: index of variable X
    :param Y: index of variable Y
    :param Z: index of variable Z
    :param p: total number of variables
    :param separators: separator matrix from step 1 of the PC Algorithm
    :return is_collider: boolean
    """
    # should never go to else unless separators is None for X, Y
    sep = separators[Y, X] if separators[X, Y] is None else separators[X, Y]

    # Z is not a collider if it is a part of the separating set
    if Z in sep:
        return False

    # create new subset of separating set S + Z
    ss = [str(e) for e in sep]
    ss_z = ss + [str(Z)]

    # check on independence with ('w') and without ('wo') Z in conditioning set
    wo = float(df.partial_corr(x=str(X), y=str(Y), covar=ss)['r'])
    w = float(df.partial_corr(x=str(X), y=str(Y), covar=ss_z)['r'])

    if verbose:
        print('Checking collider: {0} -> {2} <- {1} with sep: {5}, wo {3} w {4}'.format(X, Y, Z, round(wo, 4), round(w, 4), sep))

    return np.abs(wo) < threshold(p - 2) and np.abs(w) > threshold(p - 3)

def get_colliders(A, data, separators, verbose=False):
    """
    Get colliders in graph A using raw data. The implementation
    is greedy and two-way but wit a small sample size it's fine.

    :param A: adjancency matrix
    :param data: raw data
    :param separators: separator matrix from step 1 of the PC Algorithm
    :return colliders: list of triple-tuples
    """
    graph = nx.from_numpy_matrix(np.asarray(A), create_using=nx.Graph)

    p = np.asarray(A).shape[0]
    df = pd.DataFrame(data=data, columns=[str(e) for e in list(range(p))])

    # initialize edges
    edges = []

    for x in graph.nodes:
        for z in graph.neighbors(x):
            for y in graph.neighbors(z):

                if y == x or y in graph.neighbors(x):
                    continue

                if is_collider(df, x, y, z, p, separators, verbose):
                    edges.append( (x, z,) )
                    edges.append( (y, z,) )
                    A = add_collider(A, x, y, z, verbose)

    return A, edges

B, v_edges = get_colliders(A, data, separators, verbose=True)
gr = nx.from_numpy_array(B, create_using=nx.DiGraph)

nx.draw_circular(gr.to_undirected())
nx.draw_circular(gr, edgelist=v_edges, edge_color='red', with_labels=True, node_color='b', font_color='w')

meek_edges = [(6, 3,), (3, 0,), (3, 7,), (3, 2,), (4, 3,), (0, 5,), (4, 0,), (0, 2,), (4, 5,)]
labels = dict(zip(meek_edges, ['({})'.format(i) for i in range(1, 10)]))

nx.draw_circular(gr, edgelist=v_edges, edge_color='red', with_labels=True, node_color='b', font_color='w')
nx.draw_networkx_edge_labels(gr, pos=nx.circular_layout(gr), edge_labels=labels, label_pos=0.4)
nx.draw_circular(gr, edgelist=meek_edges, edge_color='green', with_labels=True, node_color='b', font_color='w')