import numpy as np
import networkx as nx

"""Schema
A[i,j] = 1: i -> j
A[i, j] = A[j,i] = -1 : i <-> j
A[i,j] = 2: i -o j
A[i,j] = A[j,i] = 1: i - j
"""

def is_pot_parent(pag):
    A =np.zeros(pag.shape)
    d = len(A)

    # i - j, i -> j, i -o j
    # i o- j, i o-> j, i o-o j 
    for i in range(d):
        for j in range(i+1, d):
            # Directed edge
            if pag[i, j] == 1 and pag[j,i] == 0:
                A[i, j] = 1
            if pag[i, j] == 0 and pag[j,i] == 1:
                A[j, i] = 1

            # i -o j
            if pag[i, j] == 2 and pag[j,i] == 0:
                A[i, j] = 1
                A[j, i] = 1

            # i o- j
            if pag[i, j] == 0 and pag[j,i] == 2:
                A[i, j] = 1
                A[j, i] = 1

            # i o-o j
            if pag[i, j] == 2 and pag[j,i] == 2:
                A[i, j] = 1
                A[j, i] = 1

            # i o-> j
            if pag[i, j] == 1 and pag[j,i] == 2:
                A[i, j] = 1

            # i <-o j
            if pag[i, j] == 2 and pag[j,i] == 1:
                A[j, i] = 1

            # Undirected edge
            if pag[i, j] == 1 and pag[j,i] == 1:
                A[i, j] = 1
                A[j, i] = 1

    return A
                

def is_ancestor(pag):
    P = np.zeros(pag.shape) # preliminary matrix to search paths
    d = len(P)

    # i - j, i -> j, i -o j: P[i,j] = 1
    # i - j, j -> i, i -o j: P[j,i] = 1
    for i in range(d):
        for j in range(i+1, d):
            # Directed edge
            if pag[i, j] == 1 and pag[j,i] == 0:
                P[i, j] = 1
            if pag[i, j] == 0 and pag[j,i] == 1:
                P[j, i] = 1

            # i -o j
            if pag[i, j] == 2 and pag[j,i] == 0:
                P[i, j] = 1

            # i o- j
            if pag[i, j] == 0 and pag[j,i] == 2:
                P[j, i] = 1

            # Undirected edge
            if pag[i, j] == 1 and pag[j,i] == 1:
                P[i, j] = 1
                P[j, i] = 1

    # Find is_ancestor adjacency
    A = np.zeros(P.shape)
    G = nx.from_numpy_array(P, create_using=nx.DiGraph)
    for i in range(d):
        for j in range(i+1, d):
            if nx.has_path(G, i, j):
                A[i, j] = 1
            if nx.has_path(G, j, i):
                A[j, i] = 1
    return A


def is_pot_ancestor(pag):
    P = is_pot_parent(pag) # Preliminary matrix
    d = len(P)

    # Find is_pot_ancestor adjacency
    A = np.zeros(P.shape)
    G = nx.from_numpy_array(P, create_using=nx.DiGraph)
    for i in range(d):
        for j in range(i+1, d):
            if nx.has_path(G, i, j):
                A[i, j] = 1
            if nx.has_path(G, j, i):
                A[j, i] = 1
    return A
