from __future__ import annotations

import numpy as np

from zen_fronts.utils.boundary import reflect01


def _reflect01_vec(x: np.ndarray) -> np.ndarray:
    """Vectorized definition matching `zen_fronts.utils.boundary.reflect01`."""

    x = np.abs(x.astype(np.float64))
    x = np.mod(x, 2.0)
    x = np.where(x > 1.0, 2.0 - x, x)
    return x


def test_reflect01_matches_vectorized_definition() -> None:
    rng = np.random.default_rng(0)
    vals = rng.normal(size=200).astype(np.float64) * 3.0
    got = np.array([reflect01(float(v)) for v in vals], dtype=np.float64)
    exp = _reflect01_vec(vals)
    assert np.allclose(got, exp)


def test_clip_creates_sticky_boundary_mass_reflect_does_not() -> None:
    """Regression test for the demo's "cluster at 0/1" artifact.

    Simple clipping produces *exact* 0.0/1.0 values whenever a mutation crosses
    the boundary. Over many steps this creates noticeable mass at boundaries.
    Reflection avoids that by folding values back into the box.
    """

    rng = np.random.default_rng(123)

    n = 20_000
    steps = 40
    sigma = 0.25

    # Keep away from exact endpoints to make the "reflection has ~0 exact boundaries" check stable.
    x0 = (0.01 + 0.98 * rng.random(n)).astype(np.float64)
    eps = rng.normal(0.0, sigma, size=(steps, n)).astype(np.float64)

    x_clip = x0.copy()
    x_reflect = x0.copy()
    for t in range(steps):
        x_clip = np.clip(x_clip + eps[t], 0.0, 1.0)
        x_reflect = _reflect01_vec(x_reflect + eps[t])

    boundary_mass_clip = float(np.mean((x_clip == 0.0) | (x_clip == 1.0)))
    boundary_mass_reflect = float(np.mean((x_reflect == 0.0) | (x_reflect == 1.0)))

    # Clipping should produce a non-trivial amount of exact boundary points.
    assert boundary_mass_clip > 0.01
    # Reflection should almost never land exactly on boundaries.
    assert boundary_mass_reflect < 1e-8
