# Copyright 2023 The etils 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.

"""Tree API backends."""

from __future__ import annotations

import abc
import collections
import collections.abc
import itertools
import types
from typing import Any, Callable, TypeVar

from etils import epy
from etils.etree.typing import Tree  # pylint: disable=g-importing-member

_T = TypeVar('_T')
_Tin = Any  # TypeVar('_Tin')
_Tout = Any  # TypeVar('_Tout')

# Structure which allow to reconstruct the tree
# * jax: TreeDef
# * tf.nest/tree: The original Tree
_TreeDef = Any


class Backend(abc.ABC):
  """Tree API backend.

  Note: The backend lazy-import the module on first call. This
  allow to use `etree` with Jax even if TF isn't installed (and
  vice-versa).
  """

  @epy.cached_property
  def module(self) -> types.ModuleType:
    """Module used by the backend."""
    try:
      module = self.import_module()
    except ImportError as e:
      epy.reraise(e, suffix=f'etree backend require {self.MODULE_NAME!r}.')
    return module

  @abc.abstractmethod
  def import_module(self) -> types.ModuleType:
    """Import and return the module."""
    raise NotImplementedError

  @abc.abstractmethod
  def map(
      self,
      map_fn: Callable[..., _Tout],  # Callable[[_Tin0, _Tin1,...], Tout]
      *trees: Tree[_Tin],  # _Tin0, _Tin1,...
  ) -> Tree[_Tout]:
    """Like `tf.nest.map_structure`."""
    raise NotImplementedError

  @abc.abstractmethod
  def flatten(self, tree: Tree[_T]) -> tuple[list[_T], _TreeDef]:
    """Like `tf.nest.flatten`."""
    raise NotImplementedError

  @abc.abstractmethod
  def unflatten(self, structure: _TreeDef, flat_sequence: list[_T]) -> Tree[_T]:
    raise NotImplementedError

  @abc.abstractmethod
  def assert_same_structure(
      self,
      tree0: Tree[Any],
      tree1: Tree[Any],
  ) -> None:
    raise NotImplementedError


class Jax(Backend):
  """`jax.tree_util` backend."""

  def import_module(self):
    import jax  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error

    return jax.tree_util

  def map(self, map_fn, *trees):
    return self.module.tree_map(map_fn, *trees)

  def flatten(self, tree):
    flat_vals, treedef = self.module.tree_flatten(tree)
    return flat_vals, treedef

  def unflatten(self, structure, flat_sequence):
    return structure.unflatten(flat_sequence)

  def assert_same_structure(
      self,
      tree0: Tree[Any],
      tree1: Tree[Any],
  ):
    treedef0 = self.module.tree_structure(tree0)
    treedef1 = self.module.tree_structure(tree1)
    if treedef0 != treedef1:
      raise ValueError(
          "The two structures don't have the same nested structure.\n"
          f"Left: {treedef0}\nRight: {treedef1}"
      )


class Optree(Jax):
  """`optree` backend."""

  def import_module(self):
    import optree  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error

    return optree


class DmTree(Backend):
  """`tree` backend."""

  def import_module(self):
    import tree  # pylint: disable=g-import-not-at-top  # type: ignore

    return tree

  def map(self, map_fn, *trees):
    return self.module.map_structure(map_fn, *trees)

  def flatten(self, tree):
    return self.module.flatten(tree), tree

  def unflatten(self, structure, flat_sequence):
    return self.module.unflatten_as(structure, flat_sequence)

  def assert_same_structure(
      self,
      tree0: Tree[Any],
      tree1: Tree[Any],
  ):
    self.module.assert_same_structure(tree0, tree1)


class Nest(Backend):
  """`tf.nest` backend."""

  def import_module(self):
    import tensorflow as tf  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error

    return tf.nest

  def map(self, map_fn, *trees):
    return self.module.map_structure(map_fn, *trees)

  def flatten(self, tree):
    return self.module.flatten(tree), tree

  def unflatten(self, structure, flat_sequence):
    return self.module.pack_sequence_as(structure, flat_sequence)

  def assert_same_structure(
      self,
      tree0: Tree[Any],
      tree1: Tree[Any],
  ):
    self.module.assert_same_structure(tree0, tree1)


