Benchmark_eval/TS_Caption_test/_build/ts_align/data.py

from __future__ import annotations

import random
import subprocess
from collections import OrderedDict
from collections.abc import Sequence
from pathlib import Path
from typing import Any

import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DistributedSampler, Sampler

TS_START_TOKEN = "<ts>"
TS_END_TOKEN = "</ts>"
SCALE_START_TOKEN = "<scale>"
SCALE_END_TOKEN = "</scale>"

PromptMode = str

STAGE1_PROMPT_VARIANTS_UNIVAR: tuple[str, ...] = (
    "请描述这个时间序列：{ts_start} {ts_end}",
    "请概括这个时间序列：{ts_start} {ts_end}",
    "请总结这个时间序列：{ts_start} {ts_end}",
    "请简要描述这个时间序列：{ts_start} {ts_end}",
    "请简要概括这个时间序列：{ts_start} {ts_end}",
    "请总结这个时间序列的主要特征：{ts_start} {ts_end}",
    "请描述这个时间序列的主要表现：{ts_start} {ts_end}",
    "请概括这个时间序列的整体情况：{ts_start} {ts_end}",
    "请总结这个时间序列的整体特征：{ts_start} {ts_end}",
    "请简述这个时间序列的主要模式：{ts_start} {ts_end}",
    "请描述该时间序列：{ts_start} {ts_end}",
    "请概括该时间序列：{ts_start} {ts_end}",
    "请总结该时间序列：{ts_start} {ts_end}",
    "请描述该时间序列的主要特征：{ts_start} {ts_end}",
    "请概括该时间序列的整体特征：{ts_start} {ts_end}",
    "请总结该时间序列的主要表现：{ts_start} {ts_end}",
    "请对这个时间序列做简要描述：{ts_start} {ts_end}",
    "请对这个时间序列做简要概括：{ts_start} {ts_end}",
    "请对这个时间序列做简要总结：{ts_start} {ts_end}",
    "请简要总结这个时间序列的主要情况：{ts_start} {ts_end}",
)

STAGE1_PROMPT_VARIANTS_BIVAR: tuple[str, ...] = (
    "请描述这个双变量时间序列：{ts_start} {ts_end}",
    "请概括这个双变量时间序列：{ts_start} {ts_end}",
    "请总结这个双变量时间序列：{ts_start} {ts_end}",
    "请简要描述这个双变量时间序列的整体表现：{ts_start} {ts_end}",
    "请简要概括这个双变量时间序列：{ts_start} {ts_end}",
    "请概括这个双变量时间序列的主要关系特征：{ts_start} {ts_end}",
    "请总结这个双变量时间序列的整体变化与相互关系：{ts_start} {ts_end}",
    "请总结这个双变量时间序列的主要特征：{ts_start} {ts_end}",
    "请描述这个双变量时间序列的主要表现：{ts_start} {ts_end}",
    "请概括这个双变量时间序列的整体情况：{ts_start} {ts_end}",
    "请总结这个双变量时间序列的整体特征：{ts_start} {ts_end}",
    "请简述这个双变量时间序列的主要关系模式：{ts_start} {ts_end}",
    "请描述该双变量时间序列：{ts_start} {ts_end}",
    "请概括该双变量时间序列：{ts_start} {ts_end}",
    "请描述该双变量时间序列的主要表现：{ts_start} {ts_end}",
    "请概括该双变量序列的整体特征：{ts_start} {ts_end}",
    "请总结该双变量序列的主要模式：{ts_start} {ts_end}",
    "请对这对时间序列做简要描述：{ts_start} {ts_end}",
    "请对这对时间序列做简要概括：{ts_start} {ts_end}",
    "请简要总结这对时间序列的主要情况：{ts_start} {ts_end}",
)

STAGE1_PROMPT_VARIANTS_MULTIVAR: tuple[str, ...] = (
    "请描述这个多变量时间序列：{ts_start} {ts_end}",
    "请概括这个多变量时间序列：{ts_start} {ts_end}",
    "请总结这个多变量时间序列：{ts_start} {ts_end}",
    "请简要描述这个多变量系统的整体表现：{ts_start} {ts_end}",
    "请简要概括这个多变量时间序列：{ts_start} {ts_end}",
    "请概括该多变量系统的主要结构特征：{ts_start} {ts_end}",
    "请总结该多变量序列的整体变化模式：{ts_start} {ts_end}",
    "请总结这个多变量时间序列的主要特征：{ts_start} {ts_end}",
    "请描述这个多变量系统的主要表现：{ts_start} {ts_end}",
    "请概括这个多变量系统的整体情况：{ts_start} {ts_end}",
    "请总结这个多变量系统的整体特征：{ts_start} {ts_end}",
    "请简述这个多变量系统的主要模式：{ts_start} {ts_end}",
    "请描述该多变量时间序列：{ts_start} {ts_end}",
    "请概括该多变量时间序列：{ts_start} {ts_end}",
    "请描述该多变量时间序列的主要表现：{ts_start} {ts_end}",
    "请概括该多变量时间序列的整体特征：{ts_start} {ts_end}",
    "请总结该多变量系统的主要模式：{ts_start} {ts_end}",
    "请对这个多变量时间序列做简要描述：{ts_start} {ts_end}",
    "请对这个多变量时间序列做简要概括：{ts_start} {ts_end}",
    "请简要总结这个多变量系统的主要情况：{ts_start} {ts_end}",
)

STAGE1_PROMPT_VARIANTS_BY_MODE: dict[PromptMode, tuple[str, ...]] = {
    "univar": STAGE1_PROMPT_VARIANTS_UNIVAR,
    "bivar": STAGE1_PROMPT_VARIANTS_BIVAR,
    "multivar": STAGE1_PROMPT_VARIANTS_MULTIVAR,
}

STAGE2_SYSTEM_PROMPT = "你是专业的时间序列分析助手，请仅根据给定时间序列完成分析。"

STAGE2_PROMPT_FAMILIES_UNIVAR: dict[str, tuple[str, ...]] = {
    "overall": (
        "请对这个时间序列做深入分析：{ts_start} {ts_end}",
        "请结合整体形态，对该时间序列做较完整的分析：{ts_start} {ts_end}",
        "请围绕整体表现和变化脉络，对该序列进行深入分析：{ts_start} {ts_end}",
    ),
    "pattern": (
        "请分析该时间序列的主要模式及其相互关系：{ts_start} {ts_end}",
        "请围绕趋势、周期性和局部波动，对该序列进行综合分析：{ts_start} {ts_end}",
        "请从主要模式及其关联的角度，对该时间序列进行分析：{ts_start} {ts_end}",
    ),
    "stability": (
        "请分析这个时间序列的稳定性与可预测性：{ts_start} {ts_end}",
        "请判断该序列的变化是否稳定，并说明其可预测性的来源：{ts_start} {ts_end}",
        "请从稳定性和可预测性的角度，对该时间序列进行分析：{ts_start} {ts_end}",
    ),
    "risk": (
        "请分析该时间序列是否存在结构变化风险，并说明依据：{ts_start} {ts_end}",
        "请从长期趋势、波动变化和潜在结构切换的角度分析该序列：{ts_start} {ts_end}",
        "请评估该时间序列的结构风险，并结合整体变化给出分析：{ts_start} {ts_end}",
    ),
}

STAGE2_PROMPT_FAMILIES_BIVAR: dict[str, tuple[str, ...]] = {
    "overall": (
        "请对这个双变量时间序列做深入分析，重点概括整体变化与两条序列的关系：{ts_start} {ts_end}",
        "请结合整体形态与变量间联系，对该双变量时间序列做较完整的分析：{ts_start} {ts_end}",
        "请围绕整体表现、协同变化和差异关系，对这对时间序列进行深入分析：{ts_start} {ts_end}",
    ),
    "pattern": (
        "请分析该双变量时间序列的主要关系模式及其相互作用：{ts_start} {ts_end}",
        "请围绕趋势一致性、节律同步和局部波动联动，对这对序列进行综合分析：{ts_start} {ts_end}",
        "请从关系结构与模式特征的角度，对该双变量时间序列进行分析：{ts_start} {ts_end}",
    ),
    "stability": (
        "请分析这对时间序列关系结构的稳定性与可预测性：{ts_start} {ts_end}",
        "请判断该双变量序列的协同变化是否稳定，并说明其可预测性的来源：{ts_start} {ts_end}",
        "请从关系稳定性和联合可预测性的角度，对这对时间序列进行分析：{ts_start} {ts_end}",
    ),
    "risk": (
        "请分析该双变量时间序列是否存在关系结构变化风险，并说明依据：{ts_start} {ts_end}",
        "请从长期趋势、波动联动和潜在结构切换的角度分析这对序列：{ts_start} {ts_end}",
        "请评估这对时间序列的耦合风险，并结合整体变化给出分析：{ts_start} {ts_end}",
    ),
}

