# Copyright 2021 The NetKet Authors - All rights reserved.
#
# 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.

import abc

from typing import Tuple, Optional, Union, Iterable

import jax.numpy as jnp
import numpy as np

from netket.utils import warn_deprecation

max_states = np.iinfo(np.int32).max
"""int: Maximum number of states that can be indexed"""


class NoneTypeT:
    pass


NoneType = NoneTypeT()


legacy_warn_str = (
    "This choice of arguments for `hilbert.random_state` is deprecated and "
    "will be removed in a future version.\n"
    "The new syntax is\n"
    "    hilbert.random_state(prngkey, size=1, dtype=jnp.float32)\n"
    "and, like in `jax.random`, the first argument is mandatory and must be a valid jax `PRNGKey`."
    "Results may differ between the states generated by the legacy, "
    "deprecated code and the new one. Note that the old code does not "
    "support defining custom methods."
)


class AbstractHilbert(abc.ABC):
    """Abstract class for NetKet hilbert objects.

    This class definese the common interface that can be used to
    interact with hilbert spaces.

    Hilbert Spaces are immutable.
    """

    def __init__(self):
        self._hash = None

    @property
    @abc.abstractmethod
    def size(self) -> int:
        r"""The number number of degrees of freedom in the basis of this
        Hilbert space."""
        raise NotImplementedError()  # pragma: no cover

    # after removing legacy:
    # signature must be the following
    # def random_state(self, key, size=None, dtype=np.float32):
    def random_state(
        self,
        key=NoneType,
        size: Optional[int] = NoneType,
        dtype=np.float32,
        out: Optional[np.ndarray] = None,
        rgen=None,
    ) -> jnp.ndarray:
        r"""Generates either a single or a batch of uniformly distributed random states.
        Runs as :code:`random_state(self, key, size=None, dtype=np.float32)` by default.

        Args:
            key: rng state from a jax-style functional generator.
            size: If provided, returns a batch of configurations of the form
                  :code:`(size, N)` if size is an integer or :code:`(*size, N)` if it is
                  a tuple and where :math:`N` is the Hilbert space size.
                  By default, a single random configuration with shape
                  :code:`(#,)` is returned.
            dtype: DType of the resulting vector.
            out: Deprecated. Will be removed in v3.1
            rgen: Deprecated. Will be removed in v3.1

        Returns:
            A state or batch of states sampled from the uniform distribution on the
            hilbert space.

        Example:

            >>> import netket, jax
            >>> hi = netket.hilbert.Qubit(N=2)
            >>> k1, k2 = jax.random.split(jax.random.PRNGKey(1))
            >>> print(hi.random_state(key=k1))
            [1. 0.]
            >>> print(hi.random_state(key=k2, size=2))
            [[0. 0.]
             [0. 1.]]
        """
        # legacy support
        # TODO: Remove in 3.1
        # if no positional arguments, and key is unspecified -> legacy
        if key is NoneType:
            warn_deprecation(legacy_warn_str)
            # legacy sure
            if size is NoneType:
                return self._random_state_legacy(size=None, out=out, rgen=rgen)
            else:
                return self._random_state_legacy(size=size, out=out, rgen=rgen)
        elif isinstance(key, tuple) or isinstance(key, int) and size is NoneType:
            # if one positional argument legacy typee...
            warn_deprecation(legacy_warn_str)
            return self._random_state_legacy(size=key, out=out, rgen=rgen)
        else:
            from netket.hilbert import random

            size = size if size is not NoneType else None

            return random.random_state(self, key, size, dtype=dtype)

    def ptrace(self, sites: Union[int, Iterable]) -> "AbstractHilbert":
        """Returns the hilbert space without the selected sites.

        Not all hilbert spaces support this operation.

        Args:
            sites: a site or list of sites to trace away

        Returns:
            The partially-traced hilbert space. The type of the resulting hilbert space
            might be different from the starting one.
        """
        raise NotImplementedError("Ptrace not implemented for this hilbert space type.")

    @property
    def is_indexable(self) -> bool:
        """Whever the space can be indexed with an integer"""
        return False

    @property
    @abc.abstractmethod
    def _attrs(self) -> Tuple:
        """
        Tuple of hashable attributs, used to compute the immutable
        hash of this Hilbert space
        """

    def __eq__(self, other) -> bool:
        if isinstance(other, type(self)):
            return self._attrs == other._attrs

        return False

    def __hash__(self):
        if self._hash is None:
            self._hash = hash(self._attrs)

        return self._hash
