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

"""Implementation of algorithms from the paper."""

import random
from typing import List, Mapping, Tuple, Optional, Union

from cbug import discrete_scm
from cbug import scm
from cbug import utils
import numpy as np


class Results:
  """Stores results generated from running action-elimination algorithms.

  In particular, both the run_modl algorithm and run_se returns Results objects.

  Attributes:
    s_t: A list of sets of plausibly best actions, as produced by running an
      action-elimination algorithm.
    num_actions_t: A list of ints of the number of actions remaining.
    gamma_t: A list of floats of the confidence-widths.
    n_samples_t: A list of ints of the tutal number of samples used.
    mean_est_t: A list of np.ndarrays of the mean estimates.
    true_value_t: A list of floats of the true reward generated by the
      empirically best action so far.
    best_action_t: A list of dictionaries of the empirically best action.
  """

  def __init__(self):
    # _t denotes a sequence indexed by rounds.
    self.s_t = []
    self.num_actions_t = []
    self.gamma_t = []
    self.n_samples_t = []
    self.mean_est_t = None
    self.true_value_t = []
    self.best_action_t = []

  def update(
      self,
      action_set: Union[Mapping[str, List[int]], List[Mapping[str, int]]],
      num_actions: int,
      gamma: float,
      n_samples: int,
      mean_est: np.ndarray,
      true_value: float,
      best_action: Mapping[str, int],
  ):
    """Adds an additional observation to the Results class.

    Args:
      action_set: A representation of the plausibly best actions left, either
        through a marginal action set or a list of actions.
      num_actions: The number of actions remaining.
      gamma: The confidence width.
      n_samples: The number of samples used.
      mean_est: A vector of mean-estimates for every value of every variable.
      true_value: The true value of the empirically best action.
      best_action: The empirically best action.
    """
    self.s_t.append(action_set)
    self.num_actions_t.append(num_actions)
    self.gamma_t.append(gamma)
    self.n_samples_t.append(n_samples)
    if self.mean_est_t is None:
      self.mean_est_t = mean_est.reshape(-1, 1)
    else:
      self.mean_est_t = np.hstack((self.mean_est_t, mean_est.reshape(-1, 1)))
    self.true_value_t.append(float(true_value))
    self.best_action_t.append(best_action)

  def __str__(self):
    final_action = {var: list(self.s_t[-1][var]) for var in self.s_t[-1]}
    string = (
        f"Displaying experimental results with {len(self.s_t)} rounds and total"
        f" samples {sum(self.n_samples_t)}.\nThe samples needed were"
        f" {self.n_samples_t}\nat gammas {self.gamma_t}.\nThe expected outcome"
        f" of the empirically best values are {self.true_value_t}\nand the"
        f" final actions are {final_action}\nwith a parameter"
        f" estimate\n{self.mean_est_t[:, -1]}"
    )
    return string


def reduce_actions(
    s_marginal: Mapping[str, np.ndarray],
    name_to_idx: Mapping[str, List[int]],
    mean_est: np.ndarray,
    gamma: float,
) -> Tuple[Mapping[str, List[int]], Mapping[str, int]]:
  """Uses a parameter estimate to prune the marginal action set.

  Starting from a marginal action set, this method computes the gaps for each
  variable separately then eliminates all values from the variable that have a
  gap bigger than gamma.

  Args:
    s_marginal: A dictionary of remaining actions for each variable in a
      sorted-in-ascending-order list.
    name_to_idx: A dictionary specifying which positions of mean_est correspond
      to each variable. This encoding is the same one-hot encoding by
      utils.form_data_matrix.
    mean_est: An nd.array of mean estimates. Assumed to be in the one-hot
      encoding order.
    gamma: The width of the confidence intervals for the means estimates.

  Returns:
    A new, pruned s_marginal and a guess of the best action.
  """
  new_s_marginal = {}
  best_action = {}

  # Eliminate actions variable-by-variable.
  for var in s_marginal:
    # Compute gaps.
    marginal_means = mean_est[name_to_idx[var]][s_marginal[var]]
    gaps = np.array(max(marginal_means) - marginal_means).flatten()

    idx_to_keep = np.argwhere(gaps < gamma)
    best_pos = np.argwhere(gaps == 0)[0]
    new_s_marginal[var] = s_marginal[var][idx_to_keep.flatten()]
    best_action[var] = s_marginal[var][best_pos]

  return (new_s_marginal, best_action)