STAGE2_PROMPT_FAMILIES_MULTIVAR: dict[str, tuple[str, ...]] = {
    "overall": (
        "请对这个多变量时间序列系统做深入分析，重点概括整体结构与动态模式：{ts_start} {ts_end}",
        "请结合系统整体形态与变量间协同关系，对该多变量时间序列做较完整的分析：{ts_start} {ts_end}",
        "请围绕整体表现、系统结构和变量间互动，对该多变量序列进行深入分析：{ts_start} {ts_end}",
    ),
    "pattern": (
        "请分析该多变量时间序列系统的主要模式特征及其相互关系：{ts_start} {ts_end}",
        "请围绕因子结构、同步协同、领先-滞后与局部波动，对该多变量系统进行综合分析：{ts_start} {ts_end}",
        "请从系统模式与变量间关联的角度，对该多变量时间序列进行分析：{ts_start} {ts_end}",
    ),
    "stability": (
        "请分析这个多变量时间序列系统的结构稳定性与可预测性：{ts_start} {ts_end}",
        "请判断该多变量系统的协同结构是否稳定，并说明其可预测性的来源：{ts_start} {ts_end}",
        "请从系统稳定性和联合可预测性的角度，对该多变量时间序列进行分析：{ts_start} {ts_end}",
    ),
    "risk": (
        "请分析该多变量时间序列系统是否存在结构变化风险，并说明依据：{ts_start} {ts_end}",
        "请从长期趋势、相关结构变化、波动联动和潜在状态切换的角度分析该多变量系统：{ts_start} {ts_end}",
        "请评估该多变量时间序列的结构脆弱性与异常风险，并结合整体变化给出分析：{ts_start} {ts_end}",
    ),
}

STAGE2_PROMPT_FAMILIES_BY_MODE: dict[PromptMode, dict[str, tuple[str, ...]]] = {
    "univar": STAGE2_PROMPT_FAMILIES_UNIVAR,
    "bivar": STAGE2_PROMPT_FAMILIES_BIVAR,
    "multivar": STAGE2_PROMPT_FAMILIES_MULTIVAR,
}

# Backward-compatible aliases: in the multivar package, the default exported prompt
# set should reflect the multivariate training path.
STAGE1_PROMPT_VARIANTS = STAGE1_PROMPT_VARIANTS_MULTIVAR
STAGE2_PROMPT_FAMILIES = STAGE2_PROMPT_FAMILIES_MULTIVAR

DEFAULT_STAGE2_PROMPT_FAMILY_WEIGHTS: dict[str, float] = {
    "overall": 0.4,
    "pattern": 0.25,
    "stability": 0.2,
    "risk": 0.15,
}

DEFAULT_STAGE2_LEVEL_WEIGHTS: dict[str, float] = {
    # Phase 3 (2026-06-10): enable level_1/2 captions to give ScaleEncoder
    # + LLM strong supervision for absolute (mu, sigma) decoding. Phase 2
    # used only level_3/4 (analytical captions), starving the signal for
    # simple-stats metrics (mean/std/min/max/median/...). See phase3_design
    # §4.1.1. When the runner doesn't provide --level12-jsonl, train_stage1
    # forces level_1/2 weights to 0 to preserve backward compatibility.
    "level_1": 0.5,
    "level_2": 0.3,
    "level_3": 2.0,
    "level_4": 1.0,
}


def register_ts_special_tokens(tokenizer, ts_token: str = "<ts>") -> int:
    token_ids = register_alignment_special_tokens(
        tokenizer,
        ts_start_token=ts_token,
        ts_end_token=TS_END_TOKEN,
        scale_start_token=SCALE_START_TOKEN,
        scale_end_token=SCALE_END_TOKEN,
    )
    return token_ids["ts_start"]


def register_alignment_special_tokens(
    tokenizer,
    *,
    ts_start_token: str = TS_START_TOKEN,
    ts_end_token: str = TS_END_TOKEN,
    scale_start_token: str = SCALE_START_TOKEN,
    scale_end_token: str = SCALE_END_TOKEN,
) -> dict[str, int]:
    tokens = [ts_start_token, ts_end_token, scale_start_token, scale_end_token]
    vocab = tokenizer.get_vocab()
    missing_tokens = [token for token in tokens if token not in vocab]
    if missing_tokens:
        tokenizer.add_special_tokens({"additional_special_tokens": missing_tokens})
    return {
        "ts_start": tokenizer.convert_tokens_to_ids(ts_start_token),
        "ts_end": tokenizer.convert_tokens_to_ids(ts_end_token),
        "scale_start": tokenizer.convert_tokens_to_ids(scale_start_token),
        "scale_end": tokenizer.convert_tokens_to_ids(scale_end_token),
    }


def collate_fn(
    batch: list[dict[str, Any]],
    tokenizer,
    ts_token: str = TS_START_TOKEN,
    ts_end_token: str = TS_END_TOKEN,
    ignore_index: int = -100,
    max_length: int | None = None,
) -> dict[str, Any]:
    if not batch:
        raise ValueError("collate_fn requires a non-empty batch.")

    special_token_ids = register_alignment_special_tokens(
        tokenizer,
        ts_start_token=ts_token,
        ts_end_token=ts_end_token,
    )
    batch_type = _detect_batch_type(batch)

    raw_ts_sequences: list[torch.Tensor] = []
    channel_counts: list[int] = []
    ts_lengths_list: list[int] = []
    input_ids_list: list[torch.Tensor] = []
    attention_masks: list[torch.Tensor] = []
    labels_list: list[torch.Tensor] = []
    text_valid_lengths: list[int] = []
    ts_start_positions: list[int] = []
    ts_end_positions: list[int] = []
    for sample in batch:
        raw_ts = _normalize_raw_ts(sample["raw_ts"])
        raw_ts_sequences.append(raw_ts)
        channel_counts.append(raw_ts.shape[0])
        ts_lengths_list.append(raw_ts.shape[-1])

        if batch_type == "text":
            input_ids, attention_mask, labels = _build_text_training_example(
                sample=sample,
                tokenizer=tokenizer,
                ignore_index=ignore_index,
                max_length=max_length,
            )
        else:
            input_ids, attention_mask, labels = _build_pretokenized_example(
                sample=sample,
                ignore_index=ignore_index,
                max_length=max_length,
            )

        valid_length, start_pos, end_pos = _validate_single_placeholder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            ts_start_token_id=special_token_ids["ts_start"],
            ts_end_token_id=special_token_ids["ts_end"],
        )
        text_valid_lengths.append(valid_length)
        ts_start_positions.append(start_pos)
        ts_end_positions.append(end_pos)

        input_ids_list.append(input_ids)
        attention_masks.append(attention_mask)
        labels_list.append(labels)

    batch_size = len(raw_ts_sequences)
    max_channels = max(channel_counts)
    max_ts_len = max(ts_lengths_list)

    raw_ts = torch.zeros(
        batch_size,
        max_channels,
        max_ts_len,
        dtype=torch.float32,
    )
    raw_ts_channel_mask = torch.zeros(
        batch_size,
        max_channels,
        dtype=torch.long,
    )
    raw_ts_attention_mask = torch.zeros(
        batch_size,
        max_ts_len,
        dtype=torch.long,
    )
    for batch_index, sequence in enumerate(raw_ts_sequences):
        n_channels, seq_len = sequence.shape
        raw_ts[batch_index, :n_channels, :seq_len] = sequence
        raw_ts_channel_mask[batch_index, :n_channels] = 1
        raw_ts_attention_mask[batch_index, :seq_len] = 1
    input_ids = pad_sequence(
        input_ids_list,
        batch_first=True,
        padding_value=_get_pad_token_id(tokenizer),
    )
    attention_mask = pad_sequence(
        attention_masks,
        batch_first=True,
        padding_value=0,
    )
    labels = pad_sequence(
        labels_list,
        batch_first=True,
        padding_value=ignore_index,
    )

    return {
        "raw_ts": raw_ts,
        "raw_ts_channel_mask": raw_ts_channel_mask,
        "raw_ts_attention_mask": raw_ts_attention_mask,
        "text_valid_lengths": text_valid_lengths,
        "ts_start_positions": ts_start_positions,
        "ts_end_positions": ts_end_positions,
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels,
    }


