import numpy as np
import pandas as pd
import graphviz
import lingam
from lingam.utils import make_prior_knowledge, make_dot

print([np.__version__, pd.__version__, graphviz.__version__, lingam.__version__])

np.set_printoptions(precision=3, suppress=True)
np.random.seed(0)

"""
Explanation: How to use prior knowledge in DirectLiNGAM
Import and settings
In this example, we need to import numpy, pandas, and graphviz in addition to lingam.
End of explanation
"""


def make_prior_knowledge_graph(prior_knowledge_matrix):
    d = graphviz.Digraph(engine='dot')
    
    labels = [f'x{i}' for i in range(prior_knowledge_matrix.shape[0])]
    for label in labels:
        d.node(label, label)

    dirs = np.where(prior_knowledge_matrix > 0)
    for to, from_ in zip(dirs[0], dirs[1]):
        d.edge(labels[from_], labels[to])

    dirs = np.where(prior_knowledge_matrix < 0)
    for to, from_ in zip(dirs[0], dirs[1]):
        if to != from_:
            d.edge(labels[from_], labels[to], style='dashed')
    return d

"""
Explanation: Utility function
We define a utility function to draw the directed acyclic graph.
End of explanation
"""


x3 = np.random.uniform(size=10000)
x0 = 3.0*x3 + np.random.uniform(size=10000)
x2 = 6.0*x3 + np.random.uniform(size=10000)
x1 = 3.0*x0 + 2.0*x2 + np.random.uniform(size=10000)
x5 = 4.0*x0 + np.random.uniform(size=10000)
x4 = 8.0*x0 - 1.0*x2 + np.random.uniform(size=10000)
X = pd.DataFrame(np.array([x0, x1, x2, x3, x4, x5]).T ,columns=['x0', 'x1', 'x2', 'x3', 'x4', 'x5'])
X.head()

m = np.array([[0.0, 0.0, 0.0, 3.0, 0.0, 0.0],
              [3.0, 0.0, 2.0, 0.0, 0.0, 0.0],
              [0.0, 0.0, 0.0, 6.0, 0.0, 0.0],
              [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
              [8.0, 0.0,-1.0, 0.0, 0.0, 0.0],
              [4.0, 0.0, 0.0, 0.0, 0.0, 0.0]])

make_dot(m)

"""
Explanation: Test data
We create test data consisting of 6 variables.
End of explanation
"""


prior_knowledge = make_prior_knowledge(
    n_variables=6,
    sink_variables=[0, 1, 4],
)
print(prior_knowledge)

# Draw a graph of prior knowledge
make_prior_knowledge_graph(prior_knowledge)

"""
Explanation: Make Prior Knowledge Matrix
We create prior knowledge so that x0, x1 and x4 are sink variables.
The elements of prior knowledge matrix are defined as follows:
* 0 : :math:x_i does not have a directed path to :math:x_j
* 1 : :math:x_i has a directed path to :math:x_j
* -1 : No prior knowledge is available to know if either of the two cases above (0 or 1) is true.
End of explanation
"""


model = lingam.DirectLiNGAM(prior_knowledge=prior_knowledge)
model.fit(X)
print(model.causal_order_)
print(model.adjacency_matrix_)

"""
Explanation: Causal Discovery
To run causal discovery using prior knowledge, we create a DirectLiNGAM object with the prior knowledge matrix.
End of explanation
"""


make_dot(model.adjacency_matrix_)

"""
Explanation: We can see that x0, x1, and x4 are output as sink variables, as specified in the prior knowledge.
End of explanation
"""


prior_knowledge = make_prior_knowledge(
    n_variables=6,
    exogenous_variables=[0],
)

model = lingam.DirectLiNGAM(prior_knowledge=prior_knowledge)
model.fit(X)

make_dot(model.adjacency_matrix_)

"""
Explanation: Next, let's specify the prior knowledge so that x0 is an exogenous variable.
End of explanation
"""


