# Copyright 2024 The Flax 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.

# Copyright 2023 The Flax 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.
from __future__ import annotations

from collections.abc import MutableMapping
import typing as tp

import jax
import jax.tree_util as jtu

from flax.nnx.nnx import traversals
from flax.nnx.nnx import filterlib, reprlib
from flax.typing import PathParts

A = tp.TypeVar('A')
K = tp.TypeVar('K', bound=tp.Hashable)
V = tp.TypeVar('V')

FlatState = dict[PathParts, V]


class NestedStateRepr(reprlib.Representable):
  def __init__(self, state: State):
    self.state = state

  def __nnx_repr__(self):
    yield reprlib.Object('', value_sep=': ', start='{', end='}')

    for r in self.state.__nnx_repr__():
      if isinstance(r, reprlib.Object):
        continue
      yield r

  def __penzai_repr__(self, path, subtree_renderer):
    children = {}
    for k, v in self.state.items():
      if isinstance(v, State):
        v = NestedStateRepr(v)
      children[k] = v
    # Render as the dictionary itself at the same path.
    return subtree_renderer(children, path=path)


class State(MutableMapping[K, V], reprlib.Representable):
  """A pytree-like structure that contains a ``Mapping`` from strings or
  integers to leaves. A valid leaf type is either :class:`Variable`,
  ``jax.Array``, ``numpy.ndarray`` or nested ``State``'s. A ``State``
  can be generated by either calling :func:`split` or :func:`state` on
  the :class:`Module`."""

  def __init__(
    self,
    mapping: tp.Union[
      tp.Mapping[K, tp.Mapping | V],
      tp.Iterator[tuple[K, tp.Mapping | V]],
    ],
    /,
    *,
    _copy: bool = True,
  ):
    if _copy:
      _mapping = dict(mapping)
    else:
      if not isinstance(mapping, dict):
        raise ValueError(
          'Expected a dictionary when `_copy=False`, '
          f'got {type(mapping)} instead.'
        )
      _mapping = mapping

    if tp.TYPE_CHECKING:
      self._mapping = _mapping
    else:
      super().__setattr__('_mapping', _mapping)

  @property
  def raw_mapping(self) -> tp.Mapping[K, tp.Mapping[K, tp.Any] | V]:
    return self._mapping  # type: ignore

  def __contains__(self, key) -> bool:
    return key in self._mapping

  def __getitem__(self, key: K) -> State | V:  # type: ignore
    value = self._mapping[key]
    if isinstance(value, tp.Mapping):
      return State(value, _copy=False)
    return value

  def __getattr__(self, key: K) -> State | V:  # type: ignore[misc]
    if '_mapping' not in vars(self) or key not in self._mapping:
      raise AttributeError(f"No attribute '{key}' in State")
    return self[key]

  def __setitem__(self, key: K, value: State | V) -> None:
    if isinstance(value, State):
      self._mapping[key] = value._mapping
    else:
      self._mapping[key] = value

  __setattr__ = __setitem__  # type: ignore

  def __delitem__(self, key: K) -> None:
    del self._mapping[key]

  def __iter__(self) -> tp.Iterator[K]:
    return iter(self._mapping)

  def __len__(self) -> int:
    return len(self._mapping)

  def __nnx_repr__(self):
    yield reprlib.Object(type(self), value_sep=': ', start='({', end='})')

    for k, v in self.items():
      if isinstance(v, State):
        v = NestedStateRepr(v)
      yield reprlib.Attr(repr(k), v)

  def __penzai_repr__(self, path, subtree_renderer):
    from penzai.treescope import repr_lib as pz_repr_lib  # type: ignore[import-not-found,import-untyped]

    children = {}
    for k, v in self.items():
      if isinstance(v, State):
        v = NestedStateRepr(v)
      children[k] = v
    return pz_repr_lib.render_dictionary_wrapper(
      object_type=type(self),
      wrapped_dict=children,
      path=path,
      subtree_renderer=subtree_renderer,
    )

  def flat_state(self) -> FlatState[V]:
    return traversals.flatten_mapping(self._mapping)

  @classmethod
  def from_flat_path(cls, flat_state: tp.Mapping[PathParts, V], /) -> State:
    nested_state = traversals.unflatten_mapping(flat_state)
    return cls(nested_state)

  @tp.overload
  def split(self, first: filterlib.Filter, /) -> State[K, V]: ...

  @tp.overload
  def split(
    self,
    first: filterlib.Filter,
    second: filterlib.Filter,
    /,
    *filters: filterlib.Filter,
  ) -> tuple[State[K, V], ...]: ...

  def split(
    self, first: filterlib.Filter, /, *filters: filterlib.Filter
  ) -> tp.Union[State[K, V], tuple[State[K, V], ...]]:
    """Split a ``State`` into one or more ``State``'s. The
    user must pass at least one ``Filter`` (i.e. :class:`Variable`),
    and the filters must be exhaustive (i.e. they must cover all
    :class:`Variable` types in the ``State``).

    Example usage::

      >>> from flax import nnx

      >>> class Model(nnx.Module):
      ...   def __init__(self, rngs):
      ...     self.batchnorm = nnx.BatchNorm(2, rngs=rngs)
      ...     self.linear = nnx.Linear(2, 3, rngs=rngs)
      ...   def __call__(self, x):
      ...     return self.linear(self.batchnorm(x))

      >>> model = Model(rngs=nnx.Rngs(0))
      >>> state = nnx.state(model)
      >>> param, batch_stats = state.split(nnx.Param, nnx.BatchStat)

    Arguments:
      first: The first filter
      *filters: The optional, additional filters to group the state into mutually exclusive substates.
    Returns:
      One or more ``States`` equal to the number of filters passed.
    """
    filters = (first, *filters)
    *states_, rest = _split_state(self, *filters)

    if rest:
      raise ValueError(
        'Non-exhaustive filters, got a non-empty remainder: '
        f'{rest}.\nUse `...` to match all remaining elements.'
      )

    states: State | tuple[State, ...]
    if len(states_) == 1:
      states = states_[0]
    else:
      states = tuple(states_)
    return states  # type: ignore[bad-return-type]

  @tp.overload
  def filter(
    self,
    first: filterlib.Filter,
    /,
  ) -> State[K, V]: ...

  @tp.overload
  def filter(
    self,
    first: filterlib.Filter,
    second: filterlib.Filter,
    /,
    *filters: filterlib.Filter,
  ) -> tuple[State[K, V], ...]: ...

  def filter(
    self,
    first: filterlib.Filter,
    /,
    *filters: filterlib.Filter,
  ) -> tp.Union[State[K, V], tuple[State[K, V], ...]]:
    """Filter a ``State`` into one or more ``State``'s. The
    user must pass at least one ``Filter`` (i.e. :class:`Variable`).
    This method is similar to :meth:`split() <flax.nnx.nnx.State.state.split>`,
    except the filters can be non-exhaustive.

    Example usage::

      >>> from flax import nnx

      >>> class Model(nnx.Module):
      ...   def __init__(self, rngs):
      ...     self.batchnorm = nnx.BatchNorm(2, rngs=rngs)
      ...     self.linear = nnx.Linear(2, 3, rngs=rngs)
      ...   def __call__(self, x):
      ...     return self.linear(self.batchnorm(x))

      >>> model = Model(rngs=nnx.Rngs(0))
      >>> state = nnx.state(model)
      >>> param = state.filter(nnx.Param)
      >>> batch_stats = state.filter(nnx.BatchStat)
      >>> param, batch_stats = state.filter(nnx.Param, nnx.BatchStat)

    Arguments:
      first: The first filter
      *filters: The optional, additional filters to group the state into mutually exclusive substates.
    Returns:
      One or more ``States`` equal to the number of filters passed.
    """
    *states_, _rest = _split_state(self, first, *filters)

    assert len(states_) == len(filters) + 1

    states: State | tuple[State, ...]
    if len(states_) == 1:
      states = states_[0]
    else:
      states = tuple(states_)

    return states  # type: ignore[bad-return-type]

  @staticmethod
  def merge(state: State[K, V], /, *states: State[K, V]) -> State[K, V]:
    """The inverse of :meth:`split() <flax.nnx.nnx.State.state.split>`.

    ``merge`` takes one or more ``State``'s and creates
    a new ``State``.

    Example usage::

      >>> from flax import nnx
      >>> import jax.numpy as jnp

      >>> class Model(nnx.Module):
      ...   def __init__(self, rngs):
      ...     self.batchnorm = nnx.BatchNorm(2, rngs=rngs)
      ...     self.linear = nnx.Linear(2, 3, rngs=rngs)
      ...   def __call__(self, x):
      ...     return self.linear(self.batchnorm(x))

      >>> model = Model(rngs=nnx.Rngs(0))
      >>> params, batch_stats = nnx.state(model, nnx.Param, nnx.BatchStat)
      >>> params.linear.bias.value += 1

      >>> state = nnx.State.merge(params, batch_stats)
      >>> nnx.update(model, state)
      >>> assert (model.linear.bias.value == jnp.array([1, 1, 1])).all()

    Args:
      state: A ``State`` object.
      *states: Additional ``State`` objects.
    Returns:
      The merged ``State``.
    """
    states = (state, *states)

    if len(states) == 1:
      return states[0]

    new_state: FlatState[V] = {}

    for state in states:
      new_state.update(state.flat_state())  # type: ignore[attribute-error] # pytype is wrong here

    return State.from_flat_path(new_state)

  def __or__(self, other: State[K, V]) -> State[K, V]:
    if not other:
      return self
    return State.merge(self, other)

  def __sub__(self, other: State[K, V]) -> State[K, V]:
    if not other:
      return self

    self_flat = self.flat_state()
    other_flat = other.flat_state()
    diff = {k: v for k, v in self_flat.items() if k not in other_flat}

    return State.from_flat_path(diff)


