# Copyright 2022 Google LLC.
#
# 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 optimizing your media based on media mix models."""
import functools
from typing import Optional, Tuple, Union
from absl import logging
import jax
import jax.numpy as jnp
from scipy import optimize

from lightweight_mmm import lightweight_mmm
from lightweight_mmm import preprocessing


@functools.partial(
    jax.jit,
    static_argnames=("media_mix_model", "media_input_shape", "target_scaler",
                     "media_scaler"))
def _objective_function(extra_features: jnp.ndarray,
                        media_mix_model: lightweight_mmm.LightweightMMM,
                        media_input_shape: Tuple[int,
                                                 int], media_gap: Optional[int],
                        target_scaler: Optional[preprocessing.CustomScaler],
                        media_scaler: preprocessing.CustomScaler,
                        geo_ratio: jnp.array,
                        seed: Optional[int],
                        media_values: jnp.ndarray) -> jnp.float64:
  """Objective function to calculate the sum of all predictions of the model.

  Args:
    extra_features: Extra features the model requires for prediction.
    media_mix_model: Media mix model to use. Must have a predict method to be
      used.
    media_input_shape: Input shape of the data required by the model to get
      predictions. This is needed since optimization might flatten some arrays
      and they need to be reshaped before running new predictions.
    media_gap: Media data gap between the end of training data and the start of
      the out of sample media given. Eg. if 100 weeks of data were used for
      training and prediction starts 2 months after training data finished we
      need to provide the 8 weeks missing between the training data and the
      prediction data so data transformations (adstock, carryover, ...) can take
      place correctly.
    target_scaler: Scaler that was used to scale the target before training.
    media_scaler: Scaler that was used to scale the media data before training.
    geo_ratio: The ratio to split channel media across geo. Should sum up to 1
      for each channel and should have shape (c, g).
    seed: Seed to use for PRNGKey during sampling. For replicability run
      this function and any other function that gets predictions with the same
      seed.
    media_values: Media values required by the model to run predictions.

  Returns:
    The negative value of the sum of all predictions.
  """
  if hasattr(media_mix_model, "n_geos") and media_mix_model.n_geos > 1:
    media_values = geo_ratio * jnp.expand_dims(media_values, axis=-1)
  media_values = jnp.tile(
      media_values / media_input_shape[0], reps=media_input_shape[0])
  # Distribute budget of each channels across time.
  media_values = jnp.reshape(a=media_values, newshape=media_input_shape)
  media_values = media_scaler.transform(media_values)
  return -jnp.sum(
      media_mix_model.predict(
          media=media_values.reshape(media_input_shape),
          extra_features=extra_features,
          media_gap=media_gap,
          target_scaler=target_scaler,
          seed=seed).mean(axis=0))


@jax.jit
def _budget_constraint(media: jnp.ndarray,
                       prices: jnp.ndarray,
                       budget: jnp.ndarray) -> jnp.float64:
  """Calculates optimization constraint to keep spend equal to the budget.

  Args:
    media: Array with the values of the media for this iteration.
    prices: Prices of each media channel at any given time.
    budget: Total budget of the optimization.

  Returns:
    The result from substracting the total spending and the budget.
  """
  media = media.reshape((-1, len(prices)))
  return jnp.sum(media * prices) - budget


def _get_lower_and_upper_bounds(
    media: jnp.ndarray,
    n_time_periods: int,
    lower_pct: jnp.ndarray,
    upper_pct: jnp.ndarray,
    media_scaler: Optional[preprocessing.CustomScaler] = None
) -> optimize.Bounds:
  """Gets the lower and upper bounds for optimisation based on historic data.

  It creates an upper bound based on a percentage above the mean value on
  each channel and a lower bound based on a relative decrease of the mean
  value.

  Args:
    media: Media data to get historic mean.
    n_time_periods: Number of time periods to optimize for. If model is built on
      weekly data, this would be the number of weeks ahead to optimize.
    lower_pct: Relative percentage decrease from the mean value to consider as
      new lower bound.
    upper_pct: Relative percentage increase from the mean value to consider as
      new upper bound.
    media_scaler: Scaler that was used to scale the media data before training.

  Returns:
    A list of tuples with the lower and upper bound for each media channel.
  """
  if media.ndim == 3:
    lower_pct = jnp.expand_dims(lower_pct, axis=-1)
    upper_pct = jnp.expand_dims(upper_pct, axis=-1)

  mean_data = media.mean(axis=0)
  lower_bounds = jnp.maximum(mean_data * (1 - lower_pct), 0)
  upper_bounds = mean_data * (1 + upper_pct)

  if media_scaler:
    lower_bounds = media_scaler.inverse_transform(lower_bounds)
    upper_bounds = media_scaler.inverse_transform(upper_bounds)

  if media.ndim == 3:
    lower_bounds = lower_bounds.sum(axis=-1)
    upper_bounds = upper_bounds.sum(axis=-1)

  return optimize.Bounds(lb=lower_bounds * n_time_periods,
                         ub=upper_bounds * n_time_periods)


