# Copyright 2022 The PyGlove Authors
#
# 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.
"""Categorical hyper primitives."""

import numbers
import typing
from typing import Any, Callable, Iterable, List, Optional, Tuple, Union

from pyglove.core import geno
from pyglove.core import symbolic
from pyglove.core import typing as pg_typing
from pyglove.core import utils
from pyglove.core.hyper import base
from pyglove.core.hyper import object_template


@symbolic.members([
    ('num_choices', pg_typing.Int(min_value=0).noneable(),
     ('Number of choices to make. If set to None, any number of choices is '
      'acceptable.')),
    ('candidates', pg_typing.List(pg_typing.Any()),
     ('Candidate values, which may contain nested hyper values.'
      'Candidate can customize its display value (literal) by implementing the '
      '`pg.Formattable` interface.')),
    ('choices_distinct', pg_typing.Bool(True), 'Whether choices are distinct.'),
    ('choices_sorted', pg_typing.Bool(False), 'Whether choices are sorted.'),
    ('where', pg_typing.Callable([pg_typing.Object(base.HyperPrimitive)],
                                 returns=pg_typing.Bool()).noneable(),
     ('Callable object to filter nested hyper values. If None, all nested '
      'hyper value will be included in the encoding/decoding process. '
      'Otherwise only the hyper values on which `where` returns True will be '
      'included. `where` can be useful to partition a search space into '
      'separate optimization processes. '
      'Please see `ObjectTemplate` docstr for details.'))
])
class Choices(base.HyperPrimitive):
  """Categorical choices from a list of candidates.

  Example::

    # A single categorical choice:
    v = pg.oneof([1, 2, 3])

    # A multiple categorical choice as a list:
    vs = pg.manyof(2, [1, 2, 3])

    # A hierarchical categorical choice:
    v2 = pg.oneof([
        'foo',
        'bar',
        pg.manyof(2, [1, 2, 3])
    ])

  See also:

    * :class:`pyglove.hyper.OneOf`
    * :class:`pyglove.hyper.ManyOf`
    * :func:`pyglove.oneof`
    * :func:`pyglove.manyof`
    * :func:`pyglove.permutate`
  """

  def _on_bound(self):
    """On members are bound."""
    super()._on_bound()
    if self.num_choices > len(self.candidates) and self.choices_distinct:
      raise ValueError(
          f'{len(self.candidates)} candidates cannot produce '
          f'{self.num_choices} distinct choices.')
    self._candidate_templates = [
        object_template.ObjectTemplate(c, where=self.where)
        for c in self.candidates]
    # ValueSpec for candidate.
    self._value_spec = None

  def _update_children_paths(
      self, old_path: utils.KeyPath, new_path: utils.KeyPath
  ):
    """Customized logic to update children paths."""
    super()._update_children_paths(old_path, new_path)
    for t in self._candidate_templates:
      t.root_path = self.sym_path

  @property
  def candidate_templates(self):
    """Returns candidate templates."""
    return self._candidate_templates

  @property
  def is_leaf(self) -> bool:
    """Returns whether this is a leaf node."""
    for t in self._candidate_templates:
      if not t.is_constant:
        return False
    return True

  def dna_spec(self, location: Optional[utils.KeyPath] = None) -> geno.Choices:
    """Returns corresponding DNASpec."""
    return geno.Choices(
        num_choices=self.num_choices,
        candidates=[ct.dna_spec() for ct in self._candidate_templates],
        literal_values=[
            self._literal_value(c) for i, c in enumerate(self.candidates)
        ],
        distinct=self.choices_distinct,
        sorted=self.choices_sorted,
        hints=self.hints,
        name=self.name,
        location=location or utils.KeyPath(),
    )

  def _literal_value(
      self, candidate: Any, max_len: int = 120) -> Union[int, float, str]:
    """Returns literal value for candidate."""
    if isinstance(candidate, numbers.Number):
      return candidate

    literal = utils.format(
        candidate,
        compact=True,
        hide_default_values=True,
        hide_missing_values=True,
        strip_object_id=True,
    )
    if len(literal) > max_len:
      literal = literal[:max_len - 3] + '...'
    return literal

  def _decode(self) -> List[Any]:
    """Decode a DNA into a list of object."""
    dna = self._dna
    if self.num_choices == 1:
      # Single choice.
      if not isinstance(dna.value, int):
        raise ValueError(
            utils.message_on_path(
                'Did you forget to specify values for conditional choices?\n'
                f'Expect integer for {self.__class__.__name__}. '
                f'Encountered: {dna!r}.',
                self.sym_path,
            )
        )
      if dna.value >= len(self.candidates):
        raise ValueError(
            utils.message_on_path(
                f'Choice out of range. Value: {dna.value!r}, '
                f'Candidates: {len(self.candidates)}.',
                self.sym_path,
            )
        )
      choices = [self._candidate_templates[dna.value].decode(
          geno.DNA(None, dna.children))]
    else:
      # Multi choices.
      if len(dna.children) != self.num_choices:
        raise ValueError(
            utils.message_on_path(
                'Number of DNA child values does not match the number of '
                f'choices. Child values: {dna.children!r}, '
                f'Choices: {self.num_choices}.',
                self.sym_path,
            )
        )
      if self.choices_distinct or self.choices_sorted:
        sub_dna_values = [s.value for s in dna]
        if (self.choices_distinct
            and len(set(sub_dna_values)) != len(dna.children)):
          raise ValueError(
              utils.message_on_path(
                  'DNA child values should be distinct. '
                  f'Encountered: {sub_dna_values}.',
                  self.sym_path,
              )
          )
        if self.choices_sorted and sorted(sub_dna_values) != sub_dna_values:
          raise ValueError(
              utils.message_on_path(
                  'DNA child values should be sorted. '
                  f'Encountered: {sub_dna_values}.',
                  self.sym_path,
              )
          )
      choices = []
      for i, sub_dna in enumerate(dna):
        if not isinstance(sub_dna.value, int):
          raise ValueError(
              utils.message_on_path(
                  f'Choice value should be int. Encountered: {sub_dna.value}.',
                  utils.KeyPath(i, self.sym_path),
              )
          )
        if sub_dna.value >= len(self.candidates):
          raise ValueError(
              utils.message_on_path(
                  f'Choice out of range. Value: {sub_dna.value}, '
                  f'Candidates: {len(self.candidates)}.',
                  utils.KeyPath(i, self.sym_path),
              )
          )
        choices.append(self._candidate_templates[sub_dna.value].decode(
            geno.DNA(None, sub_dna.children)))
    return choices

  def encode(self, value: List[Any]) -> geno.DNA:
    """Encode a list of values into DNA.

    Example::

        # DNA of an object containing a single OneOf.
        # {'a': 1} => DNA(0)
        {
           'a': one_of([1, 2])
        }


        # DNA of an object containing multiple OneOfs.
        # {'b': 1, 'c': bar} => DNA([0, 1])
        {
           'b': pg.oneof([1, 2]),
           'c': pg.oneof(['foo', 'bar'])
        }

        # DNA of an object containing conditional space.
        # {'a': {'b': 1} => DNA(0, 0, 0)])
        # {'a': {'b': [4, 7]} => DNA(1, [(0, 1), 2])
        # {'a': {'b': 'bar'} => DNA(2)
        {
           'a': {
              'b': pg.oneof([
                pg.oneof([
                  pg.oneof([1, 2]),
                  pg.oneof(3, 4)]),
                  pg.manyof(2, [
                    pg.oneof([4, 5]),
                    6,
                    7
                  ]),
                ]),
                'bar',
              ])
           }
        }

    Args:
      value: A list of value that can match choice candidates.

    Returns:
      Encoded DNA.

    Raises:
      ValueError if value cannot be encoded.
    """
    if not isinstance(value, list):
      raise ValueError(
          utils.message_on_path(
              'Cannot encode value: value should be a list type. '
              f'Encountered: {value!r}.',
              self.sym_path,
          )
      )
    choices = []
    if self.num_choices is not None and len(value) != self.num_choices:
      raise ValueError(
          utils.message_on_path(
              'Length of input list is different from the number of choices '
              f'({self.num_choices}). Encountered: {value}.',
              self.sym_path,
          )
      )
    for v in value:
      choice_id = None
      child_dna = None
      for i, b in enumerate(self._candidate_templates):
        succeeded, child_dna = b.try_encode(v)
        if succeeded:
          choice_id = i
          break
      if child_dna is None:
        raise ValueError(
            utils.message_on_path(
                'Cannot encode value: no candidates matches with '
                f'the value. Value: {v!r}, Candidates: {self.candidates}.',
                self.sym_path,
            )
        )
      choices.append(geno.DNA(choice_id, [child_dna]))
    return geno.DNA(None, choices)