def _state_flatten_with_keys(x: State):
  items = sorted(x._mapping.items())
  children = tuple((jtu.DictKey(key), value) for key, value in items)
  return children, tuple(key for key, _ in items)


def _state_unflatten(
  static: tuple[K, ...],
  leaves: tuple[V, ...] | tuple[dict[K, V]],
):
  return State(zip(static, leaves))


jax.tree_util.register_pytree_with_keys(
  State,
  _state_flatten_with_keys,
  _state_unflatten,  # type: ignore[arg-type]
)


def _split_state(
  state: State[K, V],
  *filters: filterlib.Filter,
) -> tuple[State[K, V], ...]:
  for i, filter_ in enumerate(filters):
    if filter_ in (..., True) and i != len(filters) - 1:
      remaining_filters = filters[i + 1 :]
      if not all(f in (..., True) for f in remaining_filters):
        raise ValueError(
          '`...` or `True` can only be used as the last filters, '
          f'got {filter_} it at index {i}.'
        )
  predicates = tuple(map(filterlib.to_predicate, filters))

  flat_state = state.flat_state()

  # we have n + 1 states, where n is the number of predicates
  # the last state is for values that don't match any predicate
  flat_states: tuple[FlatState[V], ...] = tuple(
    {} for _ in range(len(predicates) + 1)
  )

  for path, value in flat_state.items():
    for i, predicate in enumerate(predicates):
      if predicate(path, value):
        flat_states[i][path] = value  # type: ignore[index] # mypy is wrong here?
        break
    else:
      # if we didn't break, set leaf to last state
      flat_states[-1][path] = value  # type: ignore[index] # mypy is wrong here?

  return tuple(State.from_flat_path(flat_state) for flat_state in flat_states)
