"""Counter-based random numbers shared with the browser version (web/feedloop.js).

Every random number is a hash of (seed, stream, counter), so there is no state to keep in
step: the Python code can draw a whole block with numpy and the JavaScript can draw the
same numbers one at a time, and both get identical values. It also gives "common random
numbers": the click coin for (round t, user u, slot r) is the same whatever policy is being
run, so two runs from the same seed differ only where the policies actually chose
differently.

Streams (see STREAMS): 0 builds the world, 1 is the click coins, 2 the Thompson noise,
3 and 4 the epsilon-greedy coins and item picks, 5 the tie-breaking jitter, 6 the ALS start.
"""
from __future__ import annotations

import numpy as np

STREAMS = {"world": 0, "click": 1, "thompson": 2, "eps_coin": 3, "eps_item": 4, "jitter": 5, "als": 6}

_M1 = np.uint32(0x21F0AAAD)
_M2 = np.uint32(0x735A2D97)
_M3 = np.uint32(0x9E3779B9)
_M4 = np.uint32(0x85EBCA6B)
_M5 = np.uint32(0xC2B2AE35)


def _mix(x: np.ndarray) -> np.ndarray:
    """A 32-bit finaliser (two rounds of murmur-style mixing). x is a uint32 array."""
    x = x ^ (x >> np.uint32(16))
    x = x * _M1
    x = x ^ (x >> np.uint32(15))
    x = x * _M2
    x = x ^ (x >> np.uint32(15))
    x = x ^ (x >> np.uint32(16))
    x = x * _M4
    x = x ^ (x >> np.uint32(13))
    x = x * _M5
    x = x ^ (x >> np.uint32(16))
    return x


def stream_key(seed: int, stream: int) -> int:
    k = _mix(np.array([np.uint32((seed * 0x9E3779B9 + stream * 0x85EBCA6B + 0x1234567) & 0xFFFFFFFF)], dtype=np.uint32))
    return int(k[0])


def hash_u32(key: int, counters) -> np.ndarray:
    c = np.asarray(counters, dtype=np.uint32)
    with np.errstate(over="ignore"):
        return _mix((c + _M3) ^ np.uint32(key))


def uniforms(key: int, counters) -> np.ndarray:
    """Uniforms in (0, 1): (h + 0.5) / 2^32, so log(u) is always finite."""
    return (hash_u32(key, counters).astype(np.float64) + 0.5) / 4294967296.0


def normals(key: int, counters) -> np.ndarray:
    """Standard normals by Box-Muller; counter n uses uniforms 2n and 2n+1 of the stream."""
    c = np.asarray(counters, dtype=np.uint32)
    u1 = uniforms(key, c * np.uint32(2))
    u2 = uniforms(key, c * np.uint32(2) + np.uint32(1))
    return np.sqrt(-2.0 * np.log(u1)) * np.cos(2.0 * np.pi * u2)


class Seq:
    """A sequential uniform stream (counter increments on every call), for building the world."""

    def __init__(self, seed: int, stream: int = 0):
        self.key = stream_key(seed, stream)
        self.n = 0

    def uniform(self) -> float:
        u = uniforms(self.key, [self.n])[0]
        self.n += 1
        return float(u)

    def block(self, count: int) -> np.ndarray:
        u = uniforms(self.key, np.arange(self.n, self.n + count, dtype=np.uint32))
        self.n += count
        return u