@symbolic.members(
    [],
    init_arg_list=[
        'num_choices', 'candidates', 'choices_distinct',
        'choices_sorted', 'hints'
    ],
    # TODO(daiyip): Change to 'ManyOf' once existing code migrates to ManyOf.
    serialization_key='hyper.ManyOf',
    additional_keys=['pyglove.generators.genetic.ChoiceList']
)
class ManyOf(Choices):
  """N Choose K.

  Example::

    # Chooses 2 distinct candidates.
    v = pg.manyof(2, [1, 2, 3])

    # Chooses 2 non-distinct candidates.
    v = pg.manyof(2, [1, 2, 3], distinct=False)

    # Chooses 2 distinct candidates sorted by their indices.
    v = pg.manyof(2, [1, 2, 3], sorted=True)

    # Permutates the candidates.
    v = pg.permutate([1, 2, 3])

    # A hierarchical categorical choice:
    v2 = pg.manyof(2, [
        'foo',
        'bar',
        pg.oneof([1, 2, 3])
    ])

  See also:

    * :func:`pyglove.manyof`
    * :func:`pyglove.permutate`
    * :class:`pyglove.hyper.Choices`
    * :class:`pyglove.hyper.OneOf`
    * :class:`pyglove.hyper.Float`
    * :class:`pyglove.hyper.CustomHyper`
  """

  def custom_apply(
      self,
      path: utils.KeyPath,
      value_spec: pg_typing.ValueSpec,
      allow_partial: bool,
      child_transform: Optional[
          Callable[[utils.KeyPath, pg_typing.Field, Any], Any]
      ] = None,
  ) -> Tuple[bool, 'Choices']:
    """Validate candidates during value_spec binding time."""
    # Check if value_spec directly accepts `self`.
    if value_spec.value_type and isinstance(self, value_spec.value_type):
      return (False, self)

    if self._value_spec:
      src_spec = self._value_spec
      dest_spec = value_spec
      if not dest_spec.is_compatible(src_spec):
        raise TypeError(
            utils.message_on_path(
                f'Cannot bind an incompatible value spec {dest_spec!r} '
                f'to {self.__class__.__name__} with bound spec {src_spec!r}.',
                path,
            )
        )
      return (False, self)

    list_spec = typing.cast(
        pg_typing.List,
        pg_typing.ensure_value_spec(
            value_spec, pg_typing.List(pg_typing.Any()), path))
    if list_spec:
      for i, c in enumerate(self.candidates):
        list_spec.element.value.apply(
            c,
            self._allow_partial,
            root_path=path + f'candidates[{i}]')
    self._value_spec = list_spec
    return (False, self)


