theme_generation/engines/markov.py

"""Variable-order Markov engine backed by vo_regular_bp."""

from __future__ import annotations

import random
import sys
from pathlib import Path
from typing import Any

from ..common import START_SYMBOL, Symbol
from ..constraints import make_theme_acceptor


VO_REGULAR_PATH = Path("/Users/francoispachet/IdeaProjects/vo_regular_bp")
if VO_REGULAR_PATH.exists():
    sys.path.insert(0, str(VO_REGULAR_PATH))

try:
    from vo_regular_bp import (
        ConstraintSet,
        LongestFeasiblePolicy,
        OrderStackModel,
        prepare_constrained_order_stack,
    )
except ImportError as exc:  # pragma: no cover - depends on local sibling repo.
    ConstraintSet = None
    LongestFeasiblePolicy = None
    OrderStackModel = None
    prepare_constrained_order_stack = None
    VO_REGULAR_IMPORT_ERROR = exc
else:
    VO_REGULAR_IMPORT_ERROR = None


def require_vo_regular() -> None:
    if VO_REGULAR_IMPORT_ERROR is not None:
        raise RuntimeError(
            "The Markov engine requires the local vo_regular_bp package. "
            "Expected it at /Users/francoispachet/IdeaProjects/vo_regular_bp."
        ) from VO_REGULAR_IMPORT_ERROR


def generate_markov(
    *,
    sequences: list[tuple[Symbol, ...]],
    length: int,
    samples: int,
    max_order: int,
    start_weights: dict[int, float],
    end_weights: dict[int, float],
    endpoint_strength: float,
    enforce_triplet_groups: bool,
    seed: int,
) -> tuple[list[tuple[Symbol, ...]], dict[str, Any]]:
    require_vo_regular()

    model = OrderStackModel.from_sequences(
        sequences,
        max_order=max_order,
        start_symbol=START_SYMBOL,
    )
    acceptor = make_theme_acceptor(
        length=length,
        alphabet=model.alphabet,
        start_weights=start_weights,
        end_weights=end_weights,
        strength=endpoint_strength,
        enforce_triplet_groups=enforce_triplet_groups,
    )
    backend = prepare_constrained_order_stack(
        model,
        ConstraintSet(regular_acceptors=(acceptor,)),
        length=length,
        prefix=(START_SYMBOL,),
        policy=LongestFeasiblePolicy(),
    )
    rng = random.Random(seed)
    generated = [backend.sample(rng=rng) for _ in range(samples)]
    diagnostics = {
        "constrained success mass": f"{backend.diagnostics.success_mass:.6g}",
        "engine": "vo_regular variable-order Markov",
    }
    return generated, diagnostics