def _generate_starting_values(
    n_time_periods: int, media: jnp.ndarray,
    media_scaler: preprocessing.CustomScaler,
    budget: Union[float, int],
    prices: jnp.ndarray,
) -> jnp.ndarray:
  """Generates starting values based on historic allocation and budget.

  In order to make a comparison we can take the allocation of the last
  `n_time_periods` and scale it based on the given budget. Given this, one can
  compare how this initial values (based on average historic allocation) compare
  to the output of the optimisation in terms of sales/KPI.

  Args:
    n_time_periods: Number of time periods the optimization will be done with.
    media: Historic media data the model was trained with.
    media_scaler: Scaler that was used to scale the media data before training.
    budget: Total budget to allocate during the optimization time.
    prices: An array with shape (n_media_channels,) for the cost of each media
      channel unit.

  Returns:
    An array with the starting value for each media channel for the
      optimization.
  """
  previous_allocation = media.mean(axis=0) * n_time_periods
  if media_scaler:  # Scale before sum as geo scaler has shape (c, g).
    previous_allocation = media_scaler.inverse_transform(previous_allocation)

  if media.ndim == 3:
    previous_allocation = previous_allocation.sum(axis=-1)

  avg_spend_per_channel = previous_allocation * prices
  pct_spend_per_channel = avg_spend_per_channel / avg_spend_per_channel.sum()
  budget_per_channel = budget * pct_spend_per_channel
  media_unit_per_channel = budget_per_channel / prices
  return media_unit_per_channel