def get_stage1_prompt_variants(
    ts_start_token: str = TS_START_TOKEN,
    ts_end_token: str = TS_END_TOKEN,
    prompt_variants: Sequence[str] | None = None,
    prompt_mode: PromptMode = "multivar",
) -> list[str]:
    variants = prompt_variants or STAGE1_PROMPT_VARIANTS_BY_MODE[prompt_mode]
    formatted = [
        _format_prompt_variant(
            variant,
            ts_start_token=ts_start_token,
            ts_end_token=ts_end_token,
        )
        for variant in variants
    ]
    if not formatted:
        raise ValueError("At least one stage1 prompt variant is required.")
    return formatted


def sample_stage1_prompt(
    *,
    ts_start_token: str = TS_START_TOKEN,
    ts_end_token: str = TS_END_TOKEN,
    prompt_variants: Sequence[str] | None = None,
    prompt_mode: PromptMode = "multivar",
    rng: random.Random | None = None,
) -> str:
    variants = get_stage1_prompt_variants(
        ts_start_token=ts_start_token,
        ts_end_token=ts_end_token,
        prompt_variants=prompt_variants,
        prompt_mode=prompt_mode,
    )
    chooser = rng.choice if rng is not None else random.choice
    return chooser(variants)


def build_stage1_training_samples(
    record: dict[str, Any],
    *,
    ts_start_token: str = TS_START_TOKEN,
    ts_end_token: str = TS_END_TOKEN,
    prompt_variants: Sequence[str] | None = None,
    rng: random.Random | None = None,
) -> list[dict[str, Any]]:
    _validate_stage1_record(record)

    samples: list[dict[str, Any]] = []
    for level_key in ("level_1", "level_2"):
        samples.append(
            _build_stage1_sample(
                record=record,
                level_key=level_key,
                ts_start_token=ts_start_token,
                ts_end_token=ts_end_token,
                prompt_variants=prompt_variants,
                rng=rng,
            )
        )
    return samples


class Stage1AlignmentDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        records: Sequence[dict[str, Any]],
        *,
        ts_start_token: str = TS_START_TOKEN,
        ts_end_token: str = TS_END_TOKEN,
        prompt_variants: Sequence[str] | None = None,
        seed: int = 0,
        dynamic: bool = True,
    ) -> None:
        self.records = list(records)
        self.ts_start_token = ts_start_token
        self.ts_end_token = ts_end_token
        self.prompt_variants = prompt_variants
        self.seed = seed
        self.dynamic = dynamic
        self.epoch = 0

        for record in self.records:
            _validate_stage1_record(record)

    def __len__(self) -> int:
        return len(self.records) * 2

    def set_epoch(self, epoch: int) -> None:
        self.epoch = epoch

    def __getitem__(self, index: int) -> dict[str, Any]:
        if index < 0:
            index += len(self)
        if index < 0 or index >= len(self):
            raise IndexError("Stage1AlignmentDataset index out of range.")

        record_index, level_index = divmod(index, 2)
        level_key = ("level_1", "level_2")[level_index]
        sample_seed = self.seed + index
        if self.dynamic:
            sample_seed += self.epoch * max(len(self), 1)
        rng = random.Random(sample_seed)
        return _build_stage1_sample(
            record=self.records[record_index],
            level_key=level_key,
            ts_start_token=self.ts_start_token,
            ts_end_token=self.ts_end_token,
            prompt_variants=self.prompt_variants,
            rng=rng,
        )


class Stage2AlignmentDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        records: Sequence[dict[str, Any]],
        *,
        system_prompt: str = STAGE2_SYSTEM_PROMPT,
        ts_start_token: str = TS_START_TOKEN,
        ts_end_token: str = TS_END_TOKEN,
        prompt_families: dict[str, Sequence[str]] | None = None,
        prompt_family_weights: dict[str, float] | None = None,
        level_weights: dict[str, float] | None = None,
        seed: int = 0,
        dynamic: bool = True,
    ) -> None:
        self.records = list(records)
        self.system_prompt = system_prompt
        self.ts_start_token = ts_start_token
        self.ts_end_token = ts_end_token
        self.prompt_families = prompt_families
        self.prompt_family_weights = prompt_family_weights
        self.level_weights = level_weights
        self.seed = seed
        self.dynamic = dynamic
        self.epoch = 0

        for record in self.records:
            _validate_stage2_record(record)

    def __len__(self) -> int:
        return len(self.records)

    def set_epoch(self, epoch: int) -> None:
        self.epoch = epoch

    def __getitem__(self, index: int) -> dict[str, Any]:
        if index < 0:
            index += len(self)
        if index < 0 or index >= len(self):
            raise IndexError("Stage2AlignmentDataset index out of range.")

        sample_seed = self.seed + index
        if self.dynamic:
            sample_seed += self.epoch * max(len(self.records), 1)
        rng = random.Random(sample_seed)
        return build_stage2_training_sample(
            self.records[index],
            system_prompt=self.system_prompt,
            ts_start_token=self.ts_start_token,
            ts_end_token=self.ts_end_token,
            prompt_families=self.prompt_families,
            prompt_family_weights=self.prompt_family_weights,
            level_weights=self.level_weights,
            rng=rng,
        )


class QAWarmupDataset(torch.utils.data.Dataset):
    """SFT warm-up dataset that consumes pre-built (prompt, answer) QA pairs.

    Records are produced by build_qa_warmup_jsonl.py: each one already carries
    a fully-rendered Chinese prompt (with `<ts> </ts>` placeholders) and the
    deterministic reference answer (`metric_anchor = value;` joined by `；`).
    """

    def __init__(
        self,
        records: Sequence[dict[str, Any]],
        *,
        system_prompt: str | None = None,
        seed: int = 0,
        dynamic: bool = False,
    ) -> None:
        self.records = list(records)
        self.system_prompt = system_prompt
        self.seed = seed
        self.dynamic = dynamic
        self.epoch = 0

        for record in self.records:
            _validate_qa_warmup_record(record)

    def __len__(self) -> int:
        return len(self.records)

    def set_epoch(self, epoch: int) -> None:
        self.epoch = epoch

    def __getitem__(self, index: int) -> dict[str, Any]:
        if index < 0:
            index += len(self)
        if index < 0 or index >= len(self):
            raise IndexError("QAWarmupDataset index out of range.")
        record = self.records[index]
        sample: dict[str, Any] = {
            "raw_ts": record["raw_ts"],
            "prompt": record["prompt"],
            "target_text": record["answer"],
            "target_level": "qa_warmup",
            "source_id": record.get("id"),
        }
        # Prefer the record's own system_prompt (mixed pool: TSQA rows carry
        # TSQA_SYSTEM_PROMPT); fall back to the dataset-level default for rows
        # without one (our metric_qa) — identical to the previous behavior.
        sp = record.get("system_prompt", self.system_prompt)
        if sp is not None:
            sample["system_prompt"] = sp
        if "metric_keys" in record:
            sample["metric_keys"] = record["metric_keys"]
        return sample