def xy_optimal_design(
    s_marginal: Mapping[str, List[int]],
    support_sizes: Mapping[str, int],
    gamma: float,
    delta: float,
    cov: float,
) -> Mapping[str, np.ndarray]:
  """Solves an XY optimal design problem.

  The optimal design problem chooses a sequences of actions that will reveal as
  much information as possible about the gaps between values of the actions.

  If there is only one action remaining for a variable, then a uniform
  distribution over all the variable's actions is assumed, which provides a
  more informative mean estimate from the linear regression.

  Args:
    s_marginal: A marginal action set; composed of a dictionary of lists, each
      list representing the actions of the corresponding variable that haven't
      been eliminated yet.
    support_sizes: Dictionary of support sizes of every variable.
    gamma: Error tolerance.
    delta: Probability bound.
    cov: Scale of noise.

  Returns:
    A dictionary of a sequence of actions for each variable.
  """
  num_s = sum([len(s_marginal[var]) for var in s_marginal])
  num_samples = int(np.ceil(4 * cov * num_s * np.log(1 / delta) / gamma**2))

  marginal_points = {}
  for var in s_marginal:
    support = len(s_marginal[var])
    if support == 1:  # Make this uniform.
      num_cycles = int(np.ceil(num_samples / support_sizes[var]))
      samples = []
      for _ in range(num_cycles):
        new_list = list(np.arange(support_sizes[var]))
        np.random.shuffle(new_list)
        samples.append(new_list)
    else:
      num_cycles = int(np.ceil(num_samples / support))
      samples = []
      for _ in range(num_cycles):
        new_list = list(s_marginal[var])
        np.random.shuffle(new_list)
        samples.append(new_list)

    marginal_points[var] = np.array(samples).flatten()[:num_samples]

  return marginal_points


def run_modl(
    delta: float,
    epsilon: float,
    model: discrete_scm.DiscreteSCM,
    cov: float,
    outcome_bound: float,
    lambda_ridge: Optional[float] = .1,
    opt_scope: Optional[List[str]] = None,
    num_parents_bound: Optional[int] = None,
) -> Results:
  """Runs the MODL algorithm, Algorithm 1 of the paper.

  Args:
    delta: High probability requirement.
    epsilon: Error bound.
    model: A scm.SCM.
    cov: Noise scale parameter.
    outcome_bound: Bound of sum_i f_i.
    lambda_ridge: The regularization parameter for ridge regression.
    opt_scope: A list of variable names to include in the optimization. If
      unspecified, all variables are included.
    num_parents_bound: An upper bound on the number of parents.

  Returns:
    A Results object summarizing the experiment.
  """
  if opt_scope is None:
    opt_scope = model.var_names.copy()
    opt_scope.remove(model.outcome_variable)

  s_marginal_size = np.sum([model.support_sizes[var] for var in opt_scope])

  name_to_idx = utils.get_name_to_idx(
      {var: model.support_sizes[var] for var in opt_scope}
  )

  # Construct a full marginal.
  s_marginal = {
      name: np.arange(model.support_sizes[name]) for name in opt_scope
  }
  num_actions = sum([len(s_marginal[var]) for var in s_marginal])

  results = Results()
  mean_est = np.empty(s_marginal_size)
  mean_est[:] = np.nan

  num_phases = np.ceil(np.log2(outcome_bound / epsilon))
  current_phase = 0
  true_value = 0
  best_action = {}
  while current_phase <= num_phases:
    gamma = epsilon * 2**(num_phases - current_phase - 1)
    # Solve optimization problem.
    covariates = xy_optimal_design(
        s_marginal,
        model.support_sizes,
        gamma,
        delta / num_phases,
        cov,
    )
    # Collect data.
    n = len(covariates[opt_scope[0]])
    if n < num_actions:  # We cannot have enough data coverage.
      current_phase += 1
      results.update(
          s_marginal, num_actions, gamma, n, mean_est, true_value, best_action
      )
      continue
    samples = model.sample(n)
    data_matrix = utils.form_data_matrix(
        samples, model.support_sizes, opt_scope
    )
    # Update linear parameter vector.
    mean_est = np.linalg.inv(
        data_matrix.transpose().dot(data_matrix)
        + lambda_ridge * np.eye(s_marginal_size)
    ).dot(
        data_matrix.transpose().dot(
            samples[model.outcome_variable].reshape(-1, 1)
        )
    )
    # Eliminate the arms with dynamic algorithm.
    (s_marginal, best_action) = reduce_actions(
        s_marginal,
        name_to_idx,
        mean_est,
        gamma,
    )
    num_actions = sum([len(s_marginal[var]) for var in s_marginal])

    # Estimate the true value of the best_action.
    true_value = estimate_value(model, best_action, model.outcome_variable)

    results.update(s_marginal, num_actions, gamma, n, mean_est, true_value,
                   best_action)
    if num_actions == len(opt_scope):  # We have eliminated all but one action.
      break
    if num_parents_bound is not None:
      # Count the number of variables with a single remaining action.
      single_action = [len(s_marginal[var]) == 1 for var in s_marginal]
      if sum(single_action) >= num_parents_bound:  # All the parents are found.
        return results
    current_phase += 1

  return results


