"""A module defining sources of randomness for TFHE."""

import abc
from collections.abc import Sequence
import functools
import itertools
import operator
import random
from typing import Any, Callable

import jax
import jax.numpy as jnp


def _shape_generator(
    generator: Callable[[], Any],
    shape: Sequence[int],
    dtype: jnp.dtype = jnp.uint32,
) -> jax.Array:
  """Generates an array of data generated by a given function.

  Utility function to generate a JAX ndarray of values given by repeatedly
  invoking the passed in generator function.  Used for RNGs that run on CPU
  serially.

  Args:
    generator: A parameterless function that returns data
    shape: The shape of the JAX ndarray to generate
    dtype: The data type of the JAX ndarray to generate

  Returns:
    A JAX ndarray of the given shape, using data generated by repeatedly calling
    the given generator function.  The data is cast to the given dtype if
    possible.
  """
  if any(d < 0 for d in shape):
    raise ValueError(f"Shape dimensions must be positive: got {shape}.")
  num_elements = functools.reduce(operator.mul, shape, 1)
  return jnp.reshape(
      jnp.array([generator() for _ in range(num_elements)], dtype), shape
  )


class RandomSource(abc.ABC):
  """An interface for a source of randomness."""

  @abc.abstractmethod
  def uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    raise NotImplementedError("Subclasses *must* implement `uniform`.")

  @abc.abstractmethod
  def rounded_normal(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    raise NotImplementedError("Subclasses *must* implement `rounded_normal`.")

  @abc.abstractmethod
  def sk_uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    """Generates random numbers in {0, 1}, for use in generating secret keys."""
    raise NotImplementedError("Subclasses *must* implement `sk_uniform`.")


class CycleRng(RandomSource):
  """An RNG that cycles through a deterministic set of values."""

  def __init__(self, const_normal_noise: int = 0) -> None:
    self.const_normal_noise = const_normal_noise
    # This results in `uniform` having binary-valued outputs
    self.uniform_data = tuple(map(int, bin(13**39**2 // 53)[2:]))
    self.noise = iter(itertools.cycle(self.uniform_data))

  def uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return _shape_generator(lambda: next(self.noise), shape, dtype)

  def sk_uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return self.uniform(shape, dtype=dtype)

  def rounded_normal(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return _shape_generator(lambda: self.const_normal_noise, shape, dtype)


class SystemRandomSource(RandomSource):
  """A cryptographically secure random source that runs on CPU serially.

  Implemented as a wrapper on top of random.SystemRandom().
  """

  def __init__(
      self,
      uniform_bounds: tuple[int, int] = (0, 2**32 - 1),
      normal_std: int = 1,
  ) -> None:
    self.uniform_bounds = uniform_bounds
    self.normal_std = normal_std
    self.rng = random.SystemRandom()

  def uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return _shape_generator(
        lambda: self.rng.randint(*self.uniform_bounds), shape, dtype
    )

  def rounded_normal(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return _shape_generator(
        lambda: round(self.rng.normalvariate(0, self.normal_std)), shape, dtype
    )

  def sk_uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return _shape_generator(lambda: self.rng.randint(0, 1), shape, dtype)


class PseudorandomSource(RandomSource):
  """An insecure random source based on the Python stdlib."""

  def __init__(
      self,
      uniform_bounds: tuple[int, int] = (0, 2**32 - 1),
      normal_std: int = 1,
      seed: int = 1,
  ) -> None:
    self.uniform_bounds = uniform_bounds
    self.normal_std = normal_std
    self.rng = random.Random(seed)

  def uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return _shape_generator(
        lambda: self.rng.randint(*self.uniform_bounds), shape, dtype
    )

  def rounded_normal(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return _shape_generator(
        lambda: round(self.rng.normalvariate(0, self.normal_std)), shape, dtype
    )

  def sk_uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return _shape_generator(lambda: self.rng.randint(0, 1), shape, dtype)


class NormalOnlyRng(RandomSource):
  """A test-only RNG.

  Uniform returns zeros, rounded_normal returns real RNG samples.
  """

  def __init__(self, normal_std: int = 1, seed: int = 1) -> None:
    self.normal_std = normal_std
    self.rng = random.Random(seed)

  def uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return jnp.zeros(shape, dtype=dtype)

  def rounded_normal(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return _shape_generator(
        lambda: round(self.rng.normalvariate(0, self.normal_std)), shape, dtype
    )

  def sk_uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return _shape_generator(lambda: self.rng.randint(0, 1), shape, dtype)


class ConstantUniformRng(RandomSource):
  """A test-only RNG that sets the output of uniform() to a fixed value."""

  def __init__(
      self, const_uniform: int = 1, normal_std: int = 1, seed: int = 1
  ) -> None:
    self.const_uniform = const_uniform
    self.normal_std = normal_std
    self.rng = random.Random(seed)

  def uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return jnp.ones(shape, dtype=dtype) * self.const_uniform

  def rounded_normal(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return _shape_generator(
        lambda: round(self.rng.normalvariate(0, self.normal_std)), shape, dtype
    )

  def sk_uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return _shape_generator(lambda: self.rng.randint(0, 1), shape, dtype)


class ZeroRng(RandomSource):
  """An random generator that always outputs zero."""

  def uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return jnp.zeros(shape, dtype=dtype)

  def rounded_normal(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return jnp.zeros(shape, dtype=dtype)

  def sk_uniform(
      self, shape: Sequence[int] = (1,), dtype: jnp.dtype = jnp.uint32
  ) -> jax.Array:
    return jnp.zeros(shape, dtype=dtype)


_varying_params = (
    ((0, 1), 0),
    ((0, 1), 1),
    ((0, 2**5), 1),
    ((0, 2**15), 1),
    ((0, 2**30), 1),
    ((0, 2**32 - 1), 1),
    ((0, 2**5), 2**5),
    ((0, 2**15), 2**5),
    ((0, 2**30), 2**5),
    ((0, 2**32 - 1), 2**5),
    ((0, 2**32 - 1), 2**7),
)
VARYING_MAGNITUDE_TEST_RNGS = [
    PseudorandomSource(uniform_bounds=uniform_bounds, normal_std=std, seed=1)
    for (uniform_bounds, std) in _varying_params
]

ALL_RNGS = [
    ConstantUniformRng,
    CycleRng,
    NormalOnlyRng,
    PseudorandomSource,
    SystemRandomSource,
    ZeroRng,
]