def load_qa_warmup_records_from_jsonl(
    *,
    raw_values_path: str | Path,
    qa_pairs_path: str | Path,
    limit: int | None = None,
) -> list[dict[str, Any]]:
    """Join raw_values + qa_pairs by id. One record per QA pair (multiple per id)."""
    raw_records = _read_jsonl_file(raw_values_path)
    raw_by_id: dict[int, dict[str, Any]] = {}
    for record in raw_records:
        if "id" not in record:
            continue
        if not _has_valid_ts_values(record.get("values")):
            continue
        raw_by_id[int(record["id"])] = record

    qa_records = _read_jsonl_file(qa_pairs_path, limit=limit)
    records: list[dict[str, Any]] = []
    for qa in qa_records:
        record_id = qa.get("id")
        if record_id is None:
            continue
        raw = raw_by_id.get(int(record_id))
        if raw is None:
            continue
        prompt = qa.get("prompt")
        answer = qa.get("answer")
        if not isinstance(prompt, str) or not isinstance(answer, str):
            continue
        if not prompt.strip() or not answer.strip():
            continue
        rec = {
            "id": int(record_id),
            "qa_index": qa.get("qa_index"),
            "raw_ts": raw["values"],
            "prompt": prompt,
            "answer": answer,
            "metric_keys": qa.get("metric_keys", []),
        }
        # Carry per-row system_prompt / task_type through so a mixed warmup pool
        # (metric_qa + caption + TSQA) keeps each task's own system prompt. Rows
        # without these fields (our metric_qa sft_pairs) fall back to the
        # dataset-level default in QAWarmupDataset — behavior unchanged.
        if qa.get("system_prompt") is not None:
            rec["system_prompt"] = qa["system_prompt"]
        if qa.get("task_type") is not None:
            rec["task_type"] = qa["task_type"]
        records.append(rec)
    return records


def _validate_qa_warmup_record(record: dict[str, Any]) -> None:
    missing = [key for key in ("raw_ts", "prompt", "answer") if key not in record]
    if missing:
        raise ValueError(
            f"QAWarmup record is missing required fields: {', '.join(missing)}."
        )
    if not _has_valid_ts_values(record["raw_ts"]):
        raise ValueError("QAWarmup record raw_ts must be a non-empty sequence without null values.")
    prompt = record["prompt"]
    if not isinstance(prompt, str) or TS_START_TOKEN not in prompt or TS_END_TOKEN not in prompt:
        raise ValueError(
            f"QAWarmup prompt must contain both {TS_START_TOKEN} and {TS_END_TOKEN}."
        )


def build_univariate_stage1_records(
    sample_records: Sequence[dict[str, Any]],
    level_records: Sequence[dict[str, Any]],
) -> list[dict[str, Any]]:
    samples_by_id = {
        int(record["id"]): record
        for record in sample_records
    }
    levels_by_id = OrderedDict(
        (int(record["id"]), record)
        for record in level_records
    )

    records: list[dict[str, Any]] = []
    for record_id, level_record in levels_by_id.items():
        if record_id not in samples_by_id:
            raise ValueError(f"Missing sample values for record id {record_id}.")

        sample_record = samples_by_id[record_id]
        if not _has_valid_ts_values(sample_record.get("values")):
            continue
        records.append(
            {
                "id": record_id,
                "dataset": sample_record.get("dataset"),
                "channel": sample_record.get("channel"),
                "raw_ts": sample_record["values"],
                "level_1": (
                    level_record.get("level_1_revised")
                    or level_record.get("original_level_1")
                    or level_record.get("level_1")
                ),
                "level_2": (
                    level_record.get("level_2_revised")
                    or level_record.get("original_level_2")
                    or level_record.get("level_2")
                ),
            }
        )
    return records


def build_univariate_stage2_records(
    sample_records: Sequence[dict[str, Any]],
    level3_records: Sequence[dict[str, Any]],
    level4_records: Sequence[dict[str, Any]],
    level12_records: Sequence[dict[str, Any]] | None = None,
) -> list[dict[str, Any]]:
    samples_by_id = {
        int(record["id"]): record
        for record in sample_records
    }
    level3_by_id = OrderedDict(
        (int(record["id"]), record)
        for record in level3_records
    )
    level4_by_id = {
        int(record["id"]): record
        for record in level4_records
    }
    # Phase 3: optional level_1/level_2 supply. Records only get level_1/level_2
    # populated when level12_records is provided; otherwise the sampler in
    # train_stage1 forces their weights to 0 so the missing field is never read.
    level12_by_id: dict[int, dict[str, Any]] = {}
    if level12_records is not None:
        level12_by_id = {
            int(record["id"]): record
            for record in level12_records
        }

    records: list[dict[str, Any]] = []
    for record_id, level3_record in level3_by_id.items():
        if record_id not in samples_by_id:
            raise ValueError(f"Missing sample values for record id {record_id}.")
        if record_id not in level4_by_id:
            raise ValueError(f"Missing level_4 text for record id {record_id}.")

        sample_record = samples_by_id[record_id]
        if not _has_valid_ts_values(sample_record.get("values")):
            continue

        level4_record = level4_by_id[record_id]
        record: dict[str, Any] = {
            "id": record_id,
            "dataset": sample_record.get("dataset"),
            "channel": sample_record.get("channel"),
            "raw_ts": sample_record["values"],
            "level_3": level3_record["level_3"],
            "level_4": level4_record["level_4"],
            "level_3_prompt": level3_record.get("prompt"),
            "level_4_prompt": level4_record.get("prompt"),
        }
        level12_record = level12_by_id.get(record_id)
        if level12_record is not None:
            lvl1 = (
                level12_record.get("level_1_revised")
                or level12_record.get("original_level_1")
                or level12_record.get("level_1")
            )
            lvl2 = (
                level12_record.get("level_2_revised")
                or level12_record.get("original_level_2")
                or level12_record.get("level_2")
            )
            if lvl1 is not None:
                record["level_1"] = lvl1
                record["level_1_prompt"] = level12_record.get("level_1_prompt") or level12_record.get("prompt")
            if lvl2 is not None:
                record["level_2"] = lvl2
                record["level_2_prompt"] = level12_record.get("level_2_prompt") or level12_record.get("prompt")
        records.append(record)
    return records


def load_stage1_records_from_jsonl(
    *,
    samples_path: str | Path,
    level12_path: str | Path,
    limit: int | None = None,
) -> list[dict[str, Any]]:
    sample_records = _read_jsonl_file(samples_path, limit=limit)
    level_records = _read_jsonl_file(level12_path, limit=limit)
    return build_univariate_stage1_records(sample_records, level_records)


def load_stage1_records_from_univar_tar(
    archive_path: str | Path,
    *,
    limit: int | None = None,
) -> list[dict[str, Any]]:
    sample_records = _read_jsonl_from_tar_zst(
        archive_path,
        member_path="univar/samples.jsonl",
        limit=limit,
    )
    level_records = _read_jsonl_from_tar_zst(
        archive_path,
        member_path="univar/level12.jsonl",
        limit=limit,
    )
    return build_univariate_stage1_records(sample_records, level_records)


def load_stage2_records_from_jsonl(
    *,
    samples_path: str | Path,
    level3_path: str | Path,
    level4_path: str | Path,
    level12_path: str | Path | None = None,
    limit: int | None = None,
) -> list[dict[str, Any]]:
    sample_records = _read_jsonl_file(samples_path, limit=limit)
    level3_records = _read_jsonl_file(level3_path, limit=limit)
    level4_records = _read_jsonl_file(level4_path, limit=limit)
    level12_records = (
        _read_jsonl_file(level12_path, limit=limit) if level12_path else None
    )
    return build_univariate_stage2_records(
        sample_records,
        level3_records,
        level4_records,
        level12_records=level12_records,
    )


def load_stage2_records_from_univar_tar(
    samples_archive_path: str | Path,
    level34_archive_path: str | Path,
    *,
    limit: int | None = None,
) -> list[dict[str, Any]]:
    sample_records = _read_jsonl_from_tar_zst(
        samples_archive_path,
        member_path="univar/samples.jsonl",
        limit=limit,
    )
    level3_records = _read_jsonl_from_tar_zst(
        level34_archive_path,
        member_path="univar_level_34/level3.jsonl",
        limit=limit,
    )
    level4_records = _read_jsonl_from_tar_zst(
        level34_archive_path,
        member_path="univar_level_34/level4.jsonl",
        limit=limit,
    )
    return build_univariate_stage2_records(sample_records, level3_records, level4_records)


