# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""A simple Structural Causal Model (SCM) specified with stochastic functions.

The building block is a stochastic function, which is a potentially random
mapping from values of parents to values of a variable. These functions
generalize:
  -sampling from conditional probability distributions
  -structural functions in SCMs.

Stochastic functions (stoc_fns) are either parent-less, in which case they only
output a random sample, or are specified by the values of their parents with
the option of additional randomness (e.g. they may be a linear function of the
parents with Gaussian noise). In this way, these function differs form the
structural functions found in SCMs which are always deterministic.

If the stochastic function has no parents, then it must take exactly one
integer-values argument, n_samples. For example, a 3-D standard normal would
be generated by
    lamdda n_samples: np.random.normal(size=(n_samples,3 )).
All samples are always of the shape (n_samples,) or (n_samples, dimension).

If the stochastic function has parents, then the parents can be specified in
two ways:
1. Using the signature; e.g.
    def fn(x1, x2):
      return x1 + x2
2. Using the StocFnRecipe class, which specifies the function with kwargs
   and the parents as metadata, e.g.
    StocFnRecipe(lambda **kwargs; kwargs['x1'] + kwargs['x2'], ['x1', 'x2'])

The second method is helpful when you want to create procedural functions,
e.g. when the graph is generated first so the parents of a random variable
are not known apriori.

All data are assumed to be np.ndarrays, and the number of samples returned by
a stoc_fn with parents is infered from the size of the parents which is why no
n_samples argument is provided.

Finally, a SCM is specified with a dictionary of stochastic functions.

Example: We have the linear Gaussian SCM with X<-Z->Y and X->Y. We would call
  stoc_fns = {
      'Z': lambda n_samples: np.random.normal(size=n_samples)},
      'X': lambda Z: 2*Z + np.random.normal(size=Z.shape)},
      'Y': lambda X, Z: Z - 2 * X,
  }
  scm = SCM(stoc_fns)
Alternatively, we could use:
  def f_y(**kwargs):
    return kwargs['Z'] - 2 * kwargs['X']
  stoc_fns = {
      'Z': lambda n_samples: np.random.normal(size=n_samples)},
      'X': lambda Z: 2*Z + np.random.normal(size=X.shape)},
      'Y': StocFnRecipe(f_y, ['Z'' 'X'])
  }
  scm = SCM(stoc_fns).

  N.B. If you are defining functions with mutable parameters, e.g. for some
  mutable variable a
    def f_y(**kwargs):
      return kwargs['Z'] - a * kwargs['X']
  you will need to either ensure that the value of a never changes (by making
  a copy specifically for this function) or by definit the function in a
  closure:
    def f_y(a):
      def f(**kwargs):
        return kwargs['Z'] - a * kwargs['X']
      return f
  which ensures that a local copy of a in created and bound to the function f_y.
"""

import dataclasses
import inspect
from typing import Any, Callable, Mapping, Optional, Union, List

import networkx as nx
import numpy as np

StocFn = Callable[..., Any]  # [[np.ndarray, ..., np.ndarray], np.ndarray]

N_SAMPLES_NAME = 'n_samples'  # Constant for functions that do not have parents.


@dataclasses.dataclass
class StocFnRecipe:
  """Provides metadata needed to construct a stoc_fn."""
  stoc_fn: StocFn
  stoc_fn_inputs: Optional[List[str]] = None


class SCM:
  """Implements a simple SCM using stochastic functions.

  Attributes:
    var_names: A list of variable names.
    parents: Mapping[str, List[str]]
      parents[var] is a list of all names of all the parents of the variable
      var. If var has no parents, parents[var] is an empty list.
    stoc_fns: a dictionary of Callable objects.
      With samples: Mapping[str, np.ndarray] is a dictionary of values,
        stoc_fns[var](**parent_dict),
      where parent_dict is a dictionary of the parent values, will provide the
      values of the variable var. Recall that this function need not be
      deterministic.
  """

  def __init__(self,
               stoc_fns: Mapping[str, Union[StocFnRecipe, StocFn]],
               ):
    self._stoc_fns = {}
    self._stoc_fn_inputs = {}

    for var, stoc_fn in stoc_fns.items():
      if callable(stoc_fn):
        sig = inspect.signature(stoc_fn)
        self._stoc_fn_inputs[var] = list(sig.parameters.keys())
        self._stoc_fns[var] = stoc_fn
      elif isinstance(stoc_fn, StocFnRecipe):
        if stoc_fn.stoc_fn_inputs is not None:
          self._stoc_fn_inputs[var] = stoc_fn.stoc_fn_inputs
        else:
          self._stoc_fn_inputs[var] = [N_SAMPLES_NAME]
        self._stoc_fns[var] = stoc_fn.stoc_fn
      else:
        raise ValueError(
            'Must be a callable with a signature specifying the parents or a'
            ' StocFnRecipe.'
        )
    self._var_names_ordered = self._topological_sort(self._stoc_fn_inputs)

  def sample(self,
             n_samples: int,
             intervention: Optional[Union[Mapping[str, np.ndarray],
                                          Mapping[str, int],
                                          Mapping[str, float],
                                          ]] = None,
             ) -> Mapping[str, np.ndarray]:
    """Generates a sample from the SCM.

    Args:
      n_samples: The number of samples.
      intervention (optional):
        For each key, the corresponding variable is set to the value instead of
        being sampled from the SCM.

    Returns:
      Dictionary with {var_name: sample of var_name} for all var_names
    """
    samples = {N_SAMPLES_NAME: np.array([n_samples], dtype=int)}
    if intervention is None:
      intervention = {}
    for var in self._var_names_ordered:
      if var in intervention:
        val = np.array(intervention[var])
        if val.size == 1:
          samples[var] = np.repeat(val.flatten(), n_samples, axis=0)
        else:
          samples[var] = np.repeat(val.reshape((1, -1)), n_samples, axis=0)
      else:
        parent_samples = {v: samples[v] for v in self._stoc_fn_inputs[var]}
        samples[var] = self._stoc_fns[var](**parent_samples)

    samples.pop(N_SAMPLES_NAME)  # Remove integer inputs.
    return samples

  def _topological_sort(
      self, stoc_fn_inputs: Mapping[str, List[str]]
  ) -> List[str]:
    """Topologically sorts the nodes."""
    edges = []
    nodes = list(stoc_fn_inputs.keys())
    for var, curr_parent_names in stoc_fn_inputs.items():
      edges.extend([(p, var) for p in curr_parent_names])

    dag = nx.DiGraph()
    dag.add_nodes_from(nodes)
    dag.add_edges_from(edges)
    return [var for var in nx.topological_sort(dag) if var != N_SAMPLES_NAME]

  def draw(self):
    """Draws the adjacency graph of the SCM."""
    edges = []
    nodes = list(self.var_names)
    for var, curr_parent_names in self.parents.items():
      edges.extend([(p, var) for p in curr_parent_names])

    dag = nx.DiGraph()
    dag.add_nodes_from(nodes)
    dag.add_edges_from(edges)
    nx.draw(dag, with_labels=True)

  @property
  def parents(self) -> Mapping[str, List[str]]:
    return {
        v: [] if item == [N_SAMPLES_NAME] else item
        for v, item in self._stoc_fn_inputs.items()
    }

  @property
  def stoc_fns(self) -> Mapping[str, Callable[..., any]]:
    return self._stoc_fns

  @property
  def var_names(self) -> List[str]:
    """Returns all the variable names in topological order."""
    return self._var_names_ordered
