import pytest

from dowhy.causal_estimators.linear_regression_estimator import LinearRegressionEstimator
from .base import TestEstimator

@pytest.mark.usefixtures("fixed_seed")
class TestLinearRegressionEstimator(object):
    @pytest.mark.parametrize(
        [
            "error_tolerance", "Estimator", "num_common_causes", "num_instruments",
            "num_effect_modifiers", "num_treatments", "treatment_is_binary",
            "treatment_is_category", "outcome_is_binary"
        ],
        [
            (0.1, LinearRegressionEstimator, [0,1], [0,1], [0,], [1, 2],
            [True,], [False,], [False,]),
            (0.1, LinearRegressionEstimator, [0,1], [0,1], [0,], [1, 2],
            [False,], [True,], [False,]),
            (0.1, LinearRegressionEstimator, [0,1], [0,1], [0,], [1, 2],
            [False,], [False,], [False,]),
        ]
    )
    def test_average_treatment_effect(self, error_tolerance, Estimator,
            num_common_causes, num_instruments, num_effect_modifiers,
            num_treatments, treatment_is_binary, treatment_is_category,
            outcome_is_binary
            ):
        estimator_tester = TestEstimator(error_tolerance, Estimator)
        estimator_tester.average_treatment_effect_testsuite(
                num_common_causes=num_common_causes,
                num_instruments = num_instruments,
                num_effect_modifiers = num_effect_modifiers,
                num_treatments=num_treatments,
                treatment_is_binary=treatment_is_binary,
                treatment_is_category=treatment_is_category,
                outcome_is_binary=outcome_is_binary,
                confidence_intervals=[True,],
                test_significance=[True,],
                method_params={
                    'num_simulations': 10,
                    'num_null_simulations': 10
                    }
                )

