# Copyright 2025 The Levanter Authors
# SPDX-License-Identifier: Apache-2.0

import typing

import jax.numpy as jnp
import jax.random as jrandom
import numpy as np
from jaxtyping import PRNGKeyArray


PermType = typing.Literal["feistel", "linear"]


class Permutation(typing.Protocol):
    """A protocol for a pseudo-random permutation.

    This protocol defines the interface for a pseudo-random permutation. The permutation is defined by a function
    that takes an integer or an array of integers and returns the permuted value(s).
    """

    @typing.overload
    def __call__(self, indices: int) -> int:
        """Apply the permutation to a single integer.

        Args:
            indices: An integer to be permuted.

        Returns:
            The permuted value.
        """
        ...

    @typing.overload
    def __call__(self, indices: np.ndarray) -> np.ndarray:
        """Apply the permutation to an array of integers.

        Args:
            indices: An array of integers to be permuted.

        Returns:
            The permuted value(s).
        """
        ...

    def __call__(self, indices: int | np.ndarray) -> int | np.ndarray:
        """Apply the permutation to the given indices.

        Args:
            indices: An integer or an array of integers to be permuted.

        Returns:
            The permuted value(s).
        """
        pass

    @staticmethod
    def make(kind: PermType, length: int, prng_key: PRNGKeyArray, **kwargs) -> "Permutation":
        """Create a permutation of the given kind.

        Args:
            kind: The kind of permutation to create. Can be "feistel" or "linear".
            length: The length of the permutation.
            prng_key: A JAX PRNG key for random number generation.
            **kwargs: Additional arguments for the specific permutation type.

        Returns:
            An instance of the specified permutation type.
        """
        if kind == "feistel":
            return FeistelPermutation(length, prng_key, **kwargs)
        elif kind == "linear":
            return LcgPermutation(length, prng_key)
        else:
            raise ValueError(f"Unknown permutation kind: {kind}")


class LcgPermutation(Permutation):
    # Pseudo-Random Permutation Code
    """A stateless pseudo-random permutation.

    This class generates a pseudo-random permutation of a given length. The permutation is generated using a PRNG
    with a fixed key. The permutation is generated by finding a random `a` and `b` such that `gcd(a, length) == 1` and
    then computing the permutation as `p(x) = (a * x + b) % length`.
    """

    def __init__(self, length, prng_key):
        self.length = length
        # Convert jax.random.PRNGKey to numpy.random.Generator
        self.rng = np.random.Generator(np.random.PCG64(jrandom.randint(prng_key, (), 0, 2**30).item()))
        self.a, self.b = self._generate_permutation_params()  # Generate a and b in init

    def _generate_permutation_params(self):
        length = self.length
        rng = self.rng

        if length == 1:
            return 1, 0

        while True:
            a = rng.integers(1, length)
            if np.gcd(a, length) == 1:
                break

        b = rng.integers(0, length)  # b can be in [0, length-1]
        return a, b

    @typing.overload
    def __call__(self, indices: int) -> int: ...

    @typing.overload
    def __call__(self, indices: np.ndarray) -> np.ndarray: ...

    def __call__(self, indices):
        a = self.a
        b = self.b
        length = self.length

        was_int = False
        if isinstance(indices, np.ndarray | jnp.ndarray):
            if np.any(indices < 0) or np.any(indices >= self.length):
                raise IndexError(f"index {indices} is out of bounds for length {self.length}")

            if isinstance(indices, jnp.ndarray):
                indices = np.array(indices)
        else:
            if indices < 0 or indices >= self.length:
                raise IndexError(f"index {indices} is out of bounds for length {self.length}")

            indices = np.array(indices)
            was_int = True

        out = (a * indices.astype(np.uint64) + b) % length  # Compute permutation on-the-fly
        out = out.astype(indices.dtype)

        if was_int:
            return int(out)
        return out


_MAX_FEISTEL_ROUNDS = 1000


