# 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.

"""Utilities for building discreteSCM objects."""

import itertools
from typing import List, Mapping, Optional, Tuple

from cbug import discrete_scm
from cbug import scm
from cbug import stoc_fn_utils as sf
import numpy as np


def sample_discrete_additive_scm(
    num_variables: int,
    degree: float,
    num_parents: int,
    prob_outcome_child: float = .5,
    cov: float = 1,
    mean_bound: float = 1,
    support_size_min: int = 3,
    support_size_max: int = 6,
    alpha: float = 2,
    beta: float = 5,
    dir_alpha: float = .5,
    interactions: Optional[List[int]] = None,
    interaction_magnitude: Optional[float] = None,
) -> discrete_scm.DiscreteSCM:
  """Generates a discrete additive SCM randomly.

  All variables in the SCM are discrete except for the outcome variable, Y.
  The stoc_fn of the outcome variable is additive with Gaussian noise of
  covariance cov.

  This function samples an SCM with a graph generated by the Erdos-Renyi process
  with the specified degree and where all the random variables are discrete.
  The number of discinct values a variable can take, known as its support_size,
  is sampled i.i.d. for each variable unifomly betwen support_size_min and
  support_size_max. All variables with no
  parents are have a categorical distribution with a parameter vector sampled
  from a Dirichlet with parameter dir_alpha, and the conditional distributions
  are sampled from a normalized Beta distribution with parameters alpha, beta.

  The Y variable is randomly connected to num_parents parents, and Y is equal to
  a linear function of its parents, with randomly sampled parameters less than
  mean_bound and added Gaussian noise of covariance cov.

  All variables topoligically after Y are made children of Y with probability
  prob_y_child, and have the same discrete categorical distribution with a
  discretized version of Y.

  If interactions and interaction_magnitude are not None, then a non-linear
  component to the stoc_fn of Y is added. The interactions term is
  the sum of several interaction terms, as specified by the interactions
  argument. Each element in interactions will create a separate interaction term
  with size equal to the element's value, e.g. an interaction term with size 3
  will be x_1 * x_2 * x_3, where x_1, x_2, and x_3 are randomly chosen from
  y's parents.
  For example, if interactions = [2, 3], then the non-linear term could be
  x_2 * x_4 + x_0 * x_1 * x_5.

  Args:
    num_variables: The number of variables, excluding the output variable Y.
    degree: The average degree used to same on the Erdos-Renyi graph.
    num_parents: The number of parents of Y.
    prob_outcome_child: The probability of adding an edge from Y to any eligible
      variable.
    cov: covariance of Y's Gaussian noise.
    mean_bound: Maximum coefficient for Y's linear term.
    support_size_min: Lower bound on support size for each variable.
    support_size_max: Upper bound on support size for each variable.
    alpha: Y's values are sampled from a beta(alpha, beta) distribution.
    beta: Y's values are sampled from a beta(alpha, beta) distribution.
    dir_alpha: The dirichlet parameter used to sample the probabilities for
      categorical variables with no parents.
    interactions: A list of integers where each element will correspond to an
      interaction term between that many parents of Y. Only used when
      interaction_magnitude > 0.
    interaction_magnitude: The amount of model mispecification that breaks the
      additivity assumption in the outcome variable's stoc_fn.

  Returns:
    An SCM with a graph with the given properties.
  """
  cov_variables = ['X' + str(i) for i in range(num_variables)]
  outcome_variable = 'Y'

  # Create an extra column for Y.
  adj = np.zeros((num_variables, num_variables + 1))
  adj[:, :num_variables] = sample_dag_matrix(num_variables, degree)
  # The nodes are in reverse topological order; e.g. node 0 has no children,
  # node num_variables-1 has no parents.

  # Sample the parents of Y.
  outcome_parent_idx = np.random.choice(
      num_variables, size=num_parents, replace=False
  )
  outcome_parents = [cov_variables[i] for i in outcome_parent_idx]
  # Sample the children of Y. We need to ensure we add no cycles.
  max_idx = max(outcome_parent_idx)
  if max_idx < num_variables:
    # Randomly add Y as a parent to variables > Y, topologically.
    for idx in range(max_idx + 1, num_variables):
      if np.random.binomial(1, prob_outcome_child):
        adj[idx, num_variables] = 1

  # Sample the support sizes (the number of distinct values the variable can
  # take) for all variables.
  support_sizes = {
      var: np.random.randint(support_size_min, support_size_max + 1)
      for var in cov_variables
  }
  support_sizes['Y'] = int(np.ceil(mean_bound * len(outcome_parent_idx))) + 1

  # For each variable, randomly construct its stochastic function.
  stoc_fns = {}
  for i, var in enumerate(cov_variables):
    # Enumerate the parents.
    parents = []
    for j in range(num_variables):
      if adj[i, j]:
        parents.append(cov_variables[j])
    if adj[i, num_variables]:
      parents.append('Y')
    if not parents:
      probs = np.random.dirichlet([dir_alpha] * support_sizes[var])
      stoc_fns[var] = sf.create_categorical_stoc_fn(probs=probs)
      # stoc_fn has one input, n_samples, since it has no parents.
    else:
      # Sample probs.
      if len(parents) == 1:
        probs = np.ones(support_sizes[var]) / support_sizes[var]
      else:
        probs = generate_beta_cpd(support_sizes, parents, var)
      # Create the stoc_fn.
      stoc_fn = sf.create_categorical_conditional_stoc_fn(probs, parents)
      stoc_fns[var] = scm.StocFnRecipe(stoc_fn, parents)

  # Sample the linear part of the outcome stoc_fn.
  means = {}
  for var in outcome_parents:
    means[var] = mean_bound * np.random.beta(
        alpha, beta, size=support_sizes[var]
    )

  # If interaction is not present, we generate a linear function.
  if interaction_magnitude is None or interaction_magnitude == 0:
    # Find the best action and best value.
    outcome_expected_value_fn = sf.create_discrete_to_linear_stoc_fn(means)
    outcome_stoc_fn = sf.add_gaussian_noise_to_stoc_fn(
        outcome_expected_value_fn, cov
    )
    stoc_fns[outcome_variable] = scm.StocFnRecipe(
        outcome_stoc_fn, outcome_parents
    )
    interaction_mean_fn = None

  elif interaction_magnitude is not None and interaction_magnitude != 0:
    assert all(np.array(interactions) <= num_parents)
    # Each non-linear interaction term will be the product of more than one
    # of Y's parents. Here, we randomly choose between them.
    domains = [
        np.random.choice(outcome_parents, num, replace=False)
        for num in interactions
    ]

    # The stoc_fn of the outcome is formed by the sum of an interaction term,
    # a linear term, and Gaussian noise.
    interaction_mean_fn = sf.create_interaction_mean_stoc_fn(
        interaction_magnitude,
        domains,
    )
    linear_mean_fn = sf.create_discrete_to_linear_stoc_fn(means)
    outcome_expected_value_fn = sf.add_stoc_fns(
        [interaction_mean_fn, linear_mean_fn]
    )
    stoc_fns[outcome_variable] = scm.StocFnRecipe(
        sf.add_gaussian_noise_to_stoc_fn(outcome_expected_value_fn, cov),
        outcome_parents,
    )

  best_action, best_action_value = find_best_action_and_value(
      outcome_parents,
      support_sizes,
      means,
      interaction_mean_fn=interaction_mean_fn,
  )

  return discrete_scm.DiscreteSCM(
      stoc_fns=stoc_fns,
      support_sizes=support_sizes,
      outcome_variable='Y',
      best_action=best_action,
      best_action_value=best_action_value,
      outcome_expected_value_fn=outcome_expected_value_fn,
  )