def find_optimal_budgets(
    n_time_periods: int,
    media_mix_model: lightweight_mmm.LightweightMMM,
    budget: Union[float, int],
    prices: jnp.ndarray,
    extra_features: Optional[jnp.ndarray] = None,
    media_gap: Optional[jnp.ndarray] = None,
    target_scaler: Optional[preprocessing.CustomScaler] = None,
    media_scaler: Optional[preprocessing.CustomScaler] = None,
    bounds_lower_pct: Union[float, jnp.ndarray] = .2,
    bounds_upper_pct: Union[float, jnp.ndarray] = .2,
    max_iterations: int = 200,
    solver_func_tolerance: float = 1e-06,
    solver_step_size: float = 1.4901161193847656e-08,
    seed: Optional[int] = None) -> optimize.OptimizeResult:
  """Finds the best media allocation based on MMM model, prices and a budget.

  Args:
    n_time_periods: Number of time periods to optimize for. If model is built on
      weekly data, this would be the number of weeks ahead to optimize.
    media_mix_model: Media mix model to use for the optimization.
    budget: Total budget to allocate during the optimization time.
    prices: An array with shape (n_media_channels,) for the cost of each media
      channel unit.
    extra_features: Extra features needed for the model to predict.
    media_gap: Media data gap between the end of training data and the start of
      the out of sample media given. Eg. if 100 weeks of data were used for
      training and prediction starts 8 weeks after training data finished we
      need to provide the 8 weeks missing between the training data and the
      prediction data so data transformations (adstock, carryover, ...) can take
      place correctly.
    target_scaler: Scaler that was used to scale the target before training.
    media_scaler: Scaler that was used to scale the media data before training.
    bounds_lower_pct: Relative percentage decrease from the mean value to
      consider as new lower bound.
    bounds_upper_pct: Relative percentage increase from the mean value to
      consider as new upper bound.
    max_iterations: Number of max iterations to use for the SLSQP scipy
      optimizer. Default is 200.
    solver_func_tolerance: Precision goal for the value of the prediction in
      the stopping criterion. Maps directly to scipy's `ftol`. Intended only
      for advanced users. For more details see:
      https://docs.scipy.org/doc/scipy/reference/optimize.minimize-slsqp.html#optimize-minimize-slsqp.
    solver_step_size: Step size used for numerical approximation of the
      Jacobian. Maps directly to scipy's `eps`. Intended only for advanced
      users. For more details see:
      https://docs.scipy.org/doc/scipy/reference/optimize.minimize-slsqp.html#optimize-minimize-slsqp.
    seed: Seed to use for PRNGKey during sampling. For replicability run
      this function and any other function that gets predictions with the same
      seed.

  Returns:
    solution: OptimizeResult object containing the results of the optimization.
    kpi_without_optim: Predicted target based on original allocation proportion
    among channels from the historical data.
    starting_values: Budget Allocation based on original allocation proportion
    and the given total budget.
  """
  if not hasattr(media_mix_model, "media"):
    raise ValueError(
        "The passed model has not been trained. Please fit the model before "
        "running optimization.")
  jax.config.update("jax_enable_x64", True)

  if isinstance(bounds_lower_pct, float):
    bounds_lower_pct = jnp.repeat(a=bounds_lower_pct, repeats=len(prices))
  if isinstance(bounds_upper_pct, float):
    bounds_upper_pct = jnp.repeat(a=bounds_upper_pct, repeats=len(prices))

  bounds = _get_lower_and_upper_bounds(
      media=media_mix_model.media,
      n_time_periods=n_time_periods,
      lower_pct=bounds_lower_pct,
      upper_pct=bounds_upper_pct,
      media_scaler=media_scaler)
  if jnp.sum(bounds.lb * prices) > budget:
    logging.warning(
        "Budget given is smaller than the lower bounds of the constraints for "
        "optimization. This will lead to faulty optimization. Please either "
        "increase the budget or change the lower bound by increasing the "
        "percentage decrease with the `bounds_lower_pct` parameter.")
  if jnp.sum(bounds.ub * prices) < budget:
    logging.warning(
        "Budget given is larger than the upper bounds of the constraints for "
        "optimization. This will lead to faulty optimization. Please either "
        "reduce the budget or change the upper bound by increasing the "
        "percentage increase with the `bounds_upper_pct` parameter.")

  starting_values = _generate_starting_values(
      n_time_periods=n_time_periods,
      media=media_mix_model.media,
      media_scaler=media_scaler,
      budget=budget,
      prices=prices,
  )
  if not media_scaler:
    media_scaler = preprocessing.CustomScaler(multiply_by=1, divide_by=1)
  if media_mix_model.n_geos == 1:
    geo_ratio = 1.0
  else:
    average_per_time = media_mix_model.media.mean(axis=0)
    geo_ratio = average_per_time / jnp.expand_dims(
        average_per_time.sum(axis=-1), axis=-1)
  media_input_shape = (n_time_periods, *media_mix_model.media.shape[1:])
  partial_objective_function = functools.partial(
      _objective_function, extra_features, media_mix_model,
      media_input_shape, media_gap,
      target_scaler, media_scaler, geo_ratio, seed)
  solution = optimize.minimize(
      fun=partial_objective_function,
      x0=starting_values,
      bounds=bounds,
      method="SLSQP",
      jac="3-point",
      options={
          "maxiter": max_iterations,
          "disp": True,
          "ftol": solver_func_tolerance,
          "eps": solver_step_size,
      },
      constraints={
          "type": "eq",
          "fun": _budget_constraint,
          "args": (prices, budget)
      })

  kpi_without_optim = _objective_function(extra_features=extra_features,
                                          media_mix_model=media_mix_model,
                                          media_input_shape=media_input_shape,
                                          media_gap=media_gap,
                                          target_scaler=target_scaler,
                                          media_scaler=media_scaler,
                                          seed=seed,
                                          geo_ratio=geo_ratio,
                                          media_values=starting_values)
  logging.info("KPI without optimization: %r", -1 * kpi_without_optim.item())
  logging.info("KPI with optimization: %r", -1 * solution.fun)

  jax.config.update("jax_enable_x64", False)
  # TODO(yukaabe): Create an object to contain the results of this function.
  return solution, kpi_without_optim, starting_values
