import numpy as np
import pandas as pd
import graphviz
import lingam

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

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

"""
Explanation: Causal Effect
Import and settings
In this example, we need to import numpy, pandas, and graphviz in addition to lingam.
End of explanation
"""


def make_graph(adjacency_matrix, labels=None):
    idx = np.abs(adjacency_matrix) > 0.01
    dirs = np.where(idx)
    d = graphviz.Digraph(engine='dot')
    names = labels if labels else [f'x{i}' for i in range(len(adjacency_matrix))]
    for to, from_, coef in zip(dirs[0], dirs[1], adjacency_matrix[idx]):
        d.edge(names[from_], names[to], label=f'{coef:.2f}')
    return d

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


X = pd.read_csv('http://archive.ics.uci.edu/ml/machine-learning-databases/auto-mpg/auto-mpg.data-original',
                   delim_whitespace=True, header=None,
                   names = ['mpg', 'cylinders', 'displacement',
                            'horsepower', 'weight', 'acceleration',
                            'model year', 'origin', 'car name'])
X.dropna(inplace=True)
X.drop(['model year', 'origin', 'car name'], axis=1, inplace=True)
print(X.shape)
X.head()

"""
Explanation: Test data
We use 'Auto MPG Data Set' (http://archive.ics.uci.edu/ml/datasets/Auto+MPG)
End of explanation
"""


model = lingam.DirectLiNGAM()
model.fit(X)
labels = [f'{i}. {col}' for i, col in enumerate(X.columns)]
make_graph(model.adjacency_matrix_, labels)

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


from sklearn.linear_model import LassoCV

target = 0 # mpg
features = [i for i in range(X.shape[1]) if i != target]
reg = LassoCV(cv=5, random_state=0)
reg.fit(X.iloc[:, features], X.iloc[:, target])

"""
Explanation: Prediction Model
We create the linear regression model.
End of explanation
"""


ce = lingam.CausalEffect(model)
effects = ce.estimate_effects_on_prediction(X, target, reg)

df_effects = pd.DataFrame()
df_effects['feature'] = X.columns
df_effects['effect_plus'] = effects[:, 0]
df_effects['effect_minus'] = effects[:, 1]
df_effects

max_index = np.unravel_index(np.argmax(effects), effects.shape)
print(X.columns[max_index[0]])

"""
Explanation: Identification of Feature with Greatest Causal Influence on Prediction
To identify of the feature having the greatest intervention effect on the prediction, we create a CausalEffect object and call the estimate_effects_on_prediction method.
End of explanation
"""


# mpg = 15
c = ce.estimate_optimal_intervention(X, target, reg, 1, 15)
print(f'Optimal intervention: {c:.3f}')

# mpg = 21
c = ce.estimate_optimal_intervention(X, target, reg, 1, 21)
print(f'Optimal intervention: {c:.3f}')

# mpg = 30
c = ce.estimate_optimal_intervention(X, target, reg, 1, 30)
print(f'Optimal intervention: {c:.3f}')

"""
Explanation: Estimation of Optimal Intervention
To estimate of the intervention such that the expectation of the prediction of the post-intervention observations is equal or close to a specified value, we use estimate_optimal_intervention method of CausalEffect.
End of explanation
"""


