source/project/experiments/forecasting_v7/runner.py

"""Train one forecasting-v7 seed and export its immutable development sheet."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import platform
import random
import sys
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable

import numpy as np
import torch
from torch import nn

from project.data.close_distribution_v3 import (
    ADAPTER_VERSION,
    CANONICAL_HORIZONS,
    CLASS_COUNT,
    DEVELOPMENT_START,
    DEVELOPMENT_TEST,
    EVALUATION_HORIZONS,
    MINUTE_NS,
    RESEARCH_END,
    SELECTION_ASSETS,
    TOKEN_FEATURE_NAMES,
    TRAIN_END,
    VALIDATION_END,
)
from project.data.reference_level_features import (
    materialize_reference_features,
)
from project.evaluators.close_distribution_v2 import (
    ClassMap,
    class_map_to_arrays,
    encode_returns,
    fit_class_map,
)
from project.experiments.forecasting_v4.runner import (
    FEATURE_SETS,
    _context,
    _load_npz,
    _masked_nll,
    _normalization,
    _optimizer,
    _ordinal,
)
from project.experiments.forecasting_v5.runner import (
    _asset_balanced_loss_part,
    _balanced_epoch_chunks,
)
from project.models.close_distribution_v4 import (
    TypedAuxiliaryTransformer,
    build_typed_auxiliary_model,
)


RUNNER_VERSION = "forecasting-v7.0"
FEATURE_SET = "session_time"
FEATURE_NAMES = FEATURE_SETS[FEATURE_SET]
AUXILIARY_FEATURE_NAMES = (
    "prior_5_session_net_return_bps",
    "prior_5_session_rms_observed_return_bps",
    "prior_5_session_range_width_bps",
    "prior_5_session_terminal_drawdown_bps",
    "prior_5_session_available",
    "prior_20_session_net_return_bps",
    "prior_20_session_rms_observed_return_bps",
    "prior_20_session_range_width_bps",
    "prior_20_session_terminal_drawdown_bps",
    "prior_20_session_available",
    "distance_to_session_open_bps",
    "distance_to_session_max_close_so_far_bps",
    "distance_to_session_min_close_so_far_bps",
    "session_range_position_so_far",
    "session_history_available",
)
AUXILIARY_MASK_INDICES = (4, 9, 14)
DEFAULT_EPOCHS = 16
DEFAULT_CONTEXT_LENGTH = 128
DEFAULT_TARGET_PARAMETERS = 25_000
DEFAULT_ROTARY_BASE = 16.0
DEFAULT_BATCH_SIZE = 512
DEFAULT_GRADIENT_ACCUMULATION = 8
DEFAULT_INFERENCE_BATCH_SIZE = 1_024
PROVENANCE_PATHS = (
    "project/data/close_distribution_v3.py",
    "project/data/reference_level_features.py",
    "project/evaluators/close_distribution_v2.py",
    "project/experiments/forecasting_v4/runner.py",
    "project/experiments/forecasting_v5/runner.py",
    "project/experiments/forecasting_v7/runner.py",
    "project/models/close_distribution_v3.py",
    "project/models/close_distribution_v4.py",
    "project/studies/forecasting_v7.md",
    "requirements.txt",
)


@dataclass
class AssetPhaseData:
    asset: str
    asset_id: int
    features: dict[str, np.ndarray]
    bundle: dict[str, np.ndarray]
    center: np.ndarray
    scale: np.ndarray
    auxiliary_center: np.ndarray
    auxiliary_scale: np.ndarray
    auxiliary: np.ndarray
    class_map: ClassMap
    train_rows: np.ndarray
    train_labels: np.ndarray
    train_target_mask: np.ndarray
    validation_rows: np.ndarray
    validation_labels: np.ndarray
    validation_target_mask: np.ndarray


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _configure(seed: int, threads: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.set_num_threads(threads)
    torch.use_deterministic_algorithms(True)


def _selected_feature_indices() -> tuple[int, ...]:
    if tuple(FEATURE_NAMES) != (
        "close_log_return_since_last_observed_close",
        "close_return_observed",
        "log1p_elapsed_wall_clock_minutes_since_last_observed_close",
        "regular_session_progress",
        "regular_session_progress_sin",
        "regular_session_progress_cos",
    ):
        raise ValueError("the frozen v7 sequence feature set changed")
    return tuple(TOKEN_FEATURE_NAMES.index(name) for name in FEATURE_NAMES)


def _canonical_indices() -> tuple[int, ...]:
    lookup = {
        int(value): index
        for index, value in enumerate(EVALUATION_HORIZONS)
    }
    return tuple(lookup[int(value)] for value in CANONICAL_HORIZONS)


def _load_selection_data(
    study_dir: Path,
) -> tuple[
    dict[str, Any],
    dict[str, tuple[dict[str, np.ndarray], dict[str, np.ndarray]]],
]:
    manifest_path = study_dir / "data" / "prepared-manifest.json"
    manifest = json.loads(manifest_path.read_text())
    if manifest.get("adapter_version") != ADAPTER_VERSION:
        raise ValueError("v7 prepared data has the wrong adapter version")
    if manifest.get("scope") != "selection":
        raise ValueError("v7 accepts only selection-scope prepared data")
    if tuple(manifest.get("selection_assets", ())) != SELECTION_ASSETS:
        raise ValueError("v7 selection assets changed")
    source = manifest.get("source_snapshot", {})
    if source.get("research_end_exclusive") != RESEARCH_END.isoformat():
        raise ValueError("v7 prepared data has the wrong research boundary")
    if source.get("sealed_start") != RESEARCH_END.isoformat():
        raise ValueError("v7 sealed boundary is inconsistent")

    loaded = {}
    for asset in SELECTION_ASSETS:
        slug = asset.lower()
        feature_path = study_dir / "data" / "runner" / f"{slug}_features.npz"
        label_path = study_dir / "data" / "runner" / f"{slug}_labels.npz"
        for path in (feature_path, label_path):
            relative = str(path.relative_to(study_dir))
            expected = manifest["derived_sha256"].get(relative)
            if expected is None or _sha256(path) != expected:
                raise ValueError(f"v7 prepared hash mismatch: {relative}")
        loaded[asset] = (_load_npz(feature_path), _load_npz(label_path))
    return manifest, loaded


def _raw_auxiliary(features: dict[str, np.ndarray]) -> np.ndarray:
    bundle = materialize_reference_features(features)
    result = np.zeros(
        (len(bundle.generic_values), len(AUXILIARY_FEATURE_NAMES)),
        dtype=np.float32,
    )
    result[:, :10] = bundle.generic_values[:, 5:15]
    result[:, 10:15] = bundle.session_values
    if not np.isfinite(result).all():
        raise ValueError("v7 auxiliary features contain nonfinite values")
    return result


def _fit_auxiliary_scaler(
    values: np.ndarray,
    train_rows: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    center = np.zeros(values.shape[1], dtype=np.float64)
    scale = np.ones(values.shape[1], dtype=np.float64)
    selected = values[train_rows].astype(np.float64, copy=False)
    for window_index, mask_index in enumerate(AUXILIARY_MASK_INDICES):
        available = selected[:, mask_index] > 0.5
        if not np.any(available):
            continue
        start = window_index * 5
        for feature_index in range(start, start + 4):
            observed = selected[available, feature_index]
            median = float(np.median(observed))
            q25, q75 = np.quantile(observed, (0.25, 0.75))
            robust_scale = float((q75 - q25) / 1.349)
            if not np.isfinite(robust_scale) or robust_scale < 1e-8:
                robust_scale = float(np.std(observed))
            if not np.isfinite(robust_scale) or robust_scale < 1e-8:
                robust_scale = 1.0
            center[feature_index] = median
            scale[feature_index] = robust_scale
    return center.astype(np.float32), scale.astype(np.float32)


def _scale_auxiliary(
    values: np.ndarray,
    center: np.ndarray,
    scale: np.ndarray,
) -> np.ndarray:
    result = np.zeros_like(values, dtype=np.float32)
    for window_index, mask_index in enumerate(AUXILIARY_MASK_INDICES):
        available = values[:, mask_index] > 0.5
        start = window_index * 5
        result[available, start : start + 4] = (
            values[available, start : start + 4]
            - center[None, start : start + 4]
        ) / scale[None, start : start + 4]
        result[:, mask_index] = values[:, mask_index]
    if not np.isfinite(result).all():
        raise ValueError("scaled v7 auxiliary features contain nonfinite values")
    return result


def _date_mask(
    dates: np.ndarray,
    interval: tuple[str, str],
) -> np.ndarray:
    return (
        (dates >= _ordinal(interval[0]))
        & (dates < _ordinal(interval[1]))
    )


def _prepare_phase_data(
    *,
    asset: str,
    asset_id: int,
    features: dict[str, np.ndarray],
    bundle: dict[str, np.ndarray],
    train_interval: tuple[str, str],
    validation_interval: tuple[str, str] | None,
    selected_indices: tuple[int, ...],
    canonical_indices: tuple[int, ...],
    max_train_rows_per_asset: int,
) -> AssetPhaseData:
    dates = bundle["session_date"]
    train = _date_mask(dates, train_interval)
    if not np.any(train):
        raise ValueError(f"{asset} has no v7 training rows")
    train_rows = bundle["row_index"][train]
    train_y = bundle["y"][train]
    train_target_mask = bundle["target_mask"][train]
    if (
        max_train_rows_per_asset > 0
        and len(train_rows) > max_train_rows_per_asset
    ):
        selected = np.linspace(
            0,
            len(train_rows) - 1,
            max_train_rows_per_asset,
            dtype=np.int64,
        )
        train_rows = train_rows[selected]
        train_y = train_y[selected]
        train_target_mask = train_target_mask[selected]
    class_map = fit_class_map(
        train_y,
        train_target_mask,
        class_count=CLASS_COUNT,
    )
    all_train_labels = encode_returns(
        train_y,
        train_target_mask,
        class_map.edges,
    )
    center, scale = _normalization(
        features,
        train_end=train_interval[1],
        selected_indices=selected_indices,
    )
    raw_auxiliary = _raw_auxiliary(features)
    auxiliary_center, auxiliary_scale = _fit_auxiliary_scaler(
        raw_auxiliary,
        train_rows,
    )
    auxiliary = _scale_auxiliary(
        raw_auxiliary,
        auxiliary_center,
        auxiliary_scale,
    )

    if validation_interval is None:
        validation_rows = np.empty(0, dtype=np.int64)
        validation_labels = np.empty(
            (0, len(canonical_indices)),
            dtype=np.int16,
        )
        validation_target_mask = np.empty(
            (0, len(canonical_indices)),
            dtype=np.bool_,
        )
    else:
        validation = _date_mask(dates, validation_interval)
        validation_rows = bundle["row_index"][validation]
        validation_target_mask = bundle["target_mask"][validation][
            :, canonical_indices
        ]
        validation_labels = encode_returns(
            bundle["y"][validation],
            bundle["target_mask"][validation],
            class_map.edges,
        )[:, canonical_indices]
        if not len(validation_rows):
            raise ValueError(f"{asset} has no v7 validation rows")

    return AssetPhaseData(
        asset=asset,
        asset_id=asset_id,
        features=features,
        bundle=bundle,
        center=center,
        scale=scale,
        auxiliary_center=auxiliary_center,
        auxiliary_scale=auxiliary_scale,
        auxiliary=auxiliary,
        class_map=class_map,
        train_rows=train_rows,
        train_labels=all_train_labels[:, canonical_indices],
        train_target_mask=train_target_mask[:, canonical_indices],
        validation_rows=validation_rows,
        validation_labels=validation_labels,
        validation_target_mask=validation_target_mask,
    )


def _build_model(
    checkpoint: dict[str, Any],
    *,
    device: torch.device,
) -> TypedAuxiliaryTransformer:
    if tuple(checkpoint.get("assets", ())) != SELECTION_ASSETS:
        raise ValueError("locked checkpoint has the wrong assets")
    if tuple(checkpoint.get("feature_names", ())) != tuple(FEATURE_NAMES):
        raise ValueError("locked checkpoint has the wrong sequence features")
    if tuple(checkpoint.get("supervised_horizons", ())) != tuple(
        int(value) for value in CANONICAL_HORIZONS
    ):
        raise ValueError("locked checkpoint has the wrong horizons")
    if checkpoint.get("asset_conditioning") != "early_add":
        raise ValueError("locked checkpoint is not early-add")
    if int(checkpoint.get("context_length", -1)) != DEFAULT_CONTEXT_LENGTH:
        raise ValueError("locked checkpoint has the wrong context length")
    if float(checkpoint.get("rotary_base", -1.0)) != DEFAULT_ROTARY_BASE:
        raise ValueError("locked checkpoint has the wrong RoPE base")

    model = build_typed_auxiliary_model(
        context_length=DEFAULT_CONTEXT_LENGTH,
        channels=len(FEATURE_NAMES) + 1,
        output_horizons=tuple(
            float(value) for value in CANONICAL_HORIZONS
        ),
        classes=CLASS_COUNT,
        target_parameters=DEFAULT_TARGET_PARAMETERS,
        rotary_base=DEFAULT_ROTARY_BASE,
        asset_count=len(SELECTION_ASSETS),
        asset_conditioning="early_add",
        auxiliary_dim=len(AUXILIARY_FEATURE_NAMES),
    )
    incompatible = model.load_state_dict(
        checkpoint["state_dict"],
        strict=False,
    )
    if set(incompatible.missing_keys) != {"auxiliary_affine.weight"}:
        raise ValueError(
            f"unexpected missing checkpoint keys: {incompatible.missing_keys}"
        )
    if incompatible.unexpected_keys:
        raise ValueError(
            f"unexpected checkpoint keys: {incompatible.unexpected_keys}"
        )
    return model.to(device)


def _training_microbatch(
    data: tuple[AssetPhaseData, ...],
    asset_ids: np.ndarray,
    positions: np.ndarray,
    *,
    selected_indices: tuple[int, ...],
) -> tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
]:
    tokens = []
    padding = []
    auxiliary = []
    labels = []
    masks = []
    ids = []
    for item in data:
        selected = asset_ids == item.asset_id
        if not np.any(selected):
            continue
        local = positions[selected]
        rows = item.train_rows[local]
        context, context_padding = _context(
            item.features["X"],
            rows,
            context_length=DEFAULT_CONTEXT_LENGTH,
            selected_indices=selected_indices,
            center=item.center,
            scale=item.scale,
        )
        tokens.append(context)
        padding.append(context_padding)
        auxiliary.append(torch.from_numpy(item.auxiliary[rows]))
        labels.append(
            torch.from_numpy(
                item.train_labels[local].astype(np.int64, copy=False)
            )
        )
        masks.append(
            torch.from_numpy(
                item.train_target_mask[local].astype(
                    np.float32,
                    copy=False,
                )
            )
        )
        ids.append(
            torch.full(
                (len(local),),
                item.asset_id,
                dtype=torch.long,
            )
        )
    return (
        torch.cat(tokens),
        torch.cat(padding),
        torch.cat(auxiliary),
        torch.cat(labels),
        torch.cat(masks),
        torch.cat(ids),
    )


def _predict(
    model: TypedAuxiliaryTransformer,
    item: AssetPhaseData,
    *,
    rows: np.ndarray,
    selected_indices: tuple[int, ...],
    batch_size: int,
    device: torch.device,
    query_horizons: Iterable[int] | None,
) -> np.ndarray:
    horizons = (
        tuple(int(value) for value in model.output_horizons.cpu().numpy())
        if query_horizons is None
        else tuple(int(value) for value in query_horizons)
    )
    result = np.empty(
        (len(rows), len(horizons), CLASS_COUNT),
        dtype=np.float32,
    )
    query = (
        None
        if query_horizons is None
        else torch.tensor(horizons, dtype=torch.float32, device=device)
    )
    model.eval()
    with torch.inference_mode():
        for start in range(0, len(rows), batch_size):
            stop = min(start + batch_size, len(rows))
            local_rows = rows[start:stop]
            tokens, padding = _context(
                item.features["X"],
                local_rows,
                context_length=DEFAULT_CONTEXT_LENGTH,
                selected_indices=selected_indices,
                center=item.center,
                scale=item.scale,
            )
            asset_id = torch.full(
                (stop - start,),
                item.asset_id,
                dtype=torch.long,
                device=device,
            )
            logits = model(
                tokens.to(device),
                padding.to(device),
                asset_id,
                query,
                torch.from_numpy(item.auxiliary[local_rows]).to(device),
            )
            result[start:stop] = torch.softmax(logits, dim=2).cpu().numpy()
    if not np.isfinite(result).all():
        raise ValueError("v7 prediction contains nonfinite values")
    if not np.allclose(result.sum(axis=2), 1.0, atol=2e-6, rtol=0.0):
        raise ValueError("v7 probabilities do not sum to one")
    return result


def _validation_nll(
    model: TypedAuxiliaryTransformer,
    data: tuple[AssetPhaseData, ...],
    *,
    selected_indices: tuple[int, ...],
    inference_batch_size: int,
    device: torch.device,
) -> tuple[float, dict[str, float]]:
    by_asset = {}
    for item in data:
        probabilities = _predict(
            model,
            item,
            rows=item.validation_rows,
            selected_indices=selected_indices,
            batch_size=inference_batch_size,
            device=device,
            query_horizons=None,
        )
        by_asset[item.asset] = _masked_nll(
            probabilities,
            item.validation_labels,
            item.validation_target_mask,
        )
    return float(np.mean(list(by_asset.values()))), by_asset


def _train(
    model: TypedAuxiliaryTransformer,
    data: tuple[AssetPhaseData, ...],
    *,
    seed: int,
    epochs: int,
    batch_size: int,
    gradient_accumulation_steps: int,
    inference_batch_size: int,
    selected_indices: tuple[int, ...],
    device: torch.device,
    select_on_validation: bool,
) -> dict[str, Any]:
    optimizer = _optimizer(model, "adamw_constant")
    effective_batch_size = batch_size * gradient_accumulation_steps
    generator = np.random.default_rng(seed + 10_007)
    train_history = []
    validation_history = []
    best_epoch = 0
    best_nll = float("inf")
    best_state: dict[str, torch.Tensor] | None = None
    updates_per_epoch = None
    total_updates = 0

    for epoch in range(epochs):
        epoch_started = time.monotonic()
        model.train()
        updates = _balanced_epoch_chunks(
            tuple(len(item.train_rows) for item in data),
            effective_batch_size=effective_batch_size,
            generator=generator,
        )
        if updates_per_epoch is None:
            updates_per_epoch = len(updates)
        elif updates_per_epoch != len(updates):
            raise AssertionError("v7 updates per epoch changed")
        epoch_loss_sum = np.zeros(
            (len(data), len(CANONICAL_HORIZONS)),
            dtype=np.float64,
        )
        epoch_target_count = np.zeros_like(
            epoch_loss_sum,
            dtype=np.int64,
        )
        for chunks in updates:
            update_asset_ids = np.concatenate(
                [
                    np.full(len(chunk), index, dtype=np.int64)
                    for index, chunk in enumerate(chunks)
                ]
            )
            update_positions = np.concatenate(chunks)
            denominators = np.stack(
                [
                    item.train_target_mask[chunks[index]].sum(axis=0)
                    for index, item in enumerate(data)
                ]
            ).astype(np.float32)
            if np.any(denominators == 0):
                raise ValueError(
                    "a v7 asset-balanced update has no target for a horizon"
                )
            denominator_tensor = torch.from_numpy(denominators).to(device)
            optimizer.zero_grad(set_to_none=True)
            for start in range(0, len(update_positions), batch_size):
                stop = min(start + batch_size, len(update_positions))
                (
                    tokens,
                    padding,
                    auxiliary,
                    labels,
                    target_mask,
                    asset_id,
                ) = _training_microbatch(
                    data,
                    update_asset_ids[start:stop],
                    update_positions[start:stop],
                    selected_indices=selected_indices,
                )
                tokens = tokens.to(device)
                padding = padding.to(device)
                auxiliary = auxiliary.to(device)
                labels = labels.to(device)
                target_mask = target_mask.to(device)
                asset_id = asset_id.to(device)
                logits = model(
                    tokens,
                    padding,
                    asset_id,
                    auxiliary=auxiliary,
                )
                loss = _asset_balanced_loss_part(
                    logits,
                    labels,
                    target_mask,
                    asset_id,
                    denominator_tensor,
                )
                loss.backward()
                with torch.no_grad():
                    losses = nn.functional.cross_entropy(
                        logits.reshape(-1, CLASS_COUNT),
                        labels.clamp_min(0).reshape(-1),
                        reduction="none",
                    ).reshape(labels.shape)
                    for asset_index in range(len(data)):
                        selected = asset_id == asset_index
                        epoch_loss_sum[asset_index] += (
                            losses[selected] * target_mask[selected]
                        ).sum(dim=0).cpu().numpy()
                        epoch_target_count[asset_index] += (
                            target_mask[selected]
                            .sum(dim=0)
                            .cpu()
                            .numpy()
                            .astype(np.int64)
                        )
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            total_updates += 1

        train_nll = float(
            np.mean(epoch_loss_sum / epoch_target_count)
        )
        train_history.append(train_nll)
        record: dict[str, Any] = {
            "duration_seconds": time.monotonic() - epoch_started,
            "epoch": epoch + 1,
            "train_nll": train_nll,
        }
        if select_on_validation:
            equal_asset_nll, by_asset = _validation_nll(
                model,
                data,
                selected_indices=selected_indices,
                inference_batch_size=inference_batch_size,
                device=device,
            )
            record["equal_asset_validation_nll"] = equal_asset_nll
            record["validation_nll_by_asset"] = by_asset
            if equal_asset_nll < best_nll:
                best_nll = equal_asset_nll
                best_epoch = epoch + 1
                best_state = {
                    name: value.detach().cpu().clone()
                    for name, value in model.state_dict().items()
                }
        validation_history.append(record)

    if select_on_validation:
        if best_state is None:
            raise AssertionError("v7 validation selected no checkpoint")
        model.load_state_dict(best_state)
    else:
        best_epoch = epochs
    if updates_per_epoch is None:
        raise AssertionError("v7 training produced no updates")
    return {
        "best_epoch": best_epoch,
        "best_equal_asset_validation_nll": (
            best_nll if select_on_validation else None
        ),
        "completed_fixed_epoch_budget": len(train_history) == epochs,
        "epoch_history": validation_history,
        "train_nll": train_history,
        "total_updates": total_updates,
        "updates_per_epoch": updates_per_epoch,
    }


def _raw_close_vector(
    study_dir: Path,
    manifest: dict[str, Any],
    item: AssetPhaseData,
) -> np.ndarray:
    import pyarrow.parquet as pq

    record = manifest["source_snapshot"]["boundary"]["assets"][
        item.asset
    ]["outputs"]["bars"]
    path = study_dir / record["path"]
    if _sha256(path) != record["sha256"]:
        raise ValueError(f"{item.asset} raw source hash changed")
    table = pq.read_table(
        path,
        columns=["window_start_ns", "close"],
    )
    timestamp = table["window_start_ns"].to_numpy(zero_copy_only=False)
    close = table["close"].to_numpy(zero_copy_only=False).astype(np.float64)
    order = np.argsort(timestamp, kind="stable")
    timestamp = timestamp[order]
    close = close[order]
    grid_timestamp = item.features["timestamp_ns"]
    positions = np.searchsorted(timestamp, grid_timestamp)
    matched = positions < len(timestamp)
    matched[matched] &= timestamp[positions[matched]] == grid_timestamp[matched]
    result = np.full(len(grid_timestamp), np.nan, dtype=np.float64)
    result[matched] = close[positions[matched]]
    observed = item.features["X"][:, 1] > 0.5
    if np.any(observed & ~np.isfinite(result)):
        raise ValueError(f"{item.asset} observed row lacks a raw close")
    return result


def _write_prediction_sheet(
    model: TypedAuxiliaryTransformer,
    data: tuple[AssetPhaseData, ...],
    *,
    study_dir: Path,
    manifest: dict[str, Any],
    output_path: Path,
    seed: int,
    selected_indices: tuple[int, ...],
    inference_batch_size: int,
    device: torch.device,
    max_development_rows_per_asset: int,
) -> dict[str, Any]:
    import pyarrow as pa
    import pyarrow.parquet as pq

    probability_names = tuple(
        f"p{index:02d}" for index in range(CLASS_COUNT)
    )
    schema_fields = [
        pa.field("stable_row_id", pa.uint64()),
        pa.field("seed", pa.int16()),
        pa.field("asset", pa.string()),
        pa.field("asset_id", pa.int8()),
        pa.field("anchor_row_index", pa.int64()),
        pa.field("anchor_timestamp_ns", pa.int64()),
        pa.field("available_at_ns", pa.int64()),
        pa.field("session_date", pa.date32()),
        pa.field("minute_of_session", pa.int16()),
        pa.field("horizon_minutes", pa.int16()),
        pa.field("target_valid", pa.bool_()),
        pa.field("target_timestamp_ns", pa.int64()),
        pa.field("target_available_at_ns", pa.int64()),
        pa.field("anchor_raw_close", pa.float64()),
        pa.field("target_raw_close", pa.float64()),
        pa.field("realized_log_return", pa.float64()),
        pa.field("target_class", pa.int16()),
        pa.field("expected_return", pa.float32()),
        pa.field("expected_positive_probability", pa.float32()),
        pa.field("entropy", pa.float32()),
        pa.field("max_probability", pa.float32()),
        pa.field("top_class", pa.int8()),
    ]
    schema_fields.extend(
        pa.field(name, pa.float32()) for name in probability_names
    )
    schema = pa.schema(schema_fields)

    temporary = output_path.with_suffix(output_path.suffix + ".partial")
    if output_path.exists() or temporary.exists():
        raise FileExistsError(f"prediction sheet already exists: {output_path}")
    output_path.parent.mkdir(parents=True, exist_ok=True)
    writer = pq.ParquetWriter(
        temporary,
        schema,
        compression="zstd",
        use_dictionary=("asset",),
    )
    horizons = EVALUATION_HORIZONS.astype(np.int64)
    horizon_count = len(horizons)
    by_asset: dict[str, Any] = {}
    total_rows = 0
    try:
        for item in data:
            test = _date_mask(item.bundle["session_date"], DEVELOPMENT_TEST)
            test_positions = np.flatnonzero(test)
            if (
                max_development_rows_per_asset > 0
                and len(test_positions) > max_development_rows_per_asset
            ):
                selected = np.linspace(
                    0,
                    len(test_positions) - 1,
                    max_development_rows_per_asset,
                    dtype=np.int64,
                )
                test_positions = test_positions[selected]
            rows = item.bundle["row_index"][test_positions]
            returns = item.bundle["y"][test_positions]
            target_mask = item.bundle["target_mask"][test_positions]
            target_labels = encode_returns(
                returns,
                target_mask,
                item.class_map.edges,
            )
            raw_close = _raw_close_vector(study_dir, manifest, item)
            nll_sum = np.zeros(horizon_count, dtype=np.float64)
            nll_count = np.zeros(horizon_count, dtype=np.int64)
            asset_output_rows = 0
            for start in range(0, len(rows), inference_batch_size):
                stop = min(start + inference_batch_size, len(rows))
                local_rows = rows[start:stop]
                probabilities = _predict(
                    model,
                    item,
                    rows=local_rows,
                    selected_indices=selected_indices,
                    batch_size=inference_batch_size,
                    device=device,
                    query_horizons=EVALUATION_HORIZONS,
                )
                count = len(local_rows)
                flat_count = count * horizon_count
                flat_probabilities = probabilities.reshape(
                    flat_count,
                    CLASS_COUNT,
                )
                local_mask = target_mask[start:stop]
                local_returns = returns[start:stop]
                local_labels = target_labels[start:stop]
                for horizon_index in range(horizon_count):
                    valid = local_mask[:, horizon_index]
                    if np.any(valid):
                        selected_probability = probabilities[
                            valid,
                            horizon_index,
                            local_labels[valid, horizon_index],
                        ]
                        nll_sum[horizon_index] += float(
                            -np.log(
                                np.clip(
                                    selected_probability,
                                    1e-12,
                                    1.0,
                                )
                            ).sum()
                        )
                        nll_count[horizon_index] += int(np.sum(valid))

                repeated_rows = np.repeat(local_rows, horizon_count)
                tiled_horizons = np.tile(horizons, count)
                repeated_mask = local_mask.reshape(-1)
                candidate = (
                    local_rows[:, None] + horizons[None, :]
                ).reshape(-1)
                safe_candidate = np.minimum(
                    candidate,
                    len(item.features["timestamp_ns"]) - 1,
                )
                target_timestamp = item.features["timestamp_ns"][
                    safe_candidate
                ].astype(np.int64, copy=True)
                target_timestamp[~repeated_mask] = -1
                target_close = raw_close[safe_candidate].copy()
                target_close[~repeated_mask] = np.nan
                anchor_close = np.repeat(
                    raw_close[local_rows],
                    horizon_count,
                )
                flat_returns = local_returns.reshape(-1)
                flat_labels = local_labels.reshape(-1)
                expected_return = np.einsum(
                    "nhc,hc->nh",
                    probabilities.astype(np.float64),
                    item.class_map.return_means,
                ).reshape(-1)
                expected_positive = np.einsum(
                    "nhc,hc->nh",
                    probabilities.astype(np.float64),
                    item.class_map.positive_rates,
                ).reshape(-1)
                entropy = -np.sum(
                    probabilities
                    * np.log(np.clip(probabilities, 1e-12, 1.0)),
                    axis=2,
                ).reshape(-1)
                stable = (
                    np.uint64(item.asset_id + 1) * np.uint64(1 << 56)
                    + repeated_rows.astype(np.uint64)
                )
                columns: dict[str, pa.Array] = {
                    "stable_row_id": pa.array(stable, type=pa.uint64()),
                    "seed": pa.array(
                        np.full(flat_count, seed, dtype=np.int16)
                    ),
                    "asset": pa.array([item.asset] * flat_count),
                    "asset_id": pa.array(
                        np.full(
                            flat_count,
                            item.asset_id,
                            dtype=np.int8,
                        )
                    ),
                    "anchor_row_index": pa.array(repeated_rows),
                    "anchor_timestamp_ns": pa.array(
                        np.repeat(
                            item.features["timestamp_ns"][local_rows],
                            horizon_count,
                        )
                    ),
                    "available_at_ns": pa.array(
                        np.repeat(
                            item.features["available_at_ns"][local_rows],
                            horizon_count,
                        )
                    ),
                    "session_date": pa.array(
                        np.repeat(
                            item.features["session_date"][local_rows],
                            horizon_count,
                        ),
                        type=pa.date32(),
                    ),
                    "minute_of_session": pa.array(
                        np.repeat(
                            item.features["minute_of_session"][local_rows],
                            horizon_count,
                        )
                    ),
                    "horizon_minutes": pa.array(
                        tiled_horizons.astype(np.int16)
                    ),
                    "target_valid": pa.array(repeated_mask),
                    "target_timestamp_ns": pa.array(
                        target_timestamp,
                        mask=~repeated_mask,
                    ),
                    "target_available_at_ns": pa.array(
                        target_timestamp + MINUTE_NS,
                        mask=~repeated_mask,
                    ),
                    "anchor_raw_close": pa.array(anchor_close),
                    "target_raw_close": pa.array(
                        target_close,
                        mask=~repeated_mask,
                    ),
                    "realized_log_return": pa.array(
                        flat_returns,
                        mask=~repeated_mask,
                    ),
                    "target_class": pa.array(
                        flat_labels,
                        mask=~repeated_mask,
                    ),
                    "expected_return": pa.array(
                        expected_return.astype(np.float32)
                    ),
                    "expected_positive_probability": pa.array(
                        expected_positive.astype(np.float32)
                    ),
                    "entropy": pa.array(entropy.astype(np.float32)),
                    "max_probability": pa.array(
                        flat_probabilities.max(axis=1).astype(np.float32)
                    ),
                    "top_class": pa.array(
                        flat_probabilities.argmax(axis=1).astype(np.int8)
                    ),
                }
                for class_index, name in enumerate(probability_names):
                    columns[name] = pa.array(
                        flat_probabilities[:, class_index]
                    )
                writer.write_table(
                    pa.Table.from_pydict(columns, schema=schema),
                    row_group_size=flat_count,
                )
                asset_output_rows += flat_count
                total_rows += flat_count

            if np.any(nll_count == 0):
                raise ValueError(
                    f"{item.asset} development sheet lacks a horizon"
                )
            by_horizon = nll_sum / nll_count
            by_asset[item.asset] = {
                "anchor_rows": int(len(rows)),
                "macro_nll": float(np.mean(by_horizon)),
                "nll_by_horizon": {
                    str(int(horizon)): float(by_horizon[index])
                    for index, horizon in enumerate(EVALUATION_HORIZONS)
                },
                "output_rows": asset_output_rows,
                "valid_targets_by_horizon": {
                    str(int(horizon)): int(nll_count[index])
                    for index, horizon in enumerate(EVALUATION_HORIZONS)
                },
            }
    finally:
        writer.close()
    os.replace(temporary, output_path)
    return {
        "by_asset": by_asset,
        "equal_asset_macro_nll": float(
            np.mean([value["macro_nll"] for value in by_asset.values()])
        ),
        "output_rows": total_rows,
        "parquet_bytes": output_path.stat().st_size,
        "parquet_sha256": _sha256(output_path),
    }


def _write_calibration(
    path: Path,
    data: tuple[AssetPhaseData, ...],
) -> str:
    if path.exists():
        raise FileExistsError(f"calibration artifact exists: {path}")
    arrays: dict[str, np.ndarray] = {
        "evaluation_horizons": EVALUATION_HORIZONS,
        "canonical_horizons": CANONICAL_HORIZONS,
    }
    for item in data:
        prefix = item.asset.lower()
        arrays[f"{prefix}_token_center"] = item.center
        arrays[f"{prefix}_token_scale"] = item.scale
        arrays[f"{prefix}_auxiliary_center"] = item.auxiliary_center
        arrays[f"{prefix}_auxiliary_scale"] = item.auxiliary_scale
        for name, value in class_map_to_arrays(item.class_map).items():
            arrays[f"{prefix}_{name}"] = value
    np.savez_compressed(path, **arrays)
    return _sha256(path)


def _run_manifest(
    *,
    config: dict[str, Any],
    project_root: Path,
) -> dict[str, Any]:
    return {
        "command": sys.argv,
        "config": config,
        "created_at_utc": datetime.now(timezone.utc).isoformat(),
        "environment": {
            "machine": platform.machine(),
            "numpy": np.__version__,
            "platform": platform.platform(),
            "python": sys.version,
            "torch": torch.__version__,
            "torch_mps_available": torch.backends.mps.is_available(),
        },
        "provenance_sha256": {
            relative: _sha256(project_root / relative)
            for relative in PROVENANCE_PATHS
        },
        "runner_version": RUNNER_VERSION,
    }


def run(
    *,
    study_dir: Path,
    checkpoint_dir: Path,
    output_dir: Path,
    seed: int,
    epochs: int = DEFAULT_EPOCHS,
    batch_size: int = DEFAULT_BATCH_SIZE,
    gradient_accumulation_steps: int = DEFAULT_GRADIENT_ACCUMULATION,
    inference_batch_size: int = DEFAULT_INFERENCE_BATCH_SIZE,
    torch_threads: int = 6,
    device_name: str = "mps",
    max_train_rows_per_asset: int = 0,
    max_development_rows_per_asset: int = 0,
) -> dict[str, Any]:
    started = time.monotonic()
    if seed not in (0, 1, 2):
        raise ValueError("v7 seed must be 0, 1, or 2")
    if epochs <= 0 or batch_size <= 0 or inference_batch_size <= 0:
        raise ValueError("v7 training dimensions must be positive")
    if gradient_accumulation_steps <= 0:
        raise ValueError("v7 gradient accumulation must be positive")
    if max_train_rows_per_asset < 0 or max_development_rows_per_asset < 0:
        raise ValueError("v7 row caps cannot be negative")
    if (
        batch_size * gradient_accumulation_steps
        != DEFAULT_BATCH_SIZE * DEFAULT_GRADIENT_ACCUMULATION
    ):
        raise ValueError("v7 effective batch size is frozen at 4096")
    if device_name == "mps" and not torch.backends.mps.is_available():
        raise ValueError("v7 requested unavailable MPS")

    checkpoint_results_path = checkpoint_dir / "results.json"
    checkpoint_manifest_path = checkpoint_dir / "run_manifest.json"
    checkpoint_path = checkpoint_dir / "discovery_2_model.pt"
    checkpoint_results = json.loads(checkpoint_results_path.read_text())
    checkpoint_manifest = json.loads(checkpoint_manifest_path.read_text())
    if int(checkpoint_results.get("seed", -1)) != seed:
        raise ValueError("v7 seed and locked checkpoint seed differ")
    expected_checkpoint = {
        "asset_conditioning": "early_add",
        "context_length": DEFAULT_CONTEXT_LENGTH,
        "feature_set": FEATURE_SET,
        "max_train_rows_per_asset": 100_000,
        "rotary_base": DEFAULT_ROTARY_BASE,
        "supervised_horizons": [
            int(value) for value in CANONICAL_HORIZONS
        ],
        "target_parameters": DEFAULT_TARGET_PARAMETERS,
    }
    for key, expected in expected_checkpoint.items():
        if checkpoint_results.get(key) != expected:
            raise ValueError(
                f"locked checkpoint setting {key} differs from v7"
            )
    manifest_config = checkpoint_manifest.get("config", {})
    if manifest_config.get("seed") != seed:
        raise ValueError("locked checkpoint manifest has the wrong seed")

    manifest, loaded = _load_selection_data(study_dir)
    selected_indices = _selected_feature_indices()
    canonical_indices = _canonical_indices()
    checkpoint = torch.load(
        checkpoint_path,
        map_location="cpu",
        weights_only=False,
    )
    config = {
        "assets": list(SELECTION_ASSETS),
        "auxiliary_feature_names": list(AUXILIARY_FEATURE_NAMES),
        "batch_size": batch_size,
        "checkpoint_dir": str(checkpoint_dir),
        "checkpoint_model_sha256": _sha256(checkpoint_path),
        "checkpoint_results_sha256": _sha256(checkpoint_results_path),
        "context_length": DEFAULT_CONTEXT_LENGTH,
        "development_test": list(DEVELOPMENT_TEST),
        "device_name": device_name,
        "epochs": epochs,
        "evaluation_horizons": EVALUATION_HORIZONS.tolist(),
        "feature_names": list(FEATURE_NAMES),
        "gradient_accumulation_steps": gradient_accumulation_steps,
        "inference_batch_size": inference_batch_size,
        "max_development_rows_per_asset": (
            max_development_rows_per_asset
        ),
        "max_train_rows_per_asset": max_train_rows_per_asset,
        "optimizer_recipe": "adamw_constant",
        "refit_interval": [
            DEVELOPMENT_START.isoformat(),
            VALIDATION_END.isoformat(),
        ],
        "rotary_base": DEFAULT_ROTARY_BASE,
        "seed": seed,
        "snapshot_sha256": manifest["snapshot_sha256"],
        "supervised_horizons": CANONICAL_HORIZONS.tolist(),
        "target_parameters": DEFAULT_TARGET_PARAMETERS,
        "torch_threads": torch_threads,
        "train_interval": [
            DEVELOPMENT_START.isoformat(),
            TRAIN_END.isoformat(),
        ],
        "validation_interval": [
            TRAIN_END.isoformat(),
            VALIDATION_END.isoformat(),
        ],
    }
    output_dir.mkdir(parents=True, exist_ok=True)
    manifest_path = output_dir / "run_manifest.json"
    results_path = output_dir / "results.json"
    prediction_path = output_dir / f"development_predictions_seed{seed}.parquet"
    checkpoint_output_path = output_dir / f"model_seed{seed}.pt"
    calibration_path = output_dir / f"calibration_seed{seed}.npz"
    guarded = (
        manifest_path,
        results_path,
        prediction_path,
        checkpoint_output_path,
        calibration_path,
    )
    if any(path.exists() for path in guarded):
        raise FileExistsError("v7 output directory is not immutable-empty")
    project_root = Path(__file__).resolve().parents[3]
    manifest_path.write_text(
        json.dumps(
            _run_manifest(config=config, project_root=project_root),
            indent=2,
            sort_keys=True,
        )
        + "\n"
    )

    _configure(seed, torch_threads)
    device = torch.device(device_name)
    selection_data = tuple(
        _prepare_phase_data(
            asset=asset,
            asset_id=asset_id,
            features=loaded[asset][0],
            bundle=loaded[asset][1],
            train_interval=(
                DEVELOPMENT_START.isoformat(),
                TRAIN_END.isoformat(),
            ),
            validation_interval=(
                TRAIN_END.isoformat(),
                VALIDATION_END.isoformat(),
            ),
            selected_indices=selected_indices,
            canonical_indices=canonical_indices,
            max_train_rows_per_asset=max_train_rows_per_asset,
        )
        for asset_id, asset in enumerate(SELECTION_ASSETS)
    )
    selection_model = _build_model(checkpoint, device=device)
    selection_training = _train(
        selection_model,
        selection_data,
        seed=seed,
        epochs=epochs,
        batch_size=batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        inference_batch_size=inference_batch_size,
        selected_indices=selected_indices,
        device=device,
        select_on_validation=True,
    )
    selected_epoch = int(selection_training["best_epoch"])
    del selection_model, selection_data
    if device.type == "mps":
        torch.mps.empty_cache()

    refit_data = tuple(
        _prepare_phase_data(
            asset=asset,
            asset_id=asset_id,
            features=loaded[asset][0],
            bundle=loaded[asset][1],
            train_interval=(
                DEVELOPMENT_START.isoformat(),
                VALIDATION_END.isoformat(),
            ),
            validation_interval=None,
            selected_indices=selected_indices,
            canonical_indices=canonical_indices,
            max_train_rows_per_asset=max_train_rows_per_asset,
        )
        for asset_id, asset in enumerate(SELECTION_ASSETS)
    )
    refit_model = _build_model(checkpoint, device=device)
    refit_training = _train(
        refit_model,
        refit_data,
        seed=seed + 70_001,
        epochs=selected_epoch,
        batch_size=batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        inference_batch_size=inference_batch_size,
        selected_indices=selected_indices,
        device=device,
        select_on_validation=False,
    )
    torch.save(
        {
            "assets": list(SELECTION_ASSETS),
            "auxiliary_feature_names": AUXILIARY_FEATURE_NAMES,
            "context_length": DEFAULT_CONTEXT_LENGTH,
            "feature_names": FEATURE_NAMES,
            "refit_interval": config["refit_interval"],
            "rotary_base": DEFAULT_ROTARY_BASE,
            "seed": seed,
            "selected_epoch": selected_epoch,
            "state_dict": {
                name: value.detach().cpu()
                for name, value in refit_model.state_dict().items()
            },
            "supervised_horizons": tuple(
                int(value) for value in CANONICAL_HORIZONS
            ),
        },
        checkpoint_output_path,
    )
    calibration_sha256 = _write_calibration(
        calibration_path,
        refit_data,
    )
    development = _write_prediction_sheet(
        refit_model,
        refit_data,
        study_dir=study_dir,
        manifest=manifest,
        output_path=prediction_path,
        seed=seed,
        selected_indices=selected_indices,
        inference_batch_size=inference_batch_size,
        device=device,
        max_development_rows_per_asset=max_development_rows_per_asset,
    )
    result = {
        **config,
        "calibration_sha256": calibration_sha256,
        "checkpoint_sha256": _sha256(checkpoint_output_path),
        "development_test_metrics": development,
        "duration_seconds": time.monotonic() - started,
        "parameter_count": sum(
            parameter.numel()
            for parameter in refit_model.parameters()
            if parameter.requires_grad
        ),
        "prediction_path": str(prediction_path),
        "refit_training": refit_training,
        "run_manifest_sha256": _sha256(manifest_path),
        "selected_epoch": selected_epoch,
        "selection_training": selection_training,
    }
    results_path.write_text(
        json.dumps(result, indent=2, sort_keys=True) + "\n"
    )
    return result


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Train one immutable forecasting-v7 ensemble seed."
    )
    parser.add_argument("--study-dir", type=Path, required=True)
    parser.add_argument("--checkpoint-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--seed", type=int, choices=(0, 1, 2), required=True)
    parser.add_argument("--epochs", type=int, default=DEFAULT_EPOCHS)
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    parser.add_argument(
        "--gradient-accumulation-steps",
        type=int,
        default=DEFAULT_GRADIENT_ACCUMULATION,
    )
    parser.add_argument(
        "--inference-batch-size",
        type=int,
        default=DEFAULT_INFERENCE_BATCH_SIZE,
    )
    parser.add_argument("--torch-threads", type=int, default=6)
    parser.add_argument("--device-name", default="mps")
    parser.add_argument("--max-train-rows-per-asset", type=int, default=0)
    parser.add_argument(
        "--max-development-rows-per-asset",
        type=int,
        default=0,
    )
    args = parser.parse_args()
    print(json.dumps(run(**vars(args)), sort_keys=True))


if __name__ == "__main__":
    main()