@symbolic.members(
    [
        ('num_choices', 1)
    ],
    init_arg_list=['candidates', 'hints', 'where'],
    serialization_key='hyper.OneOf',
    additional_keys=['pyglove.generators.genetic.ChoiceValue']
)
class OneOf(Choices):
  """N Choose 1.

  Example::

    # A single categorical choice:
    v = pg.oneof([1, 2, 3])

    # A hierarchical categorical choice:
    v2 = pg.oneof([
        'foo',
        'bar',
        pg.oneof([1, 2, 3])
    ])

  See also:

    * :func:`pyglove.oneof`
    * :class:`pyglove.hyper.Choices`
    * :class:`pyglove.hyper.ManyOf`
    * :class:`pyglove.hyper.Float`
    * :class:`pyglove.hyper.CustomHyper`
  """

  def _on_bound(self):
    """Event triggered when members are bound."""
    super()._on_bound()
    assert self.num_choices == 1

  def _decode(self) -> Any:
    """Decode a DNA into an object."""
    return super()._decode()[0]

  def encode(self, value: Any) -> geno.DNA:
    """Encode a value into a DNA."""
    # NOTE(daiyip): Single choice DNA will automatically be pulled
    # up from children to current node. Thus we simply returns
    # encoded DNA from parent node.
    return super().encode([value])

  def custom_apply(
      self,
      path: utils.KeyPath,
      value_spec: pg_typing.ValueSpec,
      allow_partial: bool,
      child_transform: Optional[
          Callable[[utils.KeyPath, pg_typing.Field, Any], Any]
      ] = None,
  ) -> Tuple[bool, 'OneOf']:
    """Validate candidates during value_spec binding time."""
    # Check if value_spec directly accepts `self`.
    if value_spec.value_type and isinstance(self, value_spec.value_type):
      return (False, self)

    if self._value_spec:
      if not value_spec.is_compatible(self._value_spec):
        raise TypeError(
            utils.message_on_path(
                f'Cannot bind an incompatible value spec {value_spec!r} '
                f'to {self.__class__.__name__} with bound '
                f'spec {self._value_spec!r}.',
                path,
            )
        )
      return (False, self)

    for i, c in enumerate(self.candidates):
      value_spec.apply(
          c,
          self._allow_partial,
          root_path=path + f'candidates[{i}]')
    self._value_spec = value_spec
    return (False, self)


