tick-stock-panel/backend/app/backtest/regime_alignment.py

from __future__ import annotations

from collections.abc import Mapping, Sequence
from datetime import date
from typing import Any

import numpy as np

RegimePoint = tuple[str, float]

REGIME_THREE_LEVEL_MAP = {
    "strong": "strong",
    "lean_strong": "strong",
    "range": "range",
    "lean_weak": "weak",
    "weak": "weak",
}


def three_level_regime(state: str) -> str:
    return REGIME_THREE_LEVEL_MAP.get(state, state)


def _date_text(value: object) -> str:
    return str(value)[:10]


def _normalize_regime_point(value: Any) -> RegimePoint:
    if isinstance(value, Mapping):
        state = str(value.get("state", ""))
        score = float(value.get("score", 0) or 0)
        return state, score
    if isinstance(value, (tuple, list)) and len(value) >= 2:
        return str(value[0]), float(value[1] or 0)
    raise ValueError("市场环境数据格式无效")


def align_regime_t_minus_one(
    labels: Sequence[str],
    regime_by_date: Mapping[object, Any],
    required_start: date | None,
    required_end: date | None,
) -> list[RegimePoint | None]:
    """Align each label with the preceding label's regime without any I/O."""
    regime_map = {
        _date_text(key): _normalize_regime_point(value)
        for key, value in regime_by_date.items()
    }
    if not regime_map:
        raise ValueError("市场环境数据为空, 请先在数据页完成市场环境计算后再回测")

    aligned: list[RegimePoint | None] = [None] * len(labels)
    required_start_text = str(required_start) if required_start is not None else None
    required_end_text = str(required_end) if required_end is not None else None
    missing_dates: list[str] = []
    if labels and required_start_text is not None:
        first_label = _date_text(labels[0])
        if first_label >= required_start_text and (
            required_end_text is None or first_label <= required_end_text
        ):
            raise ValueError(
                f"市场环境数据覆盖不完整: 正式首日 {first_label} 缺少前一交易日环境, "
                "请把前一交易日行情包含在预热区间"
            )
    for index in range(1, len(labels)):
        current_label = _date_text(labels[index])
        previous_label = _date_text(labels[index - 1])
        point = regime_map.get(previous_label)
        if point is not None:
            aligned[index] = point
            continue
        required = (
            (required_start_text is None or current_label >= required_start_text)
            and (required_end_text is None or current_label <= required_end_text)
        )
        if required:
            missing_dates.append(previous_label)

    if missing_dates:
        first_missing = missing_dates[0]
        suffix = f" 等 {len(missing_dates)} 天" if len(missing_dates) > 1 else ""
        raise ValueError(
            f"市场环境数据覆盖不完整: 缺少前一交易日环境 {first_missing}{suffix}, "
            "请先补算对应区间"
        )
    return aligned


def build_regime_filter_mask(
    labels: Sequence[str],
    regime_filter: Mapping[str, Any] | None,
    regime_by_date: Mapping[object, Any],
    *,
    required_start: date | None = None,
    required_end: date | None = None,
) -> np.ndarray | None:
    """Build a T-1 regime filter mask from caller-supplied regime data.

    States are matched against the raw five-level labels, so each regime
    level can be filtered on its own. Callers that want the aggregated
    three-level view must list the raw states explicitly, e.g.
    ``["strong", "lean_strong"]`` for the strong bucket.
    """
    if not regime_filter:
        return None
    allowed_states = {
        str(state)
        for state in (regime_filter.get("states") or [])
    }
    min_score = regime_filter.get("min_score")
    if not allowed_states and min_score is None:
        return None

    aligned = align_regime_t_minus_one(
        labels,
        regime_by_date,
        required_start,
        required_end,
    )
    mask = np.ones(len(labels), dtype=bool)
    for index, point in enumerate(aligned):
        if point is None:
            continue
        state, score = point
        mask[index] = (
            (not allowed_states or state in allowed_states)
            and (min_score is None or score >= float(min_score))
        )
    return mask