from dowhy.datasets import linear_dataset
from dowhy import CausalModel
import econml

# 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)
# Disabling warnings output
import warnings
from sklearn.exceptions import DataConversionWarning
warnings.filterwarnings(action='ignore', category=DataConversionWarning)

"""
Explanation: Iterating over multiple refutation tests
The objective of this notebook is to compare the ability of refuters to detect the problems in a given set of estimators.
Note:
This notebook makes use of the optional dependencies:
- pygraphviz
- econml
Import Dependencies
End of explanation
"""


inspect_datasets = True
inspect_models = True
inspect_identified_estimands = True
inspect_estimates = True
inspect_refutations = True

"""
Explanation: Inspection Parameters
These parameters give us the option of inspecting the intermediate steps to sanity check the steps performed
End of explanation
"""


from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LassoCV
from sklearn.ensemble import GradientBoostingRegressor

estimator_list = ["backdoor.propensity_score_matching", "backdoor.propensity_score_weighting", "backdoor.econml.dml.DML"]
method_params= [ None, None, {"init_params":{'model_y':GradientBoostingRegressor(),
                                                              'model_t': GradientBoostingRegressor(),
                                                              "model_final":LassoCV(fit_intercept=False), 
                                                              'featurizer':PolynomialFeatures(degree=1, include_bias=False)},
                                               "fit_params":{}} ]

"""
Explanation: Estimator List
We pass a list of strings, corresponding to the estimators of interest
End of explanation
"""


refuter_list = ["bootstrap_refuter", "data_subset_refuter"]

"""
Explanation: Refuter List
A list of strings, corresponding to each refuter we wish to run
End of explanation
"""


# Parameters for creating the Dataset
TREATMENT_IS_BINARY = True
BETA = 10
NUM_SAMPLES = 5000
NUM_CONFOUNDERS = 5
NUM_INSTRUMENTS = 3
NUM_EFFECT_MODIFIERS = 2

# Creating a Linear Dataset with the given parameters
linear_data = linear_dataset(
            beta = BETA,
            num_common_causes = NUM_CONFOUNDERS,
            num_instruments = NUM_INSTRUMENTS,
            num_effect_modifiers = NUM_EFFECT_MODIFIERS,
            num_samples = NUM_SAMPLES,
            treatment_is_binary = True
        )
# Other datasets come here 


# Append them together in an array
datasets = [linear_data]


"""
Explanation: Create the Datasets
End of explanation
"""


dataset_num = 1
if inspect_datasets is True:
    for data in datasets:
        print("####### Dataset {}###########################################################################################".format(dataset_num))
        print(data['df'].head())
        print("#############################################################################################################")
        dataset_num += 1

"""
Explanation: Inspect Data
End of explanation
"""


models = []
for data in datasets:
    model = CausalModel(
                data = data['df'],
                treatment = data['treatment_name'],
                outcome = data['outcome_name'],
                graph = data['gml_graph']
            )
    models.append(model)

"""
Explanation: Create the CausalModels
End of explanation
"""


model.view_model()
from IPython.display import Image, display
display(Image(filename="causal_model.png"))

"""
Explanation: View Models
End of explanation
"""


identified_estimands = []
for model in models:
    identified_estimand = model.identify_effect(proceed_when_unidentifiable=True)
    identified_estimands.append(identified_estimand)

"""
Explanation: Identify Effect
End of explanation
"""


estimand_count = 1
for estimand in identified_estimands:
    print("####### Identified Estimand {}#####################################################################################".format(estimand_count))
    print(estimand)
    print("###################################################################################################################")
    estimand_count += 1

"""
Explanation: Identified Estimands
End of explanation
"""


estimate_list = []
for i in range(len(identified_estimands)):
    for j in range(len(estimator_list)):
        estimate = model.estimate_effect(
                        identified_estimands[i],
                        method_name=estimator_list[j],
                        method_params=method_params[j]
                  )
        estimate_list.append(estimate)

"""
Explanation: Estimate Effect
End of explanation
"""


estimand_count = 1
if inspect_estimates is True:
    for estimand in estimate_list:
        print("####### Estimand {}#######################################################################################".format(estimand_count))
        print("*** Class Name ***")
        print()
        print(estimand.params['estimator_class'])
        print()
        print(estimand)
        print("########################################################################################################")
        print()
        estimand_count += 1
    

"""
Explanation: Estimate Values
End of explanation
"""


refutation_list = []
for estimand in identified_estimands:
    for estimate in estimate_list: 
        for refuter in refuter_list:
            ref = model.refute_estimate(estimand, estimate,method_name=refuter)
            refutation_list.append(ref)

"""
Explanation: Refute Estimate
End of explanation
"""


refuter_count = 1
if inspect_refutations is True:
    for refutation in refutation_list:
        print("####### Refutation {}#######################################################################################".format(refuter_count))
        print("*** Class Name ***")
        print()
        print(refutation.refutation_type)
        print()
        print(refutation)
        print("########################################################################################################")
        print()
        refuter_count += 1

"""
Explanation: Refutation Values
End of explanation
"""


