import numpy as np
import pandas as pd
import pymc3 as pm
import arviz as az
import seaborn as sb
from platform import python_version
from causalgraphicalmodels import CausalGraphicalModel as CGM


# PRINT PACKAGE VERSIONS FOR REPRODUCIBILITY
def print_version(pkg):
    print(f'{pkg.__name__} version: {pkg.__version__}')

print(f'Python version: {python_version()}')
for pkg in [np, pd, pm, az, sb]:
    print_version(pkg)

"""
Explanation: "Correlation is not causation" goes the saying, cautioning against spurious correlation. Causally related phenomena, however, do correlate. How then can one differentiate spurious correlation from causal correlation, when confronted with a modeling problem? The following recommendations and simple simulation like the one making the subject of this post go a long way.

use causal DAGs - Directed Acyclic Graphs - to represent the current understanding of how the various pieces of data might be causally influencing each other (more on DAGs below.)
be aware of at least the four basic confounding mechanisms that lead to spurious correlation; namely, the Collider, the Pipe, the Fork, the Descendent.
be aware that spurious correlation through confounding can lead to <u>better</u> model performance that are ultimately misguiding and that can lead further inquiry astray.

<!-- TEASER_END -->
Here, I use a simple confounding-by-collider example to illustrate these points. The inspiration from this post comes from a tweet dated Jul 12, 2019, by Richard McElreath. 
First, import the required Python libraries...
End of explanation
"""


np.random.seed(42)
N = 1000
X = np.random.normal(size=N) # Gaussian centered at 0, with σ=1
Y = np.random.normal(loc=X, size=N) # Gaussian centered at X, with σ=1
W = np.random.normal(loc=X+Y, size=N) # Gaussian centered at X+Y with σ=1

"""
Explanation: <br><br>
As a simple example, we assume three measurable phenomena, identified as \( X \), \( Y \), and \( W \), for which the causal relationships are:

\( X \) influences \( Y \)
both \( X \) and \( Y \) influence \( W \)

Note that causal relationships are probabilistic, meaning that a change in a causal phenomenon will alter the probability in which the resulting effect manifests itself. Simplistically, this can be simulated as follows:
End of explanation
"""


# instantiate the causal graph object...
cgm = CGM(nodes=['X', 'Y', 'W'], 
          edges=[('X', 'Y'), ('X', 'W'), ('Y', 'W')])
#... and plot the DAG
cgm.draw()

"""
Explanation: \( \rightarrow \) Thecausal structure implied by the simulation above can be represented as follows:
End of explanation
"""


df = pd.DataFrame(np.c_[X, Y, W], columns=['X', 'Y', 'W'])

sb.pairplot(df);

"""
Explanation: The graph above is a causal Directed Acyclic Graph or DAG, where nodes represent (ideally but not necessarily) measurable phenomena, and edges represent directional causal links; a parent node is a cause, a child node is an effect. The rules in drawing DAGs are: (1) only unidirectional arrows are allowed (an effect cannot be a cause of its cause) and (2) loops are not allowed (an extension of the first rule to multiple nodes.)<br><br>
A statistical model will not care about the directionality of the causal relationships between the variables. Given \( Y = f(X) \) and \( X = g(Y) \) the link functions say nothing about which of these is the cause, which the effect. In a regression setting, it is easy to be seduced by the thought that the target variable is the effect and the predictor variable(s) the cause(s). A DAG is a useful tool for staying grounded.<br><br>
The researcher wants to predict \( Y \) from whatever data is available using a linear model.
* (s)he collects data on \( Y \), \( X \) and \( W \) without really understanding/simply ignoring causal relationships 
* (s)he proposes two models; \( M1 \) and \( M2 \) where:
    * M1: \( Y \sim f(X) \)
    * M2: \( Y \sim f(X, W) \)
* To avoid complications, \( f \) represents linear relationships such that
    * M1: predicts \( Y \) from \( X \) only  \( \Rightarrow Y \sim \beta X \)
    * M2: predicts \( Y \) from both \( X \) and \( W \ \Rightarrow\ Y \sim \beta X + \gamma W \)
<u>NOTE THAT:</u>

M1 is an UN-CONFOUNDED MODEL, where a correlation is a symptom of a real causal link.
M2 is a CONFOUNDED MODEL, where correlation is spurious, and the suspected causal link is actually not there.

\( \rightarrow \) Here's a quick illustration of the pairwise relationship for \( X,\ Y,\ W \).
End of explanation
"""


# unconfounded model
with pm.Model() as M1:
    β = pm.Normal('β', mu=0, sd=1)
    μ = β * X
    σ = pm.Exponential('σ', 1)
    Y_lik = pm.Normal('Y_lik', mu=μ, sd=σ, observed=Y)
    trace_1 = pm.sample(tune=1000)

"""
Explanation: The above does not offer any hint at what the causal relationships might be; unlike statistical inference, causal inference is not a reductionist exercise and requires some thinking, some DAG drawing, and  writing and fitting several models to tease out the correct underlying relationship. More on that in a later post.<br><br>
Below, I implement the first model \( M1 \), predicting \( Y \) on \( X \) only. I use a Bayesian framework, the details of which are not important here. Suffice it to say that the goal is to get probability distributions for the values of model parameters that measure the association between predictor(s) and target variables, conditioned on the model and the data at hand;  \( \beta \) in \( M1 \) and \( \beta \) and \( \gamma \) in \( M2 \).
End of explanation
"""


az.plot_posterior(trace_1, var_names=['β']);

"""
Explanation: \( \rightarrow \) Plot the marginal posteriors of \( M1 \) parameters
End of explanation
"""


# confounded model
with pm.Model() as M2:
    β = pm.Normal('β', mu=0, sd=1)
    γ = pm.Normal('γ', mu=0, sd=1)
    μ = β * X + γ * W
    σ = pm.Exponential('σ', 1)
    Y_lik = pm.Normal('Y_lik', mu=μ, sd=σ, observed=Y)
    trace_2 = pm.sample(tune=1000)

"""
Explanation: Recall that \( Y \) was generated as a Gaussian centered on \( X \), implying on average a  1-1 relationship. The plot above shows the model was able to recover this with \( \beta \)'s posterior mean equal to 1.0.
Below, the confounded model, predicting \( Y \) from both \( X \) and \( W \) is spelled out and sampled from.
End of explanation
"""


az.plot_posterior(trace_2, var_names=['β', 'γ']);

"""
Explanation: \( \rightarrow \) Plot the marginal posteriors of \( M2 \) parameters
End of explanation
"""


az.compare({'M2': trace_2, 'M1': trace_1})

"""
Explanation: Conditioning on W has nullified the association between \( X \) and \( Y \) \( \rightarrow \). The plot above shows the posterior of the \( \beta \) parameter, which measures the association of X with Y  in both \( M1 \) and \( M2 \), is practically centered on '0', suggesting the absence of a significant relationship between \( X \) and \( Y \) in \( M2 \). And yet the relationship is known to exist (cf.) data generation process above.
<br><br>
Will the causally accurate model be more performant?
<br><br>
Typically, ranking predictive skill of Bayesian models is done by comparing each model's WAIC (Watanabe-Akaike Information Criterion) score, where a lower score corresponds to an expectation of better prediction of model data (More on model comparison/selection at a later time.)
End of explanation
"""