# TODO(epot): Also support custom `cls` like `register_pytree_node` in Jax
_SEQUENCE_TYPES = (list, tuple)
_MAPPING_TYPES = (dict, collections.abc.Mapping)
_ALL_TYPES = _SEQUENCE_TYPES + _MAPPING_TYPES


class Python(Backend):
  """Pure `Python` backend."""

  def import_module(self):
    raise RuntimeError('Python backend do not have module')

  def map(self, map_fn, *trees):
    tree0 = trees[0]
    if isinstance(tree0, _SEQUENCE_TYPES):
      new_items = (self.map(map_fn, *v) for v in zip(*trees))
      if epy.is_namedtuple(tree0):
        return type(tree0)(*new_items)
      else:
        return type(tree0)(new_items)
    elif isinstance(tree0, _MAPPING_TYPES):
      new_items = ((k, self.map(map_fn, *v)) for k, v in epy.zip_dict(*trees))
      if isinstance(tree0, collections.defaultdict):
        new_tree = type(tree0)(tree0.default_factory)
        new_tree.update(new_items)
        return new_tree
      else:
        return type(tree0)(new_items)
    else:  # leaf
      return map_fn(*trees)

  def flatten(self, tree):
    return list(self._flatten(tree)), tree

  def _flatten(self, tree):
    if isinstance(tree, _SEQUENCE_TYPES):
      return itertools.chain.from_iterable(self._flatten(v) for v in tree)
    elif isinstance(tree, _MAPPING_TYPES):
      return itertools.chain.from_iterable(
          self._flatten(v) for _, v in sorted(tree.items())
      )
    else:  # leaf
      return [tree]

  def unflatten(self, structure, flat_sequence):
    return self._unflatten(structure, iter(flat_sequence))

  def _unflatten(self, structure, flat_iter):
    """`unflatten` recursive implementation."""
    if isinstance(structure, _SEQUENCE_TYPES):
      new_items = (self._unflatten(v, flat_iter) for v in structure)
      if epy.is_namedtuple(structure):
        return type(structure)(*new_items)
      else:
        return type(structure)(new_items)
    elif isinstance(structure, _MAPPING_TYPES):
      # Flatten sort the keys, so reconstruct the ordered sorted
      ordered_items = {
          k: self._unflatten(v, flat_iter) for k, v in sorted(structure.items())
      }
      # Restore original dict order
      new_items = ((k, ordered_items[k]) for k in structure)

      if isinstance(structure, collections.defaultdict):
        new_tree = type(structure)(structure.default_factory)
        new_tree.update(new_items)
        return new_tree
      else:
        return type(structure)(new_items)
    else:  # leaf
      return next(flat_iter)

  def assert_same_structure(
      self,
      tree0: Tree[Any],
      tree1: Tree[Any],
  ):
    try:
      self._assert_same_structure(tree0, tree1)
    except Exception as e:  # pylint: disable=broad-except
      epy.reraise(e, prefix="The two structures don't match: ")

  def _assert_same_structure(
      self,
      tree0: Tree[Any],
      tree1: Tree[Any],
  ):
    """`assert_same_structure` recursive implementation."""
    if isinstance(tree0, _ALL_TYPES):
      if type(tree0) != type(tree1):  # pylint: disable=unidiomatic-typecheck
        raise ValueError(f'{type(tree0)} != {type(tree1)}')
    if isinstance(tree0, _SEQUENCE_TYPES):
      if len(tree0) != len(tree1):
        raise ValueError(f'{len(tree0)} != {len(tree1)}')
      for i, (v0, v1) in enumerate(zip(tree0, tree1)):
        try:
          self._assert_same_structure(v0, v1)
        except Exception as e:  # pylint: disable=broad-except
          epy.reraise(e, prefix=f'In {i}: ')
    elif isinstance(tree0, _MAPPING_TYPES):
      k0 = sorted(tree0)
      k1 = sorted(tree1)
      if k0 != k1:
        raise ValueError(f'dict keys do not match: {k0} != {k1}')
      # Flatten sort the keys, so reconstruct the ordered sorted
      for k, (v0, v1) in epy.zip_dict(tree0, tree1):
        try:
          self._assert_same_structure(v0, v1)
        except Exception as e:  # pylint: disable=broad-except
          epy.reraise(e, prefix=f'In {k}: ')
    else:  # leaf
      return