def _detect_batch_type(batch: list[dict[str, Any]]) -> str:
    has_text = [("prompt" in sample and "target_text" in sample) for sample in batch]
    has_tokens = [("input_ids" in sample and "labels" in sample) for sample in batch]

    if all(has_text) and not any(has_tokens):
        return "text"
    if all(has_tokens) and not any(has_text):
        return "tokenized"
    raise ValueError(
        "Batch must contain either only prompt/target_text samples or only pretokenized samples."
    )


def _build_stage1_sample(
    *,
    record: dict[str, Any],
    level_key: str,
    ts_start_token: str,
    ts_end_token: str,
    prompt_variants: Sequence[str] | None,
    rng: random.Random | None,
) -> dict[str, Any]:
    prompt_mode = _infer_prompt_mode(record["raw_ts"])
    return {
        "raw_ts": record["raw_ts"],
        "prompt": sample_stage1_prompt(
            ts_start_token=ts_start_token,
            ts_end_token=ts_end_token,
            prompt_variants=prompt_variants,
            prompt_mode=prompt_mode,
            rng=rng,
        ),
        "target_text": record[level_key],
        "target_level": level_key,
        "source_id": record.get("id"),
    }


def get_stage2_prompt_families(
    ts_start_token: str = TS_START_TOKEN,
    ts_end_token: str = TS_END_TOKEN,
    prompt_families: dict[str, Sequence[str]] | None = None,
    prompt_mode: PromptMode = "multivar",
) -> dict[str, list[str]]:
    families = prompt_families or STAGE2_PROMPT_FAMILIES_BY_MODE[prompt_mode]
    formatted = {
        family: [
            _format_prompt_variant(
                variant,
                ts_start_token=ts_start_token,
                ts_end_token=ts_end_token,
            )
            for variant in variants
        ]
        for family, variants in families.items()
    }
    if not formatted:
        raise ValueError("At least one stage2 prompt family is required.")
    for family, variants in formatted.items():
        if not variants:
            raise ValueError(f"Stage2 prompt family '{family}' must contain at least one prompt.")
    return formatted


def sample_stage2_prompt(
    *,
    ts_start_token: str = TS_START_TOKEN,
    ts_end_token: str = TS_END_TOKEN,
    prompt_families: dict[str, Sequence[str]] | None = None,
    prompt_family_weights: dict[str, float] | None = None,
    prompt_mode: PromptMode = "multivar",
    rng: random.Random | None = None,
) -> tuple[str, str]:
    formatted_families = get_stage2_prompt_families(
        ts_start_token=ts_start_token,
        ts_end_token=ts_end_token,
        prompt_families=prompt_families,
        prompt_mode=prompt_mode,
    )
    family_weights = _resolve_weight_mapping(
        prompt_family_weights,
        defaults=DEFAULT_STAGE2_PROMPT_FAMILY_WEIGHTS,
        allowed_keys=formatted_families.keys(),
        mapping_name="stage2 prompt family weights",
    )
    chooser = rng or random
    family = _weighted_choice(family_weights, chooser)
    return family, chooser.choice(formatted_families[family])


def sample_stage2_target_level(
    *,
    level_weights: dict[str, float] | None = None,
    rng: random.Random | None = None,
) -> str:
    chooser = rng or random
    weights = _resolve_weight_mapping(
        level_weights,
        defaults=DEFAULT_STAGE2_LEVEL_WEIGHTS,
        allowed_keys=DEFAULT_STAGE2_LEVEL_WEIGHTS.keys(),
        mapping_name="stage2 target weights",
    )
    return _weighted_choice(weights, chooser)


def build_stage2_training_sample(
    record: dict[str, Any],
    *,
    system_prompt: str = STAGE2_SYSTEM_PROMPT,
    ts_start_token: str = TS_START_TOKEN,
    ts_end_token: str = TS_END_TOKEN,
    prompt_families: dict[str, Sequence[str]] | None = None,
    prompt_family_weights: dict[str, float] | None = None,
    level_weights: dict[str, float] | None = None,
    rng: random.Random | None = None,
) -> dict[str, Any]:
    _validate_stage2_record(record)
    prompt_mode = _infer_prompt_mode(record["raw_ts"])
    target_level = sample_stage2_target_level(level_weights=level_weights, rng=rng)
    if target_level in ("level_1", "level_2"):
        # Phase 3: level_1/2 are short basic-stat captions. 8cd82fa enabled the
        # data fields + weights but never wired a prompt path for them, so the
        # aligned-family machinery (DIRECT_STAGE2_SOURCE_PROMPT_RULES /
        # ALIGNED_STAGE2_PROMPT_VARIANTS) only covers level_3/4 → KeyError when
        # level_1/2 is sampled. Reuse the Stage 1 generic describe prompts
        # (STAGE1_PROMPT_VARIANTS_BY_MODE), which match the short-caption task.
        source_prompt = None
        family = "describe"
        prompt = sample_stage1_prompt(
            ts_start_token=ts_start_token,
            ts_end_token=ts_end_token,
            prompt_mode=prompt_mode,
            rng=rng,
        )
    elif (source_prompt := _get_stage2_source_prompt(record, target_level)):
        family, prompt = build_aligned_stage2_prompt(
            source_prompt=source_prompt,
            target_level=target_level,
            ts_start_token=ts_start_token,
            ts_end_token=ts_end_token,
            prompt_families=prompt_families,
            prompt_mode=prompt_mode,
            rng=rng,
        )
    else:
        family, prompt = sample_stage2_prompt(
            ts_start_token=ts_start_token,
            ts_end_token=ts_end_token,
            prompt_families=prompt_families,
            prompt_family_weights=prompt_family_weights,
            prompt_mode=prompt_mode,
            rng=rng,
        )
    # Phase 3 (2026-06-10): when env TS_ALIGN_PROMPT_TASK_TYPE=1, prepend a
    # "<task_type>caption</task_type>" prefix. Pairs with the QA-side prefix
    # injected by build_single_metric_prompt so the model can disambiguate
    # which reward branch it's being trained on. Env-gated for backward
    # compat; Phase 2 SFT builds remain unchanged.
    from .qa_templates import task_type_token_prefix
    task_type_prefix = task_type_token_prefix("caption")
    if task_type_prefix:
        prompt = task_type_prefix + prompt
    return {
        "raw_ts": record["raw_ts"],
        "system_prompt": system_prompt,
        "prompt": prompt,
        "target_text": record[target_level],
        "target_level": target_level,
        "prompt_family": family,
        "source_prompt": source_prompt,
        "source_id": record.get("id"),
    }


def _validate_stage1_record(record: dict[str, Any]) -> None:
    missing = [key for key in ("raw_ts", "level_1", "level_2") if key not in record]
    if missing:
        raise ValueError(
            f"Stage1 record is missing required fields: {', '.join(missing)}."
        )
    if not _has_valid_ts_values(record["raw_ts"]):
        raise ValueError("Stage1 record raw_ts must be a non-empty sequence without null values.")


def _validate_stage2_record(record: dict[str, Any]) -> None:
    missing = [key for key in ("raw_ts", "level_3", "level_4") if key not in record]
    if missing:
        raise ValueError(
            f"Stage2 record is missing required fields: {', '.join(missing)}."
        )
    if not _has_valid_ts_values(record["raw_ts"]):
        raise ValueError("Stage2 record raw_ts must be a non-empty sequence without null values.")


def _has_valid_ts_values(values: Any) -> bool:
    if values is None:
        return False
    try:
        tensor = _normalize_raw_ts(values)
    except (TypeError, ValueError):
        return False
    return tensor.numel() > 0 and torch.isfinite(tensor).all().item()


def _infer_prompt_mode(raw_ts: Any) -> PromptMode:
    tensor = _normalize_raw_ts(raw_ts)
    n_channels = int(tensor.shape[0])
    if n_channels <= 1:
        return "univar"
    if n_channels == 2:
        return "bivar"
    return "multivar"