#
# Helper methods for creating hyper values.
#


def oneof(candidates: Iterable[Any],
          *,
          name: Optional[str] = None,
          hints: Optional[Any] = None) -> Any:
  """N choose 1.

  Example::

    @pg.members([
      ('x', pg.typing.Int())
    ])
    class A(pg.Object):
      pass

    # A single categorical choice:
    v = pg.oneof([1, 2, 3])

    # A complex type as candidate.
    v1 = pg.oneof(['a', {'x': 1}, A(1)])

    # A hierarchical categorical choice:
    v2 = pg.oneof([
        'foo',
        'bar',
        A(pg.oneof([1, 2, 3]))
    ])

  See also:

    * :class:`pyglove.hyper.OneOf`
    * :func:`pyglove.manyof`
    * :func:`pyglove.floatv`
    * :func:`pyglove.permutate`
    * :func:`pyglove.evolve`

  .. note::

    Under symbolic mode (by default), `pg.oneof` returns a ``pg.hyper.OneOf``
    object. Under dynamic evaluation mode, which is called under the context of
    :meth:`pyglove.hyper.DynamicEvaluationContext.collect` or
    :meth:`pyglove.hyper.DynamicEvaluationContext.apply`, it evaluates to
    a concrete candidate value.

    To use conditional search space in dynamic evaluation mode, the candidate
    should be wrapped with a `lambda` function, which is not necessary under
    symbolic mode. For example::

      pg.oneof([lambda: pg.oneof([0, 1], name='sub'), 2], name='root')

  Args:
    candidates: Candidates to select from. Items of candidate can be any type,
      therefore it can have nested hyper primitives, which forms a hierarchical
      search space.
    name: A name that can be used to identify a decision point in the search
      space. This is needed when the code to instantiate the same hyper
      primitive may be called multiple times under a
      `pg.DynamicEvaluationContext.collect` context or under a
      `pg.DynamicEvaluationContext.apply` context.
    hints: An optional value which acts as a hint for the controller.

  Returns:
    In symbolic mode, this function returns a `ChoiceValue`.
    In dynamic evaluation mode, this function returns one of the items in
    `candidates`.
    If evaluated under a `pg.DynamicEvaluationContext.apply` scope,
    this function will return the selected candidate.
    If evaluated under a `pg.DynamicEvaluationContext.collect`
    scope, it will return the first candidate.
  """
  return OneOf(candidates=list(candidates), name=name, hints=hints)


def manyof(k: int,
           candidates: Iterable[Any],
           distinct: bool = True,
           sorted: bool = False,    # pylint: disable=redefined-builtin
           *,
           name: Optional[str] = None,
           hints: Optional[Any] = None,
           **kwargs) -> Any:
  """N choose K.

  Example::

    @pg.members([
      ('x', pg.typing.Int())
    ])
    class A(pg.Object):
      pass

    # Chooses 2 distinct candidates.
    v = pg.manyof(2, [1, 2, 3])

    # Chooses 2 non-distinct candidates.
    v = pg.manyof(2, [1, 2, 3], distinct=False)

    # Chooses 2 distinct candidates sorted by their indices.
    v = pg.manyof(2, [1, 2, 3], sorted=True)

    # A complex type as candidate.
    v1 = pg.manyof(2, ['a', {'x': 1}, A(1)])

    # A hierarchical categorical choice:
    v2 = pg.manyof(2, [
        'foo',
        'bar',
        A(pg.oneof([1, 2, 3]))
    ])

  .. note::

    Under symbolic mode (by default), `pg.manyof` returns a ``pg.hyper.ManyOf``
    object. Under dynamic evaluation mode, which is called under the context of
    :meth:`pyglove.hyper.DynamicEvaluationContext.collect` or
    :meth:`pyglove.hyper.DynamicEvaluationContext.apply`, it evaluates to
    a concrete candidate value.

    To use conditional search space in dynamic evaluate mode, the candidate
    should be wrapped with a `lambda` function, which is not necessary under
    symbolic mode. For example::

        pg.manyof(2, [
           lambda: pg.oneof([0, 1], name='sub_a'),
           lambda: pg.floatv(0.0, 1.0, name='sub_b'),
           lambda: pg.manyof(2, ['a', 'b', 'c'], name='sub_c')
        ], name='root')

  See also:

    * :class:`pyglove.hyper.ManyOf`
    * :func:`pyglove.manyof`
    * :func:`pyglove.floatv`
    * :func:`pyglove.permutate`
    * :func:`pyglove.evolve`

  Args:
    k: number of choices to make. Should be no larger than the length of
      `candidates` unless `choice_distinct` is set to False,
    candidates: Candidates to select from. Items of candidate can be any type,
      therefore it can have nested hyper primitives, which forms a hierarchical
      search space.
    distinct: If True, each choice needs to be unique.
    sorted: If True, choices are sorted by their indices in the
      candidates.
    name: A name that can be used to identify a decision point in the search
      space. This is needed when the code to instantiate the same hyper
      primitive may be called multiple times under a
      `pg.DynamicEvaluationContext.collect` context or a
      `pg.DynamicEvaluationContext.apply` context.
    hints: An optional value which acts as a hint for the controller.
    **kwargs: Keyword arguments for backward compatibility.
      `choices_distinct`: Old name for `distinct`.
      `choices_sorted`: Old name for `sorted`.

  Returns:
    In symbolic mode, this function returns a `Choices`.
    In dynamic evaluate mode, this function returns a list of items in
    `candidates`.
    If evaluated under a `pg.DynamicEvaluationContext.apply` scope,
    this function will return a list of selected candidates.
    If evaluated under a `pg.DynamicEvaluationContext.collect`
    scope, it will return a list of the first valid combination from the
    `candidates`. For example::

        # Evaluates to [0, 1, 2].
        manyof(3, range(5))

        # Evaluates to [0, 0, 0].
        manyof(3, range(5), distinct=False)
  """
  choices_distinct = kwargs.pop('choices_distinct', distinct)
  choices_sorted = kwargs.pop('choices_sorted', sorted)
  return ManyOf(
      num_choices=k,
      candidates=list(candidates),
      choices_distinct=choices_distinct,
      choices_sorted=choices_sorted,
      name=name,
      hints=hints)


