tests/test_vectorized_cleanroom.py

from __future__ import annotations

import importlib.util
import itertools
import json
from pathlib import Path

import numpy as np
import pytest

from cascade.interface import DataGenerator
from cascade.interface.generator import check_series
from cascade.interface.static_guard import scan_file
from cascade.interface.validation import (
    check_config,
    check_repo_layout,
    check_requirements_hash_locked,
)
from cascade.miner.verify import verify_repo
from cascade.shared.config import load_chain_config


ROOT = Path(__file__).resolve().parents[1]
CASCADE = Path("/root/91/cascade")
HEAT_BUDGET = 6_660_000_000
FINAL_BUDGET = 39_960_000_000


def _module():
    spec = importlib.util.spec_from_file_location(
        "fullclock_zero_cleanroom_generator", ROOT / "generator.py"
    )
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def _rows(seed: int, n: int, root: Path = ROOT):
    return list(_module().Generator(str(root), seed=seed).generate(n))


def test_contract_exact_partial_chunk_and_determinism():
    module = _module()
    gen = module.Generator(str(ROOT), seed=73)
    assert isinstance(gen, DataGenerator)
    n = gen._chunk_size + 3
    first = list(gen.generate(n))
    second = list(module.Generator(str(ROOT), seed=73).generate(n))
    assert len(first) == len(second) == n
    for index, (left, right) in enumerate(zip(first, second, strict=True)):
        check_series(
            left, min_length=64, max_length=4096, max_channels=1, index=index
        )
        assert left.shape == (4096,)
        assert left.flags.c_contiguous
        assert np.array_equal(left, right)
        assert np.isfinite(left).all()
        assert np.ptp(left) > 0.0
    keys = {row.tobytes() for row in first}
    assert len(keys) == n


def test_huge_request_is_bounded_by_fixed_chunk(monkeypatch):
    module = _module()
    gen = module.Generator(str(ROOT), seed=91)
    observed: list[int] = []
    original = gen._family_batch

    def wrapped(rng, family, count):
        observed.append(int(count))
        return original(rng, family, count)

    monkeypatch.setattr(gen, "_family_batch", wrapped)
    rows = list(itertools.islice(gen.generate(624_375_002), gen._chunk_size + 1))
    assert len(rows) == gen._chunk_size + 1
    assert observed and max(observed) <= gen._chunk_size


def test_seed_changes_corpus():
    left = _rows(10, 32)
    right = _rows(11, 32)
    assert any(
        not np.array_equal(a, b)
        for a, b in zip(left, right, strict=True)
    )


def test_rng_resets_partial_restart_and_global_state_is_untouched():
    module = _module()
    gen = module.Generator(str(ROOT), seed=104)
    first = list(itertools.islice(gen.generate(104_062_502), 67))
    second = list(itertools.islice(gen.generate(104_062_502), 67))
    assert all(
        np.array_equal(left, right)
        for left, right in zip(first, second, strict=True)
    )
    np.random.seed(611)
    expected = np.random.random(8)
    np.random.seed(611)
    list(itertools.islice(gen.generate(104_062_502), 5))
    assert np.array_equal(np.random.random(8), expected)


def test_component_lanes_cannot_advance_each_other():
    gen = _module().Generator(str(ROOT), seed=902)
    direct = gen._rng(0x0401)
    expected = [gen._source_batch(direct, 1, size)[0] for size in (2, 5, 3)]

    stock = gen._rng(0x0400)
    market = gen._rng(0x0401)
    actual = []
    for index, size in enumerate((2, 5, 3)):
        gen._source_batch(stock, 0, index + 1)
        actual.append(gen._source_batch(market, 1, size)[0])
    assert all(
        np.array_equal(left, right)
        for left, right in zip(actual, expected, strict=True)
    )


def test_zero_and_enabled_source_configs_share_every_unselected_base_row(tmp_path):
    cfg = json.loads((ROOT / "config.json").read_text(encoding="utf-8"))
    cfg["source_rate_start"] = {
        "stock_flow": 0.004,
        "market": 0.0,
        "weather_load": 0.003,
        "weekly_count": 0.003,
    }
    cfg["source_rate_final"] = {
        "stock_flow": 0.032,
        "market": 0.024,
        "weather_load": 0.016,
        "weekly_count": 0.008,
    }
    (tmp_path / "config.json").write_text(json.dumps(cfg), encoding="utf-8")
    module = _module()
    disabled = module.Generator(str(ROOT), seed=814)
    enabled = module.Generator(str(tmp_path), seed=814)
    n = 512
    left = list(enabled.generate(n))
    right = list(disabled.generate(n))
    kinds = enabled._select_sources(enabled._rng(0x0300), 0, n, n)
    assert np.any(kinds >= 0)
    assert all(
        np.array_equal(left[index], right[index])
        for index in np.flatnonzero(kinds < 0)
    )


def test_every_family_and_source_is_finite_nonconstant_and_integer_when_marked():
    module = _module()
    gen = module.Generator(str(ROOT), seed=512)
    for family in range(module.FAMILY_COUNT):
        values, integer = gen._family_batch(
            gen._rng(0x7000 + family), family, 8
        )
        assert values.shape == (8, 4096)
        assert np.isfinite(values).all()
        assert np.all(np.ptp(values, axis=1) > 0.0)
        if integer:
            assert np.array_equal(values, np.rint(values))
    for kind in range(4):
        values, integer = gen._source_batch(gen._rng(0x7100 + kind), kind, 32)
        assert values.shape == (32, 4096)
        assert np.isfinite(values).all()
        assert np.all(np.ptp(values, axis=1) > 0.0)
        integer_rows = np.full(32, integer) if np.ndim(integer) == 0 else integer
        if np.any(integer_rows):
            assert np.array_equal(
                values[integer_rows], np.rint(values[integer_rows])
            )