def _get_stage2_source_prompt(record: dict[str, Any], target_level: str) -> str | None:
    prompt_key = f"{target_level}_prompt"
    source_prompt = record.get(prompt_key)
    if isinstance(source_prompt, str) and source_prompt.strip():
        return source_prompt.strip()
    return None


def build_aligned_stage2_prompt(
    *,
    source_prompt: str,
    target_level: str,
    ts_start_token: str = TS_START_TOKEN,
    ts_end_token: str = TS_END_TOKEN,
    prompt_families: dict[str, Sequence[str]] | None = None,
    prompt_mode: PromptMode = "multivar",
    rng: random.Random | None = None,
) -> tuple[str, str]:
    source_topic = infer_stage2_source_topic(
        source_prompt=source_prompt,
        target_level=target_level,
        prompt_mode=prompt_mode,
    )
    family = SOURCE_TOPIC_TO_PROMPT_FAMILY[source_topic]
    variants = get_aligned_stage2_prompt_variants(
        prompt_mode=prompt_mode,
        source_topic=source_topic,
        ts_start_token=ts_start_token,
        ts_end_token=ts_end_token,
    )
    chooser = rng.choice if rng is not None else random.choice
    return family, chooser(variants)


def infer_stage2_prompt_family_from_source_prompt(
    *,
    source_prompt: str,
    target_level: str,
    prompt_mode: PromptMode = "multivar",
) -> str:
    source_topic = infer_stage2_source_topic(
        source_prompt=source_prompt,
        target_level=target_level,
        prompt_mode=prompt_mode,
    )
    return SOURCE_TOPIC_TO_PROMPT_FAMILY[source_topic]


def infer_stage2_source_topic(
    *,
    source_prompt: str,
    target_level: str,
    prompt_mode: PromptMode = "multivar",
) -> str:
    headline = _extract_source_prompt_headline(source_prompt)
    rules = DIRECT_STAGE2_SOURCE_PROMPT_RULES_BY_MODE[prompt_mode][target_level]
    for phrase, source_topic in rules:
        if phrase in headline:
            return source_topic
    raise ValueError(
        f"Unable to map source prompt headline to aligned topic for mode={prompt_mode}, "
        f"target_level={target_level}, headline={headline!r}."
    )


def get_aligned_stage2_prompt_variants(
    *,
    prompt_mode: PromptMode,
    source_topic: str,
    ts_start_token: str = TS_START_TOKEN,
    ts_end_token: str = TS_END_TOKEN,
) -> list[str]:
    prompt_variants = ALIGNED_STAGE2_PROMPT_VARIANTS_BY_MODE[prompt_mode][source_topic]
    return [
        _format_prompt_variant(
            variant,
            ts_start_token=ts_start_token,
            ts_end_token=ts_end_token,
        )
        for variant in prompt_variants
    ]


def _contains_any(text: str, keywords: Sequence[str]) -> bool:
    return any(keyword in text for keyword in keywords)


def _extract_source_prompt_headline(source_prompt: str) -> str:
    for line in source_prompt.splitlines():
        headline = line.strip()
        if headline:
            return headline
    return source_prompt.strip()


SOURCE_TOPIC_TO_PROMPT_FAMILY: dict[str, str] = {
    "pattern": "pattern",
    "association": "pattern",
    "summary": "overall",
    "coupling_stability": "stability",
    "frequency": "stability",
    "risk": "risk",
    "deep_summary": "overall",
}


DIRECT_STAGE2_SOURCE_PROMPT_RULES_UNIVAR: dict[str, tuple[tuple[str, str], ...]] = {
    "level_3": (
        ("综合分析其主要模式", "pattern"),
        ("主要模式特征", "pattern"),
        ("模式总结", "summary"),
        ("关联关系", "association"),
        ("核心模式", "summary"),
    ),
    "level_4": (
        ("稳定性、复杂度和结构风险", "coupling_stability"),
        ("稳定性和可预测性", "coupling_stability"),
        ("频率结构与复杂度", "frequency"),
        ("异常风险与结构脆弱性", "risk"),
        ("最具技术价值的深层特征", "deep_summary"),
    ),
}


DIRECT_STAGE2_SOURCE_PROMPT_RULES_BIVAR: dict[str, tuple[tuple[str, str], ...]] = {
    "level_3": (
        ("综合分析两序列的关系模式", "pattern"),
        ("主要关系模式", "pattern"),
        ("关系总结", "summary"),
        ("内在关联", "association"),
        ("最突出的关系特征", "summary"),
    ),
    "level_4": (
        ("综合分析两序列的深层耦合特性", "deep_summary"),
        ("耦合结构与时变稳定性", "coupling_stability"),
        ("因果结构与频域特征", "frequency"),
        ("结构脆弱性与极端联动风险", "risk"),
        ("最具技术价值的深层特征", "deep_summary"),
    ),
}


DIRECT_STAGE2_SOURCE_PROMPT_RULES_MULTIVAR: dict[str, tuple[tuple[str, str], ...]] = {
    "level_3": (
        ("综合分析其结构和动态模式", "pattern"),
        ("主要模式特征", "pattern"),
        ("总结概括", "summary"),
        ("关联关系", "association"),
        ("核心特征", "summary"),
    ),
    "level_4": (
        ("综合分析系统的动态耦合特性和结构稳定性", "deep_summary"),
        ("动态耦合结构与时变稳定性", "coupling_stability"),
        ("频率特征与季节性结构", "frequency"),
        ("结构脆弱性与异常特征", "risk"),
        ("最具技术价值的深层特征", "deep_summary"),
    ),
}


DIRECT_STAGE2_SOURCE_PROMPT_RULES_BY_MODE: dict[
    PromptMode, dict[str, tuple[tuple[str, str], ...]]
] = {
    "univar": DIRECT_STAGE2_SOURCE_PROMPT_RULES_UNIVAR,
    "bivar": DIRECT_STAGE2_SOURCE_PROMPT_RULES_BIVAR,
    "multivar": DIRECT_STAGE2_SOURCE_PROMPT_RULES_MULTIVAR,
}


ALIGNED_STAGE2_PROMPT_VARIANTS_UNIVAR: dict[str, tuple[str, ...]] = {
    "pattern": (
        "请综合分析这个时间序列的主要模式：{ts_start} {ts_end}",
        "请识别这个时间序列的主要模式特征：{ts_start} {ts_end}",
        "请分析这个时间序列的主要模式：{ts_start} {ts_end}",
    ),
    "association": (
        "请分析这个时间序列各特征之间的关联关系：{ts_start} {ts_end}",
        "请分析这个时间序列的关联关系：{ts_start} {ts_end}",
    ),
    "summary": (
        "请对这个时间序列做模式总结：{ts_start} {ts_end}",
        "请概括这个时间序列的核心模式：{ts_start} {ts_end}",
        "请总结这个时间序列的模式：{ts_start} {ts_end}",
    ),
    "coupling_stability": (
        "请分析这个时间序列的稳定性、复杂度和结构风险：{ts_start} {ts_end}",
        "请重点分析这个时间序列的稳定性和可预测性：{ts_start} {ts_end}",
        "请分析这个时间序列的稳定性和可预测性：{ts_start} {ts_end}",
    ),
    "frequency": (
        "请深入分析这个时间序列的频率结构与复杂度：{ts_start} {ts_end}",
        "请分析这个时间序列的频率结构与复杂度：{ts_start} {ts_end}",
    ),
    "risk": (
        "请评估这个时间序列的异常风险与结构脆弱性：{ts_start} {ts_end}",
        "请评估这个时间序列的结构脆弱性与异常风险：{ts_start} {ts_end}",
    ),
    "deep_summary": (
        "请概括这个时间序列最具技术价值的深层特征：{ts_start} {ts_end}",
        "请概括这个时间序列的深层特征：{ts_start} {ts_end}",
    ),
}


