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

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

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

"""
Explanation: Total Effect
Import and settings
In this example, we need to import numpy, pandas, and graphviz in addition to lingam.
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
"""


model = lingam.DirectLiNGAM()
model.fit(X)
make_dot(model.adjacency_matrix_)

"""
Explanation: Causal Discovery
To run causal discovery, we create a DirectLiNGAM object and call the fit method.
End of explanation
"""


# x3 --> x1
te = model.estimate_total_effect(X, 3, 1)
print(f'total effect: {te:.3f}')

# x2 --> x1
te = model.estimate_total_effect(X, 2, 1)
print('total effect: {:.3f}'.format(te))

# x1 --> x3 (Warning!)
# The causal order of the destination variable is earlier than the source variable.
te = model.estimate_total_effect(X, 1, 3)
print('total effect: {:.3f}'.format(te))

"""
Explanation: Total Effect
To estimate the total effect, we can call estimate_total_effect method with lingam model object.
End of explanation
"""