def get_expected_value_of_outcome(
    model: discrete_scm.DiscreteSCM,
    action: Mapping[str, int],
    n: int = 1000,
) -> float:
  """Computes or estimates the value of an action.

  If the action specifies all the parents of the outcome variable, then
  the expected value can be read off from the model's outcome_expected_value_fn.

  Otherwise, the value is approximated.

  Args:
    model: The DiscreteSCM that generates the value.
    action: A Dictionary specifying the action for each variable in model.
    n: The sample size used to estimate the value if an exact value cannot be
      calculated.

  Returns:
    Either the exact expected value of the outcome variable or an approximation.
  """
  if action.keys() >= set(model.parents[model.outcome_variable]):
    return float(model.outcome_expected_value_fn(**action))
  else:  # We need to sample.
    return np.mean(model.sample(n, action)[model.outcome_variable])


def find_best_action_and_value(
    outcome_parents: List[str],
    support_sizes: Mapping[str, int],
    means: Mapping[str, np.ndarray],
    interaction_mean_fn: Optional[scm.StocFn] = None,
) -> Tuple[Mapping[str, int], float]:
  """Finds the best mean action for a linear + optional interaction fn.

  Args:
    outcome_parents: A list of the parents of the outcome variable
    support_sizes: A mapping from variable names to the support size.
    means: Dictionary of means for the value of every parent of the
      outcome_variable.
    interaction_mean_fn: A callable that accepts an action dictionary and
      returns a float equal to the mean of the interaction term for that action.

  Returns:
    The best action (as a dictionary) and its value.
  """
  if interaction_mean_fn is None:
    best_action = {var: np.argmax(var_mean) for var, var_mean in means.items()}
    best_action_value = np.sum(
        [np.max(var_mean) for var_mean in means.values()]
    )
    return (best_action, best_action_value)

  parent_values = {
      par: np.arange(support_sizes[par]) for par in outcome_parents
  }

  best_value = -np.inf
  # Loop through all values of the parents.
  for val in itertools.product(*parent_values.values()):
    # Simulate average return.
    action = {parent: val for parent, val in zip(outcome_parents, val)}
    # Compute the value of the action.
    linear_terms = [means[parent][val] for (parent, val) in action.items()]
    # Compute the means from the non-linear part.
    interaction = interaction_mean_fn(**action) if interaction_mean_fn else 0
    value = np.sum(linear_terms) + interaction

    if value > best_value:
      best_value = value
      best_action = action
  return (best_action, best_value)