def test_market_source_returns_per_row_integer_mask():
    gen = _module().Generator(str(ROOT), seed=902)
    values, integer_rows = gen._source_batch(gen._rng(0x999), 1, 64)
    assert isinstance(integer_rows, np.ndarray)
    assert integer_rows.dtype == np.bool_
    assert integer_rows.shape == (64,)
    assert np.any(integer_rows) and np.any(~integer_rows)
    assert np.array_equal(
        values[integer_rows], np.rint(values[integer_rows])
    )
    assert np.any(values[~integer_rows] != np.rint(values[~integer_rows]))


def test_pointwise_sanitizer_and_linear_filters_are_prefix_causal():
    module = _module()
    prefix = np.asarray([1.2, np.nan, 5.7, np.inf, -4.2])
    left = module._causal_sanitize(
        np.concatenate((prefix, np.ones(10))), cap=100.0
    )
    right = module._causal_sanitize(
        np.concatenate((prefix, np.full(10, -99.0))), cap=100.0
    )
    assert np.array_equal(left[: prefix.size], right[: prefix.size])

    rng = np.random.default_rng(8)
    forcing = rng.normal(size=(7, 256))
    altered = forcing.copy()
    altered[:, 111:] = rng.normal(size=altered[:, 111:].shape)
    initial = rng.normal(size=7)
    phi = rng.uniform(-0.3, 0.9, size=7)
    a = module._ar1_rows(forcing, phi, initial)
    b = module._ar1_rows(altered, phi, initial)
    assert np.array_equal(a[:, :111], b[:, :111])


@pytest.mark.parametrize(
    ("budget", "upper", "consumed", "terminal_row"),
    [
        (
            HEAT_BUDGET,
            104_062_502,
            1_625_977,
            1_056_885,
        ),
        (
            FINAL_BUDGET,
            624_375_002,
            9_755_860,
            6_341_309,
        ),
    ],
)
def test_exact_public_source_and_base_clock_contract(
    budget, upper, consumed, terminal_row
):
    gen = _module().Generator(str(ROOT), seed=4)
    assert budget // 64 + 2 == upper
    assert gen._estimated_consumed_rows(upper) == consumed
    assert (budget + 4095) // 4096 == consumed
    assert gen._progress(terminal_row - 1, upper) < 1.0
    assert gen._progress(terminal_row, upper) == 1.0
    assert gen._source_progress(terminal_row - 1, upper) < 1.0
    assert gen._source_progress(terminal_row, upper) == 1.0
    assert gen._progress(consumed - 1, upper) == 1.0
    assert gen._progress(terminal_row, upper) == gen._source_progress(
        terminal_row, upper
    )
    assert np.array_equal(gen._source_rates(0, upper), gen._rate_start)
    assert np.array_equal(
        gen._source_rates(terminal_row, upper), gen._rate_final
    )
    probe = np.asarray([0, 1, terminal_row - 1, terminal_row])
    source_vector = gen._source_progress_chunk(0, terminal_row + 1, upper)[
        probe
    ]
    base_vector = gen._progress_chunk(0, terminal_row + 1, upper)[probe]
    source_scalar = np.asarray(
        [gen._source_progress(int(row), upper) for row in probe]
    )
    base_scalar = np.asarray([gen._progress(int(row), upper) for row in probe])
    assert np.array_equal(source_vector, source_scalar)
    assert np.array_equal(base_vector, base_scalar)
    assert np.array_equal(base_vector, source_vector)


def test_zero_overlay_component_mapping_and_rates():
    cfg = json.loads((ROOT / "config.json").read_text(encoding="utf-8"))
    assert list(cfg["source_rate_start"]) == [
        "stock_flow",
        "market",
        "weather_load",
        "weekly_count",
    ]
    assert list(cfg["source_rate_start"].values()) == [
        0.0,
        0.0,
        0.0,
        0.0,
    ]
    assert list(cfg["source_rate_final"].values()) == [
        0.0,
        0.0,
        0.0,
        0.0,
    ]
    assert sum(cfg["source_rate_start"].values()) == 0.0
    assert sum(cfg["source_rate_final"].values()) == 0.0
    gen = _module().Generator(str(ROOT), seed=3)
    kinds = gen._select_sources(gen._rng(0x0300), 0, 8192, 104_062_502)
    assert np.all(kinds == -1)


@pytest.mark.parametrize("value", [0, -1, 4097, 1.5, True])
def test_rejects_invalid_chunk_size(tmp_path, value):
    cfg = json.loads((ROOT / "config.json").read_text(encoding="utf-8"))
    cfg["chunk_size"] = value
    (tmp_path / "config.json").write_text(json.dumps(cfg), encoding="utf-8")
    with pytest.raises(ValueError, match="chunk_size"):
        _module().Generator(str(tmp_path), seed=1)


def test_static_cascade_validation():
    cfg = load_chain_config(CASCADE / "chain.toml")
    assert check_repo_layout(ROOT).ok
    assert check_config(ROOT).ok
    assert check_requirements_hash_locked(
        ROOT / "requirements.txt",
        allowed=cfg.dependencies.allowed,
        max_packages=cfg.dependencies.max_packages,
    ).ok
    assert scan_file(ROOT / "generator.py", cfg.static_guard.blocked).ok
    report = verify_repo(ROOT, cfg, skip_runtime=True)
    assert report.ok and report.runtime_skipped