def permutate(candidates: Iterable[Any],
              name: Optional[str] = None,
              hints: Optional[Any] = None) -> Any:
  """Permuatation of candidates.

  Example::

    @pg.members([
      ('x', pg.typing.Int())
    ])
    class A(pg.Object):
      pass

    # Permutates the candidates.
    v = pg.permutate([1, 2, 3])

    # A complex type as candidate.
    v1 = pg.permutate(['a', {'x': 1}, A(1)])

    # A hierarchical categorical choice:
    v2 = pg.permutate([
        'foo',
        'bar',
        A(pg.oneof([1, 2, 3]))
    ])

  .. note::

    Under symbolic mode (by default), `pg.manyof` returns a ``pg.hyper.ManyOf``
    object. Under dynamic evaluate mode, which is called under the context of
    :meth:`pyglove.hyper.DynamicEvaluationContext.collect` or
    :meth:`pyglove.hyper.DynamicEvaluationContext.apply`, it evaluates to
    a concrete candidate value.

    To use conditional search space in dynamic evaluate mode, the candidate
    should be wrapped with a `lambda` function, which is not necessary under
    symbolic mode. For example::

      pg.permutate([
         lambda: pg.oneof([0, 1], name='sub_a'),
         lambda: pg.floatv(0.0, 1.0, name='sub_b'),
         lambda: pg.manyof(2, ['a', 'b', 'c'], name='sub_c')
      ], name='root')

  See also:

    * :class:`pyglove.hyper.ManyOf`
    * :func:`pyglove.oneof`
    * :func:`pyglove.manyof`
    * :func:`pyglove.floatv`
    * :func:`pyglove.evolve`

  Args:
    candidates: Candidates to select from. Items of candidate can be any type,
      therefore it can have nested hyper primitives, which forms a hierarchical
      search space.
    name: A name that can be used to identify a decision point in the search
      space. This is needed when the code to instantiate the same hyper
      primitive may be called multiple times under a
      `pg.DynamicEvaluationContext.collect` context or a
      `pg.DynamicEvaluationContext.apply` context.
    hints: An optional value which acts as a hint for the controller.

  Returns:
    In symbolic mode, this function returns a `Choices`.
    In dynamic evaluate mode, this function returns a permutation from
    `candidates`.
    If evaluated under an `pg.DynamicEvaluationContext.apply` scope,
    this function will return a permutation of candidates based on controller
    decisions.
    If evaluated under a `pg.DynamicEvaluationContext.collect`
    scope, it will return the first valid permutation.
    For example::

      # Evaluates to [0, 1, 2, 3, 4].
      permutate(range(5), name='numbers')
  """
  candidates = list(candidates)
  return manyof(
      len(candidates), candidates,
      choices_distinct=True, choices_sorted=False, name=name, hints=hints)