def generate_beta_cpd(
    support_sizes: Mapping[str, int],
    parents: List[str],
    var_name: str,
    alpha: float = 2,
    beta: float = 5,
) -> np.ndarray:
  """Generates a random discrete probability distribution of variable name.

  Uses support_sizes and parents to calculate the appropriate parameters.
  Returns an np.ndarray where array[i,j,...,k,:] is the probability vector for
  variable name when its parents have discrete values (i,j,...,k).

  Args:
    support_sizes: A dictionary of support sizes, indexed by variable name.
    parents: A dict, indexed by variable name, with values of a list of the
      names of the parents.
    var_name: The name of the variable this cpd is for.
    alpha: Parameter of the beta distribution.
    beta: Parameter of the beta distribution.

  Returns:
    A np.ndarray describing the marginal probability of variable name
    conditioned on the support value of the parents.
  """
  sizes = [support_sizes[parent] for parent in parents]
  sizes.append(support_sizes[var_name])
  probs = np.random.beta(alpha, beta, size=sizes)
  # Normalize probability.
  return probs / np.repeat(
      np.sum(probs, axis=-1)[..., np.newaxis],
      support_sizes[var_name],
      axis=-1,
  )


def sample_dag_matrix(num_variabless: int, degree: float = 3.0) -> np.ndarray:
  """Produces an Erdos-Renyi graph's adjacency matrix.

  Args:
    num_variabless: the number of variables to include in the dag.
    degree: the average degree of the graph

  Returns:
    An adjacency matrix.
  """
  prob = float(degree) / (num_variabless - 1)
  graph = np.tril(
      (np.random.rand(num_variabless, num_variabless) < prob).astype(float),
      k=-1,
  )

  return graph
