from dowhy import CausalModel
import dowhy.datasets
import pandas as pd
import numpy as np

# Config dict to set the logging level
import logging.config
DEFAULT_LOGGING = {
    'version': 1,
    'disable_existing_loggers': False,
    'loggers': {
        '': {
            'level': 'WARN',
        },
    }
}

logging.config.dictConfig(DEFAULT_LOGGING)

"""
Explanation: A Simple Example on Creating a Custom Refutation Using User-Defined Outcome Functions
In this experiment, we define a linear dataset. In order to find the coefficients, we make use of the linear regression estimator. In order to test the effectiveness of the linear estimator, we now replace the outcome value with a dummy produced with the help of a linear expression based on the value of the confounders. This effectively means that the effect of the treatment on the outcome should be zero. This is exactly, what we should expect from the results of the refuter.
Insert Dependencies
End of explanation
"""


# Value of the coefficient [BETA]
BETA = 10
# Number of Common Causes
NUM_COMMON_CAUSES = 2
# Number of Instruments
NUM_INSTRUMENTS = 1
# Number of Samples
NUM_SAMPLES = 100000
# Treatment is Binary
TREATMENT_IS_BINARY = False
data = dowhy.datasets.linear_dataset(beta=BETA,
                                 num_common_causes=NUM_COMMON_CAUSES,
                                 num_instruments=NUM_INSTRUMENTS,
                                 num_samples=NUM_SAMPLES,
                                 treatment_is_binary=TREATMENT_IS_BINARY)
data['df'].head()

"""
Explanation: Create the Dataset
You can change the values of the hyper params to see how the effects change, as each parameter changes
Variable Guide:
| Variable Name   | Data Type |  Interpretation    |
|-----------------|-----------|--------------------|
|   $Z_i$         |  float    | Insrument Variable |
|   $W_i$         |  float    | Confounder         |
|   $v_0$         |  float    | Treatment          |
|    $y$          |  float    | Outcome            |
End of explanation
"""


model = CausalModel(
    data = data['df'],
    treatment = data['treatment_name'],
    outcome = data['outcome_name'],
    graph = data['gml_graph'],
    instruments = data['instrument_names']
)

model.view_model()

"""
Explanation: Creating the Causal Model
End of explanation
"""


identified_estimand = model.identify_effect(proceed_when_unidentifiable=True)
print(identified_estimand)

"""
Explanation: In the above figure, we have a causal graph that shows the relationships between the treatment, outcome, confounders and the instrument variable.
- The Confounders $W_0$ and $W_1$ affect both the treatment and the outcome
- The instrument variable $Z_0$ is able to effect the outcome $y$ through the treatment $x$
Identify the Estimand
End of explanation
"""


causal_estimate = model.estimate_effect( identified_estimand,
                                       method_name="iv.instrumental_variable",
                                       method_params={'iv_instrument_name':'Z0'}
                                       )
print(causal_estimate)

"""
Explanation: Estimating the Effect
End of explanation
"""


ref = model.refute_estimate(identified_estimand,
                           causal_estimate,
                           method_name="dummy_outcome_refuter"
                           )
print(ref[0])

"""
Explanation: Refuting the Estimate
Using a Randomly Generated Outcome
End of explanation
"""


coefficients = np.array([1,2])
bias = 3
def linear_gen(df):
    y_new = np.dot(df[['W0','W1']].values,coefficients) + 3
    return y_new

"""
Explanation: The result shows that the treatment does not lead to the outcome. The estimated effect is a value that tends to zero, which matches our expectation. This shows that if we replace the outcome by randomly generated data, the estimator correctly predicts that the influence if treatment is Zero.
Using a Function that Generates the Outcome from the Confounders
Let us define a simple function that generates the outcome as a linear function of the confounders.
End of explanation
"""


ref = model.refute_estimate(identified_estimand,
                           causal_estimate,
                           method_name="dummy_outcome_refuter",
                           outcome_function=linear_gen
                           )

print(ref[0])

"""
Explanation: The basic expression is of the form
$y_{new} = \beta_0W_0 + \beta_1W_1 + \gamma_0$
where,
$\beta_0=1$, $\beta_1=2$ and $\gamma_0=3$
End of explanation
"""


