tests/test_generator_contract.py

from __future__ import annotations

import importlib.util
from pathlib import Path

import numpy as np
import pytest


REPO = Path(__file__).resolve().parents[1]


def _generator(seed: int = 123):
    spec = importlib.util.spec_from_file_location(
        "profilefix_generator", REPO / "generator.py"
    )
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    module._CHUNK = 128
    return module.Generator(str(REPO), seed=seed)


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


def test_exact_count_allowed_lengths_and_finite():
    rows = list(_generator().generate(131))
    assert len(rows) == 131
    assert {row.shape for row in rows} == {(4096,)}
    assert all(row.dtype == np.float64 for row in rows)
    assert all(np.isfinite(row).all() and row.std() > 1e-9 for row in rows)


def test_deterministic_across_instances():
    left = list(_generator().generate(130))
    right = list(_generator().generate(130))
    assert all(np.array_equal(a, b) for a, b in zip(left, right, strict=True))


def test_requested_size_does_not_change_prefix():
    short = list(_generator().generate(17))
    long = list(_generator().generate(130))
    assert all(np.array_equal(a, b) for a, b in zip(short, long[:17], strict=True))


def test_nonpositive_request_is_empty():
    assert list(_generator().generate(0)) == []


def test_genuine_short_outputs_are_disabled():
    rows = list(_generator(seed=211).generate(512))
    lengths = np.asarray([row.size for row in rows])
    assert set(lengths) == {4096}


def test_family_weights_are_exact_conservative_uid131_scaling():
    module = _module()
    uid131 = {
        "trend_seasonal_ar": 0.12,
        "regime_shift": 0.12,
        "multiplicative": 0.08,
        "ar2": 0.15,
        "integrated": 0.12,
        "threshold_ar": 0.08,
        "chaotic": 0.04,
        "spectral_gp": 0.07,
        "long_memory": 0.06,
        "ou_stochastic_vol": 0.10,
        "physical_sensors": 0.02,
        "seasonal_counts": 0.02,
        "intermittent": 0.01,
        "pulse_outlier": 0.01,
    }
    assert sum(module._DEFAULT_WEIGHTS.values()) == pytest.approx(1.0)
    assert module._DEFAULT_WEIGHTS["prefix_cumulative_counter"] == 0.025
    for family, weight in uid131.items():
        assert module._DEFAULT_WEIGHTS[family] == pytest.approx(0.975 * weight)


def test_period5_added_without_changing_period96_probability():
    module = _module()
    assert module._SEASONAL_PROBS.sum() == pytest.approx(1.0)
    period_to_probability = dict(
        zip(
            module._SEASONAL_PERIODS.astype(int),
            module._SEASONAL_PROBS,
            strict=True,
        )
    )
    assert period_to_probability[5] > 0.0
    # Raw mass was reallocated only from periods 4 and 7. UID131's raw total
    # was 1.07, so period 96 retains exactly its prior conditional probability.
    assert period_to_probability[96] == pytest.approx(0.07 / 1.07)


def test_first_value_padding_keeps_suffix_and_reads_no_later_value():
    module = _module()
    raw = np.arange(3 * 32, dtype=np.float64).reshape(3, 32)
    lengths = np.asarray([8, 12, 20])
    padded = module._left_pad_first_value(raw, lengths)
    for row, active in enumerate(lengths):
        start = raw.shape[1] - int(active)
        assert np.array_equal(padded[row, start:], raw[row, start:])
        assert np.all(padded[row, :start] == raw[row, start])
    assert np.array_equal(raw, np.arange(3 * 32).reshape(3, 32))


def test_cumulative_primitive_is_prefix_causal():
    module = _module()
    steps = np.random.default_rng(223).normal(size=(3, 512))
    base = np.asarray([1.0e4, 2.0e8, 3.0e12])
    start = np.asarray([7, 31, 63])
    full = module._causal_cumulative_from_steps(steps, base, start)
    prefix = module._causal_cumulative_from_steps(
        steps[:, :257], base, start
    )
    assert np.array_equal(full[:, :257], prefix)


def test_prefix_cumulative_family_is_deterministic_large_and_padded():
    module = _module()
    left = module._prefix_cumulative_counter(
        np.random.default_rng(227), 64, 4096
    )
    right = module._prefix_cumulative_counter(
        np.random.default_rng(227), 64, 4096
    )
    assert np.array_equal(left, right)
    assert np.isfinite(left).all()
    assert (left >= 0.0).all()
    assert float(left.max()) > 1.0e8
    changed = np.diff(left, axis=1) != 0.0
    first_change = changed.argmax(axis=1) + 1
    assert (first_change >= 4096 - 1024).all()
    assert (changed.sum(axis=1) > 8).all()
    assert "prefix_cumulative_counter" in module._FAMILIES
    assert module._DEFAULT_WEIGHTS["prefix_cumulative_counter"] == 0.025


def test_ou_stochastic_vol_is_deterministic_finite_and_active():
    module = _module()
    left = module._ou_stochastic_vol(np.random.default_rng(77), 16, 512)
    right = module._ou_stochastic_vol(np.random.default_rng(77), 16, 512)
    assert np.array_equal(left, right)
    assert np.isfinite(left).all()
    assert (left.std(axis=1) > 1e-9).all()
    assert "ou_stochastic_vol" in module._FAMILIES
    assert module._DEFAULT_WEIGHTS["ou_stochastic_vol"] > 0


def test_physical_sensors_are_deterministic_finite_and_active():
    module = _module()
    left = module._physical_sensors(np.random.default_rng(91), 32, 512)
    right = module._physical_sensors(np.random.default_rng(91), 32, 512)
    assert np.array_equal(left, right)
    assert np.isfinite(left).all()
    assert (left.std(axis=1) > 1e-9).all()
    assert "physical_sensors" in module._FAMILIES
    assert module._DEFAULT_WEIGHTS["physical_sensors"] > 0


def test_seasonal_counts_are_deterministic_integer_and_active():
    module = _module()
    left = module._seasonal_counts(np.random.default_rng(109), 32, 512)
    right = module._seasonal_counts(np.random.default_rng(109), 32, 512)
    assert np.array_equal(left, right)
    assert np.isfinite(left).all()
    assert (left >= 0).all()
    assert np.array_equal(left, np.round(left))
    assert (left.std(axis=1) > 1e-9).all()
    assert "seasonal_counts" in module._FAMILIES
    assert module._DEFAULT_WEIGHTS["seasonal_counts"] > 0


def test_event_family_has_real_flat_runs_and_recovery():
    module = _module()
    left = module._pulse_outlier(np.random.default_rng(131), 32, 512)
    right = module._pulse_outlier(np.random.default_rng(131), 32, 512)
    assert np.array_equal(left, right)
    assert np.isfinite(left).all()
    assert (left.std(axis=1) > 1e-9).all()
    assert np.any(np.diff(left, axis=1) == 0.0)


def test_measurement_artifacts_are_deterministic_finite_and_shape_safe():
    module = _module()
    raw = module._trend_seasonal_ar(np.random.default_rng(149), 64, 512)
    left = module._measurement_artifacts(
        np.random.default_rng(151), raw, preserve_nonnegative=False
    )
    right = module._measurement_artifacts(
        np.random.default_rng(151), raw, preserve_nonnegative=False
    )
    assert np.array_equal(left, right)
    assert left.shape == raw.shape
    assert np.isfinite(left).all()
    assert (left.std(axis=1) > 1e-9).all()