ALIGNED_STAGE2_PROMPT_VARIANTS_BIVAR: dict[str, tuple[str, ...]] = {
    "pattern": (
        "请综合分析这对时间序列的关系模式：{ts_start} {ts_end}",
        "请识别这对时间序列的主要关系模式：{ts_start} {ts_end}",
        "请分析这对时间序列的关系模式：{ts_start} {ts_end}",
    ),
    "association": (
        "请分析这对时间序列各维度特征之间的内在关联：{ts_start} {ts_end}",
        "请分析这对时间序列的内在关联：{ts_start} {ts_end}",
    ),
    "summary": (
        "请对这对时间序列做关系总结：{ts_start} {ts_end}",
        "请概括这对时间序列最突出的关系特征：{ts_start} {ts_end}",
        "请总结这对时间序列的关系特征：{ts_start} {ts_end}",
    ),
    "coupling_stability": (
        "请分析该双变量时间序列的耦合结构与时变稳定性：{ts_start} {ts_end}",
        "请重点分析这对时间序列的耦合结构与时变稳定性：{ts_start} {ts_end}",
        "请分析这对时间序列的耦合结构与时变稳定性：{ts_start} {ts_end}",
    ),
    "frequency": (
        "请分析该双变量时间序列的因果结构与频域特征：{ts_start} {ts_end}",
        "请深入分析该双变量时间序列的因果结构与频域特征：{ts_start} {ts_end}",
        "请分析这对时间序列的因果结构与频域特征：{ts_start} {ts_end}",
    ),
    "risk": (
        "请评估该双变量时间序列的结构脆弱性与极端联动风险：{ts_start} {ts_end}",
        "请评估这对时间序列的结构脆弱性与极端联动风险：{ts_start} {ts_end}",
    ),
    "deep_summary": (
        "请概括该双变量时间序列最具技术价值的深层特征：{ts_start} {ts_end}",
        "请概括这对时间序列最具技术价值的深层特征：{ts_start} {ts_end}",
    ),
}


ALIGNED_STAGE2_PROMPT_VARIANTS_MULTIVAR: dict[str, tuple[str, ...]] = {
    "pattern": (
        "请综合分析该多变量时间序列的结构和动态模式：{ts_start} {ts_end}",
        "请识别该多变量系统的主要模式特征：{ts_start} {ts_end}",
        "请分析该多变量时间序列的结构和动态模式：{ts_start} {ts_end}",
    ),
    "association": (
        "请分析该多变量系统各特征之间的关联关系：{ts_start} {ts_end}",
        "请分析该多变量时间序列的关联关系：{ts_start} {ts_end}",
    ),
    "summary": (
        "请对该多变量系统做总结概括：{ts_start} {ts_end}",
        "请概括该多变量系统的核心特征：{ts_start} {ts_end}",
        "请总结该多变量时间序列的特征：{ts_start} {ts_end}",
    ),
    "coupling_stability": (
        "请分析该多变量时间序列的动态耦合特性和结构稳定性：{ts_start} {ts_end}",
        "请分析该多变量系统的动态耦合结构与时变稳定性：{ts_start} {ts_end}",
        "请重点分析该多变量系统的动态耦合结构与时变稳定性：{ts_start} {ts_end}",
        "请分析该多变量系统的动态耦合与结构稳定性：{ts_start} {ts_end}",
    ),
    "frequency": (
        "请分析该多变量时间序列的频率特征与季节性结构：{ts_start} {ts_end}",
        "请深入分析该多变量时间序列的频率特征与季节性结构：{ts_start} {ts_end}",
        "请分析该多变量系统的频率特征与季节性结构：{ts_start} {ts_end}",
    ),
    "risk": (
        "请评估该多变量时间序列的结构脆弱性与异常特征：{ts_start} {ts_end}",
        "请评估该多变量系统的结构脆弱性与异常特征：{ts_start} {ts_end}",
    ),
    "deep_summary": (
        "请概括该多变量时间序列最具技术价值的深层特征：{ts_start} {ts_end}",
        "请概括该多变量系统最具技术价值的深层特征：{ts_start} {ts_end}",
    ),
}


ALIGNED_STAGE2_PROMPT_VARIANTS_BY_MODE: dict[PromptMode, dict[str, tuple[str, ...]]] = {
    "univar": ALIGNED_STAGE2_PROMPT_VARIANTS_UNIVAR,
    "bivar": ALIGNED_STAGE2_PROMPT_VARIANTS_BIVAR,
    "multivar": ALIGNED_STAGE2_PROMPT_VARIANTS_MULTIVAR,
}


def _normalize_raw_ts(raw_ts: Any) -> torch.Tensor:
    tensor = torch.as_tensor(raw_ts, dtype=torch.float32)
    if tensor.ndim == 1:
        return tensor.unsqueeze(0)
    if tensor.ndim == 2:
        if tensor.shape[0] == 1:
            return tensor
        if tensor.shape[1] == 1:
            return tensor.transpose(0, 1)
        # For multivariate inputs, prefer [C, L]. If the first dimension is much
        # larger, interpret the tensor as [L, C] and transpose into channel-first layout.
        if tensor.shape[0] > tensor.shape[1]:
            return tensor.transpose(0, 1)
        return tensor
    raise ValueError("raw_ts must have shape [L], [C, L], or [L, C].")


