# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/01_data.scm.ipynb.

# %% ../../nbs/01_data.scm.ipynb 3
from __future__ import annotations
from ..import_essentials import *
from . import TabularDataModule, TabularDataModuleConfigs
import networkx as nx
try: 
    from causalgraphicalmodels import CausalGraphicalModel, StructuralCausalModel
except: 
    raise ModuleNotFoundError("`causalgraphicalmodels` is not installed. "
        "Try to install via `pip install git+https://github.com/BirkhoffG/causalgraphicalmodels.git`.")

# %% auto 0
__all__ = ['BaseDistribution', 'Normal', 'MixtureOfGaussians', 'sanity_3_lin', 'CausalModel']

# %% ../../nbs/01_data.scm.ipynb 5
class BaseDistribution(ABC):
    """Base class for all distributions."""
    def __init__(self, name: str):
        self.name = name

    @abstractmethod
    def sample(self, size: int):
        pass

    @abstractmethod
    def pdf(self, value):
        pass

# %% ../../nbs/01_data.scm.ipynb 6
class Normal(BaseDistribution):
    """Normal distribution."""
    def __init__(self, mean: int | float, var: int | float):
        super().__init__(f"Normal (mean={mean}, var={var})")
        self.mean = mean
        self.var = var

    def sample(self, size=1):
        tmp = [np.random.normal(self.mean, np.sqrt(self.var)) for _ in range(size)]
        return tmp[0] if size == 1 else tmp

    def pdf(self, value):
        return scipy.stats.norm(self.mean, self.var).pdf(value)


class MixtureOfGaussians(BaseDistribution):
    """Mixture of Gaussians distribution."""
    def __init__(self, probs, means, vars):
        if not np.sum(probs) == 1:
            raise ValueError("Mixture probabilities must sum to 1.")
        if not len(probs) == len(means) == len(vars):
            raise ValueError("Length mismatch.")

        self.probs = probs
        self.means = means
        self.vars = vars
        super().__init__(
            f"MixtureOfGaussians (mean={probs}, means={means}, vars={vars}"
        )

    def sample(self, size=1):
        tmp = [
            np.random.normal(self.means[mixture_idx], np.sqrt(self.vars[mixture_idx]))
            for mixture_idx in np.random.choice(
                len(self.probs), size=size, p=self.probs
            )
        ]
        return tmp[0] if size == 1 else tmp

    def pdf(self, value):
        return np.sum(
            [
                prob * scipy.stats.norm(mean, var).pdf(value)
                for (prob, mean, var) in zip(self.probs, self.means, self.vars)
            ]
        )


# %% ../../nbs/01_data.scm.ipynb 8
def sanity_3_lin():
    structural_equations = {
        "x1": lambda n_samples: n_samples,
        "x2": lambda n_samples, x1: -x1 + n_samples,
        "x3": lambda n_samples, x1, x2: 0.5 * (0.1 * x1 + 0.5 * x2) + n_samples,
    }
    noises_distributions = {
        "u1": MixtureOfGaussians([0.5, 0.5], [-2, +1], [1.5, 1]),
        "u2": Normal(0, 1),
        "u3": Normal(0, 1),
    }
    continuous = list(structural_equations.keys()) + list(
        noises_distributions.keys()
    )
    categorical = []
    immutables = []

    return (
        structural_equations,
        noises_distributions,
        continuous,
        categorical,
        immutables,
    )

# %% ../../nbs/01_data.scm.ipynb 9
def _load_scm_equations(
    scm_class: str # Name of the structural causal model.
) -> Tuple[
    Dict[str, Callable], 
    Dict[str, Callable], 
    list, list, list]: # return scm equations
    
    scm_func_dict = {
        'sanity_3_lin': sanity_3_lin
    }
    if not scm_class in scm_func_dict:
        raise ValueError(f'{scm_class} should be one of {scm_func_dict.keys()}')
    return scm_func_dict[scm_class]()