class FeistelPermutation(Permutation):
    """A pseudo-random permutation for an arbitrary domain using a Feistel network and cycle-walking.

    The domain [0, length) is embedded into a power-of-two domain [0, m) where m = next_power_of_two(length).
    A Feistel network is applied on the bit-level representation, and if the result falls outside [0, length),
    the network is reapplied (cycle-walking) until the result is in-range.
    """

    def __init__(self, length: int, prng_key, rounds: int = 5):
        if rounds < 1:
            raise ValueError("Number of rounds must be at least 1.")

        self.length = length
        self.m = next_power_of_two(length)  # m is a power-of-two >= length
        self.rounds = rounds
        self.bits = self.m.bit_length() - 1  # because m == 1 << bits
        # Split the bits into two halves.
        self.left_bits = self.bits // 2
        self.right_bits = self.bits - self.left_bits
        self.R_mask = np.uint64((1 << self.right_bits) - 1)

        # Convert JAX PRNG to numpy RNG and generate round keys in range [0, 2^(right_bits))
        self.rng = np.random.Generator(np.random.PCG64(jrandom.randint(prng_key, (), 0, 2**30).item()))
        self.keys = self.rng.integers(0, 1 << self.right_bits, size=rounds, dtype=np.uint64)

    def _F(self, right: np.ndarray, key: np.uint64) -> np.ndarray:
        """A simple round function that mixes the right half.

        Operates modulo 2^(right_bits).
        """
        masked_right = right & self.R_mask
        return (masked_right * np.uint64(2654435761) + key) & self.R_mask

    def _feistel(self, x: np.ndarray) -> np.ndarray:
        """Apply the Feistel network to x, where x is assumed to be in [0, m)."""
        # Split x into left and right parts.
        right = x & self.R_mask
        left = x >> np.uint64(self.right_bits)
        for key in self.keys:
            new_left = right
            new_right = left ^ self._F(right, key)
            left, right = new_left, new_right

        out = (left << np.uint64(self.right_bits)) | right
        # np will convert 0-dim arrays to scalars, so we need to ensure the output is an array.
        if not isinstance(out, np.ndarray):
            out = np.array(out)

        return out

    @typing.overload
    def __call__(self, indices: int) -> int: ...

    @typing.overload
    def __call__(self, indices: np.ndarray | jnp.ndarray) -> np.ndarray: ...

    def __call__(self, indices):

        was_int = False
        # Ensure we work with a numpy array of uint64.
        if isinstance(indices, jnp.ndarray):
            indices = np.array(indices)

        if not isinstance(indices, np.ndarray):
            if indices < 0 or indices >= self.length:
                raise IndexError(f"Index {indices} is out of bounds for length {self.length}")
            indices = np.array(indices, dtype=np.uint64)
            was_int = True
        else:
            indices = indices.astype(np.uint64)

        if np.any(indices < 0) or np.any(indices >= self.length):
            raise IndexError(f"Index {indices} is out of bounds for length {self.length}")

        if self.length == 1:
            return 0 if was_int else np.zeros_like(indices)

        # Embed indices into the larger domain.
        x = indices
        # Apply cycle-walking: reapply the Feistel network until the result is in [0, length).
        out = self._feistel(x)
        mask = out >= self.length
        cycles = 0
        while np.any(mask):
            if cycles >= _MAX_FEISTEL_ROUNDS:
                raise RuntimeError(f"Feistel cycle-walking exceeded {cycles} iterations. This should not happen.")
            n = self._feistel(out[mask])
            out[mask] = n
            mask = out >= self.length
            cycles += 1

        if was_int:
            return int(out)
        return out


def next_power_of_two(n: int) -> int:
    return 1 << (n - 1).bit_length()


if __name__ == "__main__":
    # test extremely large permutation
    length = 2**33 + 5
    prng_key = jrandom.PRNGKey(0)
    permutation = FeistelPermutation(length, prng_key)

    import time

    time_in = time.time()

    # test single index
    for i in np.random.randint(0, length, 1000):
        # print(i, permutation(i))
        permutation(i)

    # test array of indices
    for batch in np.random.randint(0, length, (1024, 1024)):
        # print(batch, permutation(batch))
        out = permutation(batch)
        assert np.all(out < length)

    print("Time taken:", time.time() - time_in)