def _build_text_training_example(
    *,
    sample: dict[str, Any],
    tokenizer,
    ignore_index: int,
    max_length: int | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    prompt_messages: list[dict[str, str]] = []
    system_prompt = sample.get("system_prompt")
    if system_prompt is not None:
        prompt_messages.append({"role": "system", "content": system_prompt})
    prompt_messages.append({"role": "user", "content": sample["prompt"]})
    full_messages = [
        *prompt_messages,
        {"role": "assistant", "content": sample["target_text"]},
    ]

    prompt_ids = _apply_chat_template(
        tokenizer,
        prompt_messages,
        add_generation_prompt=True,
    )
    input_ids = _apply_chat_template(
        tokenizer,
        full_messages,
        add_generation_prompt=False,
    )
    if len(prompt_ids) >= len(input_ids):
        raise ValueError("Chat template must leave assistant tokens after the user prompt prefix.")

    attention_mask = [1] * len(input_ids)
    labels = [ignore_index] * len(prompt_ids) + input_ids[len(prompt_ids) :]

    input_ids_tensor = torch.tensor(input_ids, dtype=torch.long)
    attention_mask_tensor = torch.tensor(attention_mask, dtype=torch.long)
    labels_tensor = torch.tensor(labels, dtype=torch.long)

    return _truncate_text_fields(
        input_ids=input_ids_tensor,
        attention_mask=attention_mask_tensor,
        labels=labels_tensor,
        max_length=max_length,
    )


def _build_pretokenized_example(
    *,
    sample: dict[str, Any],
    ignore_index: int,
    max_length: int | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    input_ids = torch.as_tensor(sample["input_ids"], dtype=torch.long)
    attention_mask = torch.as_tensor(
        sample.get("attention_mask", torch.ones_like(input_ids)),
        dtype=torch.long,
    )
    labels = torch.as_tensor(sample["labels"], dtype=torch.long)

    if input_ids.ndim != 1 or attention_mask.ndim != 1 or labels.ndim != 1:
        raise ValueError("Pretokenized input_ids, attention_mask, and labels must be 1D.")

    return _truncate_text_fields(
        input_ids=input_ids,
        attention_mask=attention_mask,
        labels=labels,
        max_length=max_length,
    )


def _truncate_text_fields(
    *,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    labels: torch.Tensor,
    max_length: int | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    if max_length is None or input_ids.numel() <= max_length:
        return input_ids, attention_mask, labels

    input_ids = input_ids[:max_length]
    attention_mask = attention_mask[:max_length]
    labels = labels[:max_length]
    return input_ids, attention_mask, labels


def _apply_chat_template(
    tokenizer,
    messages: list[dict[str, str]],
    *,
    add_generation_prompt: bool,
) -> list[int]:
    if not hasattr(tokenizer, "apply_chat_template"):
        raise ValueError("Tokenizer must support apply_chat_template for stage1 text examples.")

    input_ids = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=add_generation_prompt,
    )
    if isinstance(input_ids, torch.Tensor):
        return input_ids.tolist()
    return list(input_ids)


def _validate_single_placeholder(
    *,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    ts_start_token_id: int,
    ts_end_token_id: int,
) -> tuple[int, int, int]:
    valid_length = int(attention_mask.sum().item())
    valid_input_ids = input_ids[:valid_length]
    start_positions = (valid_input_ids == ts_start_token_id).nonzero(as_tuple=False).flatten()
    end_positions = (valid_input_ids == ts_end_token_id).nonzero(as_tuple=False).flatten()
    if start_positions.numel() != 1 or end_positions.numel() != 1:
        raise ValueError("Each sample must contain exactly one <ts> and one </ts> token.")
    start_pos = int(start_positions.item())
    end_pos = int(end_positions.item())
    if start_pos >= end_pos:
        raise ValueError("<ts> must appear before </ts> in each sample.")
    return valid_length, start_pos, end_pos


def _format_prompt_variant(
    variant: str,
    *,
    ts_start_token: str,
    ts_end_token: str,
) -> str:
    formatted = variant.format(
        ts_start=ts_start_token,
        ts_end=ts_end_token,
    )
    if formatted.count(ts_start_token) != 1 or formatted.count(ts_end_token) != 1:
        raise ValueError(
            "Each stage1 prompt variant must contain exactly one <ts> and one </ts> placeholder."
        )
    return formatted


def _resolve_weight_mapping(
    weights: dict[str, float] | None,
    *,
    defaults: dict[str, float],
    allowed_keys,
    mapping_name: str,
) -> dict[str, float]:
    resolved = dict(defaults)
    if weights is not None:
        resolved.update(weights)
    resolved = {key: float(value) for key, value in resolved.items() if key in allowed_keys}
    if not resolved:
        raise ValueError(f"{mapping_name} must contain at least one entry.")
    if any(value < 0 for value in resolved.values()):
        raise ValueError(f"{mapping_name} cannot contain negative weights.")
    total = sum(resolved.values())
    if total <= 0:
        raise ValueError(f"{mapping_name} must sum to a positive value.")
    return resolved


def _weighted_choice(weights: dict[str, float], rng: random.Random) -> str:
    total = sum(weights.values())
    threshold = rng.random() * total
    cumulative = 0.0
    last_key = next(iter(weights))
    for key, value in weights.items():
        cumulative += value
        last_key = key
        if threshold <= cumulative:
            return key
    return last_key


def _read_jsonl_file(path: str | Path, *, limit: int | None = None) -> list[dict[str, Any]]:
    import json

    records: list[dict[str, Any]] = []
    with open(path, "r", encoding="utf-8") as handle:
        for index, line in enumerate(handle):
            if limit is not None and index >= limit:
                break
            records.append(json.loads(line))
    return records


def _read_jsonl_from_tar_zst(
    archive_path: str | Path,
    *,
    member_path: str,
    limit: int | None = None,
) -> list[dict[str, Any]]:
    import json

    command = (
        f"zstd -dc {Path(archive_path)} | tar -xOf - {member_path}"
    )
    process = subprocess.Popen(
        ["bash", "-lc", command],
        text=True,
        encoding="utf-8",
        errors="replace",
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        bufsize=1,
    )

    records: list[dict[str, Any]] = []
    assert process.stdout is not None
    assert process.stderr is not None
    try:
        for index, line in enumerate(process.stdout):
            if limit is not None and index >= limit:
                process.terminate()
                break
            if not line.strip():
                continue
            records.append(json.loads(line))
    finally:
        process.stdout.close()
        stderr = process.stderr.read()
        process.wait()

    if process.returncode not in (0, -15):
        raise RuntimeError(f"Failed to extract {member_path} from {archive_path}: {stderr}")
    return records


def _get_pad_token_id(tokenizer) -> int:
    pad_token_id = getattr(tokenizer, "pad_token_id", None)
    if pad_token_id is not None:
        return int(pad_token_id)

    eos_token_id = getattr(tokenizer, "eos_token_id", None)
    if eos_token_id is not None:
        return int(eos_token_id)

    raise ValueError("Tokenizer must define pad_token_id or eos_token_id for collation.")


# ---------------------------------------------------------------------------
# Phase 3 Tier-3 §T0 (2026-06-10): resumable samplers with persistent
# (epoch, position) state, so --resume-from sees the same prompt sequence
# the pre-kill trajectory would have. Closes the only remaining post-resume
# divergence source after Tier-2 §5/§6 (trainer_state + RNG persistence) —
# without T0, even with model weights + optimizer + RNG perfectly restored
# the DataLoader would yield the FIRST batch of a fresh epoch instead of
# resuming mid-epoch, and the next ~50 micro-batches would consume entirely
# different prompts than the original run.
#
# Both classes save {epoch: int, position_in_epoch: int}; load_state_dict
# resumes from the saved position; set_epoch (called explicitly at epoch
# boundary by the train loop) resets the position. Single-rank and
# distributed variants share the same state-dict shape so train_stage3
# can save without caring which is in use.
# ---------------------------------------------------------------------------
class _ResumableSamplerMixin:
    """Shared (epoch, position) persistence for the two resumable samplers."""

    epoch: int
    _position_in_epoch: int

    def state_dict(self) -> dict[str, int]:
        return {
            "epoch": int(self.epoch),
            "position_in_epoch": int(self._position_in_epoch),
        }

    def load_state_dict(self, state: dict[str, int]) -> None:
        # Use set_epoch so any subclass-specific bookkeeping fires (e.g.
        # DistributedSampler reseeds its shuffle generator off epoch).
        self.set_epoch(int(state.get("epoch", 0)))
        self._position_in_epoch = int(state.get("position_in_epoch", 0))


class ResumableRandomSampler(_ResumableSamplerMixin, Sampler[int]):
    """Single-rank shuffle sampler with deterministic per-epoch ordering.

    Replaces the implicit RandomSampler inside DataLoader(shuffle=True) for
    Phase 3 GRPO single-GPU training. Seeds the shuffle off (seed, epoch)
    so the index sequence is reproducible across runs, and remembers how
    many indices were yielded so resume picks up mid-epoch instead of
    restarting from index 0. Drops the implicit non-determinism that
    DataLoader(shuffle=True) ships with by default.
    """

    def __init__(self, data_source, *, seed: int = 42, epoch: int = 0):
        self.data_source = data_source
        self.seed = int(seed)
        self.epoch = int(epoch)
        self._position_in_epoch = 0

    def set_epoch(self, epoch: int) -> None:
        self.epoch = int(epoch)
        self._position_in_epoch = 0

    def __iter__(self):
        generator = torch.Generator()
        generator.manual_seed(self.seed + self.epoch)
        indices = torch.randperm(len(self.data_source), generator=generator).tolist()
        skip = self._position_in_epoch
        for offset, idx in enumerate(indices[skip:], start=skip):
            # Update BEFORE yield: a Python generator suspended at `yield`
            # never resumes if the consumer breaks. If we updated after the
            # yield, an early break would leave the saved position pointing
            # at the just-yielded index (so resume would re-yield it). The
            # consumer is committed to consuming the value the moment
            # __next__ returns, so charging position += 1 first matches
            # "elements yielded" exactly.
            self._position_in_epoch = offset + 1
            yield idx
        # Iter exhausted — caller should advance epoch + call set_epoch.

    def __len__(self):
        return len(self.data_source)


class ResumableDistributedSampler(_ResumableSamplerMixin, DistributedSampler):
    """DistributedSampler with mid-epoch resume support.

    Index ordering matches the parent DistributedSampler exactly — the
    Phase 3 contribution is only that we remember how many of those indices
    have already been yielded and skip ahead on the next __iter__. Used by
    Phase 3 GRPO under FSDP / DDP where the existing code already wires a
    DistributedSampler via train_stage3_grpo.py.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._position_in_epoch = 0

    def set_epoch(self, epoch: int) -> None:
        super().set_epoch(epoch)
        self._position_in_epoch = 0

    def __iter__(self):
        # super().__iter__ yields a generator; materialise once so the
        # skip-and-resume logic can index into the deterministic order.
        base_indices = list(super().__iter__())
        skip = self._position_in_epoch
        for offset, idx in enumerate(base_indices[skip:], start=skip):
            # See ResumableRandomSampler.__iter__: update position BEFORE
            # yield to keep state consistent under early consumer-break.
            self._position_in_epoch = offset + 1
            yield idx