# Copyright 2020 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
#
#     https://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 few common classes to be used for the library."""

import dataclasses
import enum
from typing import Any, Callable, Dict, NamedTuple

import pandas as pd


class GeoAssignment(enum.IntEnum):
  """Defines the values for Treatment/Control assignment."""
  CONTROL = 2
  TREATMENT = 1
  EXCLUDED = -1


class ExperimentPeriod(enum.IntEnum):
  """Defines the values for Pre-Experiment, Experiment and Post-Experiment."""
  PRE_EXPERIMENT = 0
  EXPERIMENT = 1
  POST_EXPERIMENT = 2


class GeoXType(enum.Enum):
  """Defines the types of experimental arms."""
  CONTROL = enum.auto()  # business as usual
  GO_DARK = enum.auto()  # stop ad spend (0 during test, >0 during pretest)
  HEAVY_UP = enum.auto()  # increase ad spend
  HEAVY_DOWN = enum.auto()  # decrease ad spend
  HOLD_BACK = enum.auto()  # start ad spend during test (0 during pretest)
  GO_DARK_TREATMENT_NOT_BAU_CONTROL = enum.auto(
  )  # 0 ad spend in treatment, changed (up/down) ad spend in control


class DesignSelectionMethod(enum.Enum):
  """Defines the different methods for design selection.

  MIN_RMSE selects the design that minimizes the rmse.
  MIN_COST_ADJUSTED_RMSE selects the design that minimizes cost_adjusted_rmse.
  MIN_COST_ADJUSTED_RMSE_ONE_STD selects the largest design within one standard
    deviation of the design that minimizes cost_adjusted_rmse.
  MIN_BIAS selects the design that minimizes the bias.
  LARGEST_NO_BIAS selects the largest design with bias/std < max_bias_std.
  LARGEST_LOW_RMSE selects the largest design with rmse < max_rmse.
  LARGEST selects the largest design.
  """
  MIN_RMSE = enum.auto()
  MIN_COST_ADJUSTED_RMSE = enum.auto()
  MIN_COST_ADJUSTED_RMSE_ONE_STD = enum.auto()
  MIN_BIAS = enum.auto()
  LARGEST_NO_BIAS = enum.auto()
  LARGEST_LOW_RMSE = enum.auto()
  LARGEST = enum.auto()


class GeoLevelData(NamedTuple):
  """Geo level data."""
  geo: int
  response: float
  spend: float


class GeoLevelPotentialOutcomes(NamedTuple):
  """Two potential outcomes."""
  controlled: GeoLevelData
  treated: GeoLevelData


@dataclasses.dataclass
class TimeWindow:
  """Defines a time window using first day and last day."""
  first_day: pd.Timestamp
  last_day: pd.Timestamp

  def __post_init__(self):
    if not isinstance(self.first_day, pd.Timestamp):
      self.first_day = pd.Timestamp(self.first_day)
    if not isinstance(self.last_day, pd.Timestamp):
      self.last_day = pd.Timestamp(self.last_day)

    if self.first_day > self.last_day:
      raise ValueError('TimeWindow(): first_day > last_day: {!r}, {!r}'.format(
          self.first_day, self.last_day))

  def contains(self, dates: pd.Series) -> pd.Series:
    """Checks whether dates fall into this time window.

    Args:
        dates: a series of datetime.

    Returns:
        a series of bools.
    """
    return pd.to_datetime(dates).between(self.first_day, self.last_day)

  def is_included_in(self, other) -> bool:
    """Checks whether it falls into the 'other' time window.

    Args:
      other: another TimeWindow.

    Returns:
      True iff it falls into the 'other' time window.
    """
    return (self.first_day >= other.first_day and
            self.last_day <= other.last_day)

  def finds_latest_date_moving_window_in(self,
                                         input_dates: pd.Series) -> pd.Series:
    """Finds the latest dates in input_dates.

    The number of latest dates is determined by the number of distinct dates
    that fall into this time window. For example, if input_dates consists of
    2020-02-01, 2020-02-02, 2020-02-03, 2020-02-04, and the current time window
    is TimeWindow('2020-01-01', '2020-02-02') covering 2 dates in input_dates.
    The results will return the latest 2 dates: 2020-02-03, 2020-02-04.

    Args:
        input_dates: a series of datetime.

    Returns:
        a series of datetime (a subset of distinct entries in input_dates).

    Raises:
        ValueError: if input_dates does not overlap with current time window.
    """
    sorted_dates = input_dates.sort_values(ascending=False).drop_duplicates()
    duration = self.contains(sorted_dates).sum()
    if duration == 0:
      raise ValueError('Overlap between input_dates and current time window.')
    return sorted_dates[:duration]


def time_window_shifted_from(time_window: TimeWindow,
                             new_first_day: pd.Timestamp) -> TimeWindow:
  """Shift a TimeWindow to a new first_day."""
  delta = time_window.last_day - time_window.first_day
  return TimeWindow(new_first_day, new_first_day + delta)


@dataclasses.dataclass
class FormatOptions:
  """Defines the formatting parameters to change the format of a table.

  The class contains three attributes:
    column:column on which to apply the format.
    function: function that changes the format of the cell of a table.
    args: arguments to be passed to the function above or pandas Styling.
  """
  column: str
  function: Callable[..., Any]
  args: Dict[str, Any]

  def __post_init__(self):
    if 'axis' in self.args:
      self.args['column'] = self.column
    else:
      self.args['subset'] = self.column