# %% ../../nbs/01_data.scm.ipynb 11
class CausalModel:
    """Class with topological methods given a structural causal model."""
    def __init__(
        self,
        scm_class: str # Name of the structural causal model.
    ):
        (
            self._structural_equations,
            self._noise_distributions,
            self._continuous,
            self._categorical,
            self._immutables,
        ) = _load_scm_equations(scm_class)
        self._scm = StructuralCausalModel(self._structural_equations)
        self._cgm = self._scm.cgm

        self._endogenous = list(self._structural_equations.keys())
        self._exogenous = list(self._noise_distributions.keys())

        self._continuous_noise = list(set(self._continuous) - set(self._endogenous))
        self._categorical_noise = list(set(self._categorical) - set(self._endogenous))

        self._continuous = list(set(self._continuous) - set(self._exogenous))
        self._categorical = list(set(self._categorical) - set(self._exogenous))

    def get_topological_ordering(self, node_type: str = 'endogenous'):
        """Get topological order of nodes."""
        tmp = nx.topological_sort(self._cgm.dag)
        if node_type == 'endogenous':
            return tmp
        elif node_type == 'exogenous':
            return [f'u{node[1:]}' for node in tmp]
        else:
            raise ValueError(f'node_type should be one of ["endogenous", "exogenous"], '
                             f'but got {node_type}')
        
    def get_children(self, node: str):
        """Get children of a node."""
        return set(self._cgm.dag.successors(node))
    
    def get_parents(self, node: str, return_sorted: bool = True):
        """Get parents of a node."""
        tmp = set(self._cgm.dag.predecessors(node))
        return sorted(tmp) if return_sorted else tmp
    
    def get_ancestors(self, node: str, return_sorted: bool = True):
        """Get ancestors of a node."""
        tmp = nx.ancestors(self._cgm.dag, node)
        return sorted(tmp) if return_sorted else tmp
    
    def get_descendants(self, node: str, return_sorted: bool = True):
        """Get descendants of a node."""
        tmp = nx.descendants(self._cgm.dag, node)
        return sorted(tmp) if return_sorted else tmp
    
    def get_non_descendants(self, node: str):
        """Get non-descendants of a node."""
        return (
            set(nx.topological_sort(self._cgm.dag))
            .difference(self.get_descendants(node))
            .symmetric_difference(set([node]))
        )
    
    def visualize(self, dir_name: str = None):
        """Visualize the causal graph."""
        if dir_name:
            path = f'{dir_name}/_causal_graph'
            view_flag = False
        else:
            path = "_tmp/_causal_graph"
            view_flag = True
        self._cgm.draw().render(path, view=view_flag)

    @property
    def scm(self) -> StructuralCausalModel:
        return self._scm

    @property
    def cgm(self) -> CausalGraphicalModel:
        return self._cgm

    @property
    def scm_class(self) -> str:
        return self._scm_class

    @property
    def structural_equations(self) -> dict:
        return self._structural_equations

    @property
    def noise_distributions(self) -> dict:
        return self._noise_distributions

    @property
    def exogenous(self) -> List[str]:
        return self._exogenous

    @property
    def endogenous(self) -> List[str]:
        return self._endogenous

# %% ../../nbs/01_data.scm.ipynb 16
def _get_noise_string(node):
    if not node[0] == "x":
        raise ValueError
    return "u" + node[1:] # get node id

# %% ../../nbs/01_data.scm.ipynb 17
def _create_synthetic_data(
    scm: CausalModel, # Structural causal model 
    num_samples: int # Number of samples in the dataset.
) -> Union[pd.DataFrame, pd.DataFrame]: # Return exogenous and endogenous variables.
    """Generate synthetic data."""

    exogenous_nodes = scm.get_topological_ordering("exogenous")
    exogenous_variables = np.concatenate(
        [
            np.array(scm.noise_distributions[node].sample(num_samples)).reshape(-1, 1)
            for node in exogenous_nodes
        ],
        axis=1,
    )
    exogenous_variables = pd.DataFrame(
        exogenous_variables, columns=exogenous_nodes
    )

    endogenous_variables = exogenous_variables.copy()
    endogenous_variables = endogenous_variables.rename(
        columns=dict(
            zip(
                scm.get_topological_ordering("exogenous"),
                scm.get_topological_ordering("endogenous"),
            )
        )
    )
    # used later to make sure parents are populated when computing children
    endogenous_variables.loc[:] = np.nan
    for node in scm.get_topological_ordering("endogenous"):
        parents = scm.get_parents(node)
        if endogenous_variables.loc[:, list(parents)].isnull().values.any():
            raise ValueError(
                "parents in endogenous_variables should already be occupied"
            )
        endogenous_variables[node] = scm.structural_equations[node](
            exogenous_variables[_get_noise_string(node)],
            *[endogenous_variables[p] for p in parents],
        )

    # fix a hyperplane
    w = np.ones((endogenous_variables.shape[1], 1))
    # get the average scale of (w^T)*X, this depends on the scale of the data
    scale = 2.5 / np.mean(np.abs(np.dot(endogenous_variables, w)))
    predictions = 1 / (1 + np.exp(-scale * np.dot(endogenous_variables, w)))

    if not 0.20 < np.std(predictions) < 0.42:
        raise ValueError(f"std of labels is strange: {np.std(predictions)}")

    # sample labels from class probabilities in predictions
    uniform_rv = np.random.rand(endogenous_variables.shape[0], 1)
    labels = uniform_rv < predictions
    labels = pd.DataFrame(data=labels, columns=["label"])

    df_endogenous = pd.concat([endogenous_variables, labels], axis=1).astype("float64")
    df_exogenous = pd.concat([exogenous_variables], axis=1).astype("float64")
    return df_endogenous, df_exogenous