def run_se(
    delta: float,
    epsilon: float,
    model: discrete_scm.DiscreteSCM,
    cov: float,
    outcome_bound: float,
    opt_scope: Optional[List[str]] = None,
    batch_size: int = 100,
    sample_limit: int = 10000000,
    max_num_actions: int = 10000,
) -> Results:
  """Runs the successive elimination algorithm.

  Warning: this algorithm has exponential complexity in the number of parents:
  please be careful with the size of the input problem.

  Args:
    delta: High probability requirement.
    epsilon: Error bound.
    model: An scm.SCM; action variables are assumed to be discrete.
    cov: Noise scale parameter.
    outcome_bound: Bound of sum_i f_i.
    opt_scope: A list of variable names to include in the optimization. If
      unspecified, all variables are included.
    batch_size: The number of samples to gather between elimination rounds.
    sample_limit: The algorithm terminates after sample_limit total samples have
      been collected.
    max_num_actions: The total number of actions (arms) allowed; if the product
      of all actions support is larger, the method terminates.

  Returns:
    A Results object.
  """
  if opt_scope is None:
    opt_scope = model.var_names.copy()
    opt_scope.remove(model.outcome_variable)

  results = Results()

  # If the number of actions is too large, abort.
  num_actions = np.product(list(model.support_sizes.values()))
  if num_actions > max_num_actions:
    print(f"Warning: {num_actions} actions: termination could be slow.")

  actions = utils.get_full_action_set(model.support_sizes)

  candidate_arm_idx = list(range(len(actions)))
  means = []  # Dicts are not hashable.
  total_samples = 0

  # Pull every arm once.
  for action in actions:
    values = {var: action[var] for var in action}
    rewards = model.sample(batch_size, values)[model.outcome_variable]
    means.append(np.mean(rewards))
  total_samples += batch_size * len(actions)
  num_samples_per_arm = batch_size
  means = np.array(means)

  while len(candidate_arm_idx) > 1:
    # Pull all remaining arms.
    for arm_idx in candidate_arm_idx:
      values = {var: actions[arm_idx][var]for var in actions[arm_idx]}
      rewards = model.sample(batch_size, values)[model.outcome_variable]
      total_samples += batch_size

      means[arm_idx] = (
          means[arm_idx] * num_samples_per_arm + np.mean(rewards) * batch_size
      ) / (num_samples_per_arm + batch_size)

    num_samples_per_arm += batch_size

    # Do elimination.
    width = np.sqrt(
        cov * outcome_bound
        * np.log(2 * len(actions) * num_samples_per_arm**2 / delta)
        / num_samples_per_arm
    )
    best_mean = np.max([means[idx] for idx in candidate_arm_idx])
    new_candidate_arms = []
    for idx in candidate_arm_idx:
      if means[idx] > best_mean - 2 * width:
        new_candidate_arms.append(idx)
      if means[idx] == best_mean:
        best_action = actions[idx]
    candidate_arms = new_candidate_arms

    # Calculate the true value of the best_action.
    true_value = estimate_value(model, best_action, model.outcome_variable)

    results.update(
        [],
        len(candidate_arms),
        width,
        total_samples,
        means,
        true_value,
        best_action,
    )
    if width < epsilon / 2:  # All remaining arms are approximately optimal.
      return results
    if total_samples > sample_limit:
      print("Sample limit for successive elimination reached.")
      return results

  return results


def find_parents(
    target_var: str,
    delta: float,
    epsilon: float,
    model: discrete_scm.DiscreteSCM,
    cov: float,
    num_parents_bound: Optional[int] = None,
) -> Tuple[List[str], int]:
  """Finds the parents of Y in a discrete SCM variable-by-variable.

  For each variable, an arbitrary point in the support is chosen, and a
  hypothesis test for parenthood of that variable is conducted. The method can
  terminate early when enough parents have been found.

  This method assumes that the SCM is discrete and the causal effects on
  var_name are additive.

  Args:
    target_var: The name of the variable whose parents we want to find.
    delta: Desired error probability bound.
    epsilon: The minimum detectable effect.
    model: A DiscreteSCM.
    cov: Scale of the noise.
    num_parents_bound: An upper bound on the number of parents.

  Returns:
    A list of parents of the target variable.
  """
  opt_scope = model.var_names.copy()
  opt_scope.remove(target_var)
  num_vars = len(opt_scope)
  parents = []
  total_samples = 0
  var_order = random.sample(opt_scope, len(opt_scope))
  for var in var_order:
    # Calculate number of samples.
    n = int(
        np.ceil(
            8 * cov * np.log(2 * model.support_sizes[var] * num_vars / delta)
        )
        / epsilon**2
    )
    conf_set = (-np.inf, np.inf)
    action = {var: 0 for var in opt_scope}
    for value in range(model.support_sizes[var]):
      action[var] = value
      # Collect n samples.
      sample = model.sample(n, action)
      total_samples += n
      y_mean = np.mean(sample[target_var])
      conf_set = utils.interval_intersection(
          conf_set,
          (y_mean - epsilon / 2, y_mean + epsilon / 2),
      )
      if conf_set is None:  # Empty intersection.
        parents.append(var)
        break
    if num_parents_bound is not None:
      if len(parents) >= num_parents_bound:
        break

  return (parents, total_samples)


def estimate_value(model: scm.SCM,
                   best_action: Mapping[str, int],
                   outcome_variable: str,
                   n: int = 1000,
                   ) -> float:
  """Estimates the true value of the best_action."""
  return np.mean(model.sample(n, best_action)[outcome_variable])
