import re

import warnings

import anndata

import dowhy

import dowhy.datasets

import matplotlib.pyplot as plt

import numpy as np

import ehrapy as ep

warnings.filterwarnings("ignore")

class TestCausal:

    def setup_method(self):

        linear_data = dowhy.datasets.linear_dataset(

            beta=10,

            num_common_causes=5,

            num_instruments=2,

            num_samples=1000,

            treatment_is_binary=True,

        )

        self.linear_data = anndata.AnnData(linear_data["df"].astype(np.float32))

        self.linear_graph = linear_data["gml_graph"]

        self.outcome_name = "y"

        self.treatment_name = "v0"

    def test_dowhy_linear_dataset(self):

        estimate, refute_results = ep.tl.causal_inference(

            adata=self.linear_data,

            graph=self.linear_graph,

            treatment=self.treatment_name,

            outcome=self.outcome_name,

            estimation_method="backdoor.linear_regression",

            return_as="estimate+refute",

        )

        assert isinstance(refute_results, dict)

        assert len(refute_results) == 6

        assert isinstance(estimate, dowhy.causal_estimator.CausalEstimate)

        assert np.isclose(

            np.round(refute_results["Refute: Add a random common cause"]["test_significance"], 3), 10.002, atol=0.005

        )

        assert np.isclose(

            np.round(refute_results["Refute: Add a random common cause"]["test_significance"], 3), 10.002, atol=0.005

        )

    def test_plot_causal_effect(self):

        estimate = ep.tl.causal_inference(

            adata=self.linear_data,

            graph=self.linear_graph,

            treatment=self.treatment_name,

            outcome=self.outcome_name,

            estimation_method="backdoor.linear_regression",

            return_as="estimate",

            show_graph=False,

            show_refute_plots=False,

        )

        ax = ep.pl.causal_effect(estimate)

        assert isinstance(ax, plt.Axes)

        legend = ax.get_legend()

        assert len(legend.get_texts()) == 2  # Check the number of legend labels

        assert legend.get_texts()[0].get_text() == "Observed data"

        assert legend.get_texts()[1].get_text() == "Causal variation"

        assert re.search(r"(9\.99\d+|10\.0)", str(ax.get_title()))