from __future__ import annotations

import base64
import io
import json
import tempfile
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, cast

import numpy as np
import pandas as pd
import polars as pl
import torch
import torch.nn as nn
import xgboost as xgb
from sklearn.isotonic import IsotonicRegression
from sklearn.metrics import brier_score_loss, f1_score
from torch.utils.data import DataLoader, TensorDataset

from uforia.storage import DATASET_SPECS, DatasetName, DatasetStore

from .builders import summarize_transition_probabilities
from .common import horizon_to_seconds, parse_feature_payload, utc_now
from .constants import FUSION_SOURCE, MODEL_VERSION, PAPER_SOURCE, TABULAR_SOURCE, TCN_SOURCE
SEQUENCE_LENGTH = 32
TCN_EMBEDDING_DIM = 64
MIN_TRAIN_ROWS = 96
ROUND_TRIP_COST_BPS = 3.0
CONFIDENCE_THRESHOLD = 0.50
EDGE_THRESHOLD_BPS = 0.25
FEATURE_SCHEMA_VERSION = DATASET_SPECS[DatasetName.MICROSTRUCTURE_FEATURE_SERIES].version
LABEL_SCHEMA_VERSION = DATASET_SPECS[DatasetName.MICROSTRUCTURE_LABEL_SERIES].version
OPTIONAL_FEATURE_COLUMNS = {
    "basis_bps",
    "funding_rate",
    "oi_value",
    "liquidation_volume",
    "iv_atm",
    "iv_skew_25d",
    "iv_rv_gap",
}
TABULAR_FEATURE_COLUMNS = [
    "mid_price",
    "return_1",
    "return_5",
    "realized_vol",
    "book_imbalance",
    "microprice_tilt",
    "order_flow_imbalance",
    "cumulative_delta",
    "cum_delta_10",
    "cum_delta_30",
    "cum_delta_50",
    "basis_bps",
    "funding_rate",
    "oi_value",
    "liquidation_volume",
    "iv_atm",
    "iv_skew_25d",
    "cross_venue_spread",
    "iv_rv_gap",
    "spread_rel",
    "depth_total",
    "volume_delta",
    "ofi_mean_12",
    "ofi_std_12",
    "imbalance_mean_12",
    "spread_rel_mean_12",
    "depth_change_1",
    "microprice_tilt_change_5",
    "microprice_gap",
]
SEQUENCE_FEATURE_COLUMNS = [
    "return_1",
    "return_5",
    "realized_vol",
    "book_imbalance",
    "microprice_tilt",
    "order_flow_imbalance",
    "cum_delta_10",
    "cum_delta_30",
    "cum_delta_50",
    "basis_bps",
    "funding_rate",
    "iv_atm",
    "iv_skew_25d",
    "cross_venue_spread",
    "iv_rv_gap",
    "spread_rel",
    "depth_total",
    "volume_delta",
    "ofi_mean_12",
    "ofi_std_12",
    "imbalance_mean_12",
    "spread_rel_mean_12",
    "depth_change_1",
    "microprice_tilt_change_5",
    "microprice_gap",
]


class PurgedTimeSeriesCV:
    def __init__(self, n_splits: int = 1, embargo_bars: int = SEQUENCE_LENGTH - 1) -> None:
        self.n_splits = n_splits
        self.embargo_bars = embargo_bars

    def split(self, features: pd.DataFrame, y: np.ndarray | None = None, groups: Any = None):
        count = len(features)
        if count <= 1:
            return
        test_size = max(1, count // (self.n_splits + 1))
        for split_index in range(self.n_splits):
            test_start = count - (self.n_splits - split_index) * test_size
            test_end = min(count, test_start + test_size)
            train_end = max(0, test_start - self.embargo_bars)
            if train_end <= SEQUENCE_LENGTH:
                continue
            yield np.arange(0, train_end), np.arange(test_start, test_end)


class CausalChomp1d(nn.Module):
    def __init__(self, chomp_size: int) -> None:
        super().__init__()
        self.chomp_size = chomp_size

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.chomp_size == 0:
            return x
        return x[:, :, : -self.chomp_size]


class TemporalBlock(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        *,
        kernel_size: int,
        dilation: int,
        dropout: float,
    ) -> None:
        super().__init__()
        padding = (kernel_size - 1) * dilation
        self.net = nn.Sequential(
            nn.Conv1d(
                in_channels,
                out_channels,
                kernel_size,
                padding=padding,
                dilation=dilation,
            ),
            CausalChomp1d(padding),
            nn.BatchNorm1d(out_channels),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Conv1d(
                out_channels,
                out_channels,
                kernel_size,
                padding=padding,
                dilation=dilation,
            ),
            CausalChomp1d(padding),
            nn.BatchNorm1d(out_channels),
            nn.GELU(),
            nn.Dropout(dropout),
        )
        self.downsample = (
            nn.Conv1d(in_channels, out_channels, kernel_size=1)
            if in_channels != out_channels
            else nn.Identity()
        )
        self.activation = nn.GELU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = self.downsample(x)
        return self.activation(self.net(x) + residual)


class TcnClassifier(nn.Module):
    def __init__(self, input_dim: int, embedding_dim: int = TCN_EMBEDDING_DIM) -> None:
        super().__init__()
        channels = [32, 48, 64, 64, 64, embedding_dim]
        dilations = [1, 2, 4, 8, 16, 32]
        layers: list[nn.Module] = []
        in_channels = input_dim
        for out_channels, dilation in zip(channels, dilations, strict=True):
            layers.append(
                TemporalBlock(
                    in_channels,
                    out_channels,
                    kernel_size=3,
                    dilation=dilation,
                    dropout=0.1,
                )
            )
            in_channels = out_channels
        self.backbone = nn.Sequential(*layers)
        self.head = nn.Linear(embedding_dim, 3)

    def embed(self, x: torch.Tensor) -> torch.Tensor:
        features = self.backbone(x)
        return features[:, :, -1]

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        embedding = self.embed(x)
        logits = self.head(embedding)
        return logits, embedding


class FocalLoss(nn.Module):
    def __init__(self, gamma: float = 2.0, alpha: torch.Tensor | None = None) -> None:
        super().__init__()
        self.gamma = gamma
        self.alpha = alpha

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        log_probs = torch.log_softmax(logits, dim=1)
        probs = torch.softmax(logits, dim=1)
        ce = nn.functional.nll_loss(log_probs, targets, reduction="none")
        target_probs = probs.gather(1, targets.unsqueeze(1)).squeeze(1)
        focal_weight = (1.0 - target_probs).pow(self.gamma)
        if self.alpha is not None:
            alpha_weight = self.alpha.gather(0, targets)
        else:
            alpha_weight = torch.ones_like(target_probs)
        return (alpha_weight * focal_weight * ce).mean()


@dataclass(frozen=True)
class SequenceBatch:
    tabular: np.ndarray
    sequence: np.ndarray
    labels: np.ndarray
    forward_returns_bps: np.ndarray
    timestamps: list[pd.Timestamp]
    feature_columns: list[str]
    sequence_columns: list[str]
    mean_moves: dict[int, float]
    label_horizon_steps: int
    feature_schema_version: int
    label_schema_version: int


def train_microstructure_model(
    store: DatasetStore,
    *,
    asset: str,
    horizon: str,
    source: str = FUSION_SOURCE,
) -> tuple[pl.DataFrame, dict[str, Any]]:
    dataset = _load_training_frame(store, asset=asset, horizon=horizon)
    if dataset.empty or len(dataset) < MIN_TRAIN_ROWS:
        return pl.DataFrame(), {}

    batch = _build_training_batch(dataset)
    split = _validation_split(len(batch.labels), batch.label_horizon_steps)
    if split is None:
        return pl.DataFrame(), {}
    train_idx, val_idx = split

    tabular_model, tabular_best_iteration = _train_xgb(
        batch.tabular[train_idx],
        batch.labels[train_idx],
        batch.tabular[val_idx],
        batch.labels[val_idx],
    )
    tabular_val_probs = tabular_model.predict_proba(batch.tabular[val_idx])
    tabular_metrics = _classification_metrics(
        batch.labels[val_idx],
        tabular_val_probs,
        batch.forward_returns_bps[val_idx],
    )
    profitable_gate = tabular_metrics["paper_net_edge_bps"] > 0

    tcn_model = _train_tcn(
        batch.sequence[train_idx],
        batch.labels[train_idx],
        batch.sequence[val_idx],
        batch.labels[val_idx],
    )
    tcn_val_probs, _ = _predict_tcn(tcn_model, batch.sequence[val_idx])
    tcn_metrics = _classification_metrics(
        batch.labels[val_idx],
        tcn_val_probs,
        batch.forward_returns_bps[val_idx],
    )

    train_embed = _embed_tcn(tcn_model, batch.sequence[train_idx])
    val_embed = _embed_tcn(tcn_model, batch.sequence[val_idx])
    fusion_train = np.concatenate([batch.tabular[train_idx], train_embed], axis=1)
    fusion_val = np.concatenate([batch.tabular[val_idx], val_embed], axis=1)
    fusion_model, fusion_best_iteration = _train_xgb(
        fusion_train,
        batch.labels[train_idx],
        fusion_val,
        batch.labels[val_idx],
    )
    fusion_val_probs = fusion_model.predict_proba(fusion_val)
    calibrator = _fit_calibrator(fusion_val_probs, batch.labels[val_idx])
    calibrated_val_probs = _apply_calibrator(calibrator, fusion_val_probs)
    fusion_metrics = _classification_metrics(
        batch.labels[val_idx],
        calibrated_val_probs,
        batch.forward_returns_bps[val_idx],
    )
    validation_embargo_bars = max(SEQUENCE_LENGTH - 1, batch.label_horizon_steps)

    model = {
        "model_type": MODEL_VERSION,
        "asset": asset,
        "horizon": horizon,
        "source": source,
        "sequence_length": SEQUENCE_LENGTH,
        "tabular_feature_columns": batch.feature_columns,
        "sequence_feature_columns": batch.sequence_columns,
        "class_move_bps": {
            str(key): value for key, value in batch.mean_moves.items()
        },
        "tabular_model": tabular_model,
        "tcn_model": tcn_model,
        "fusion_model": fusion_model,
        "calibration": calibrator,
        "tabular_metrics": tabular_metrics,
        "tcn_metrics": tcn_metrics,
        "fusion_metrics": fusion_metrics,
        "profitable": profitable_gate and fusion_metrics["paper_net_edge_bps"] > 0,
        "lift_vs_tabular": fusion_metrics["mean_accuracy"] - tabular_metrics["mean_accuracy"],
        "feature_schema_version": batch.feature_schema_version,
        "label_schema_version": batch.label_schema_version,
        "validation_embargo_bars": validation_embargo_bars,
        "tabular_best_iteration": tabular_best_iteration,
        "fusion_best_iteration": fusion_best_iteration,
    }
    now = utc_now()
    run = pl.DataFrame(
        {
            "ts": [now],
            "date": [now.date()],
            "asset": [asset],
            "horizon": [horizon],
            "source": [source],
            "model_version": [MODEL_VERSION],
            "train_start": [batch.timestamps[0].to_pydatetime()],
            "train_end": [batch.timestamps[-1].to_pydatetime()],
            "split_count": [1],
            "cv_accuracy_mean": [fusion_metrics["mean_accuracy"]],
            "cv_accuracy_std": [0.0],
            "brier_score": [fusion_metrics["mean_brier_score"]],
            "calibration_error": [fusion_metrics["mean_brier_score"]],
            "lift_vs_tabular": [model["lift_vs_tabular"]],
            "profitable": [model["profitable"]],
            "artifact_uri": [""],
            "artifact_payload": [""],
            "feature_schema_version": [batch.feature_schema_version],
            "label_schema_version": [batch.label_schema_version],
            "validation_embargo_bars": [validation_embargo_bars],
            "tabular_best_iteration": [tabular_best_iteration],
            "fusion_best_iteration": [fusion_best_iteration],
            "quality_flag": ["ok" if model["profitable"] else "gated"],
        }
    )
    return run, model


def evaluate_microstructure_model(
    store: DatasetStore,
    model: dict[str, Any],
    *,
    asset: str,
    horizon: str,
    source: str = FUSION_SOURCE,
) -> pl.DataFrame:
    if "fusion_model" not in model:
        model = _load_exported_model(model)
    dataset = _load_training_frame(store, asset=asset, horizon=horizon)
    if dataset.empty or len(dataset) < MIN_TRAIN_ROWS:
        return pl.DataFrame()
    batch = _build_training_batch(dataset, feature_columns=model["tabular_feature_columns"])
    probs = _predict_fusion(model, batch)
    metrics = _classification_metrics(batch.labels, probs, batch.forward_returns_bps)
    now = utc_now()
    return pl.DataFrame(
        {
            "ts": [now],
            "date": [now.date()],
            "asset": [asset],
            "horizon": [horizon],
            "source": [source],
            "model_version": [MODEL_VERSION],
            "train_start": [batch.timestamps[0].to_pydatetime()],
            "train_end": [batch.timestamps[-1].to_pydatetime()],
            "test_start": [batch.timestamps[0].to_pydatetime()],
            "test_end": [batch.timestamps[-1].to_pydatetime()],
            "split_count": [1],
            "mean_accuracy": [metrics["mean_accuracy"]],
            "mean_brier_score": [metrics["mean_brier_score"]],
            "paper_sharpe": [metrics["paper_sharpe"]],
            "paper_net_edge_bps": [metrics["paper_net_edge_bps"]],
            "directional_f1": [metrics["directional_f1"]],
            "feature_schema_version": [model.get("feature_schema_version", FEATURE_SCHEMA_VERSION)],
            "label_schema_version": [model.get("label_schema_version", LABEL_SCHEMA_VERSION)],
            "validation_embargo_bars": [model.get("validation_embargo_bars", 0)],
            "quality_flag": ["ok" if model.get("profitable", False) else "gated"],
        }
    )


def export_microstructure_model(
    model: dict[str, Any],
    *,
    asset: str,
    horizon: str,
    output_dir: Any | None = None,
) -> dict[str, Any]:
    del output_dir
    tabular_bytes = _dump_xgb_model_bytes(cast(xgb.XGBClassifier, model["tabular_model"]))
    fusion_bytes = _dump_xgb_model_bytes(cast(xgb.XGBClassifier, model["fusion_model"]))
    tcn_state = io.BytesIO()
    torch.save(cast(TcnClassifier, model["tcn_model"]).state_dict(), tcn_state)
    payload = {
        "asset": asset,
        "horizon": horizon,
        "model_type": MODEL_VERSION,
        "source": FUSION_SOURCE,
        "artifact_uri": f"db://microstructure_training_run/{asset.lower()}/{horizon}/{MODEL_VERSION}",
        "tabular_model_b64": base64.b64encode(tabular_bytes).decode(),
        "fusion_model_b64": base64.b64encode(fusion_bytes).decode(),
        "tcn_state_b64": base64.b64encode(tcn_state.getvalue()).decode(),
        "sequence_length": model["sequence_length"],
        "tabular_feature_columns": model["tabular_feature_columns"],
        "sequence_feature_columns": model["sequence_feature_columns"],
        "class_move_bps": model["class_move_bps"],
        "calibration": model["calibration"],
        "tabular_metrics": model["tabular_metrics"],
        "tcn_metrics": model["tcn_metrics"],
        "fusion_metrics": model["fusion_metrics"],
        "lift_vs_tabular": model["lift_vs_tabular"],
        "profitable": model["profitable"],
        "feature_schema_version": model["feature_schema_version"],
        "label_schema_version": model["label_schema_version"],
        "validation_embargo_bars": model["validation_embargo_bars"],
        "tabular_best_iteration": model["tabular_best_iteration"],
        "fusion_best_iteration": model["fusion_best_iteration"],
    }
    return payload


def generate_microstructure_model_output(
    store: DatasetStore,
    artifact: dict[str, Any],
    *,
    asset: str,
    horizon: str,
    source: str = FUSION_SOURCE,
) -> tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame]:
    if artifact.get("model_type") != MODEL_VERSION:
        raise ValueError(f"Incompatible model artifact: {artifact.get('model_type')}")
    dataset = _load_training_frame(store, asset=asset, horizon=horizon)
    if dataset.empty:
        return pl.DataFrame(), pl.DataFrame(), pl.DataFrame()
    batch = _build_training_batch(
        dataset,
        feature_columns=artifact["tabular_feature_columns"],
        sequence_columns=artifact["sequence_feature_columns"],
    )
    if batch.labels.size == 0:
        return pl.DataFrame(), pl.DataFrame(), pl.DataFrame()

    model = _load_exported_model(artifact)
    probs = _predict_fusion(model, batch)
    class_move = np.asarray(
        [artifact["class_move_bps"].get(str(i), 0.0) for i in range(3)],
        dtype=np.float64,
    )
    expected_move = probs @ class_move
    confidence = np.max(probs, axis=1)
    up_move = max(float(class_move[2]), 0.0)
    down_move = abs(min(float(class_move[0]), 0.0))
    long_edge = (probs[:, 2] * up_move) - (probs[:, 0] * down_move) - ROUND_TRIP_COST_BPS
    short_edge = (probs[:, 0] * down_move) - (probs[:, 2] * up_move) - ROUND_TRIP_COST_BPS
    expected_edge = np.maximum.reduce([long_edge, short_edge, np.zeros_like(long_edge)])
    tradability = np.max(probs, axis=1) * np.clip(expected_edge / 4.0, 0.0, 3.0)
    confidence_bucket = [
        "high" if value >= 0.62 else "medium" if value >= 0.5 else "low"
        for value in confidence
    ]
    regime = store.read(
        DatasetName.MICROSTRUCTURE_REGIME_SERIES,
        filters={"asset": asset, "source": "tcn_boost_regime_v1"},
    ).sort("ts")
    regime_lookup = {
        cast(datetime, row["ts"]): cast(str, row["regime_state"])
        for row in regime.select(["ts", "regime_state"]).iter_rows(named=True)
    }
    book_quality = store.read(
        DatasetName.MICROSTRUCTURE_BOOK_STATE,
        filters={"asset": asset},
    ).sort("ts")
    book_quality_lookup = {
        cast(datetime, row["ts"]): cast(str, row["book_quality"])
        for row in book_quality.select(["ts", "book_quality"]).iter_rows(named=True)
    }
    timestamps = [ts.to_pydatetime() for ts in batch.timestamps]
    blocked_reasons: list[str | None] = []
    model_families: list[str] = []
    for edge, conf, _prob_row, ts in zip(expected_edge, confidence, probs, timestamps, strict=True):
        quality = book_quality_lookup.get(ts, "unknown")
        if quality != "complete":
            blocked_reasons.append("degraded_book")
        elif edge <= EDGE_THRESHOLD_BPS:
            blocked_reasons.append("insufficient_edge")
        elif conf < CONFIDENCE_THRESHOLD:
            blocked_reasons.append("low_confidence")
        else:
            blocked_reasons.append(None)
        model_families.append("fusion")
    outputs = pl.DataFrame(
        {
            "ts": timestamps,
            "date": [ts.date() for ts in timestamps],
            "asset": [asset] * len(timestamps),
            "horizon": [horizon] * len(timestamps),
            "source": [source] * len(timestamps),
            "prob_down": probs[:, 0].tolist(),
            "prob_flat": probs[:, 1].tolist(),
            "prob_up": probs[:, 2].tolist(),
            "expected_move_bps": expected_move.tolist(),
            "expected_edge_bps": expected_edge.tolist(),
            "tradability_score": tradability.tolist(),
            "confidence": confidence.tolist(),
            "confidence_bucket": confidence_bucket,
            "regime_state": [
                regime_lookup.get(ts, "neutral") for ts in timestamps
            ],
            "blocked_reason": blocked_reasons,
            "model_version": [MODEL_VERSION] * len(timestamps),
            "model_family": model_families,
            "quality_flag": [
                "ok" if reason is None else "blocked" for reason in blocked_reasons
            ],
        }
    )
    signal_rows = outputs.select(
        [
            "ts",
            "date",
            "asset",
            "horizon",
            pl.lit(PAPER_SOURCE).alias("source"),
            pl.when((pl.col("prob_up") > 0.5) & (pl.col("tradability_score") >= 0.25))
            .then(pl.lit("long_bias"))
            .when((pl.col("prob_down") > 0.5) & (pl.col("tradability_score") >= 0.25))
            .then(pl.lit("short_bias"))
            .otherwise(pl.lit("flat"))
            .alias("signal_state"),
            pl.when((pl.col("prob_up") > 0.5) & (pl.col("tradability_score") >= 0.25))
            .then(pl.lit("long"))
            .when((pl.col("prob_down") > 0.5) & (pl.col("tradability_score") >= 0.25))
            .then(pl.lit("short"))
                    .otherwise(pl.lit("flat"))
                    .alias("paper_position"),
                    pl.col("expected_edge_bps"),
                    (pl.col("blocked_reason").is_null()).alias("cost_filter_passed"),
                    (pl.col("quality_flag") == "ok").alias("quality_filter_passed"),
                    "blocked_reason",
                    "confidence",
                    "tradability_score",
                    "quality_flag",
        ]
    )
    states = signal_rows["signal_state"].to_list()
    probs_map, median_duration = summarize_transition_probabilities(states)
    regime_transition = pl.DataFrame(
        {
            "ts": [timestamps[-1]],
            "date": [timestamps[-1].date()],
            "asset": [asset],
            "source": ["vol_regime_transitions"],
            "signal_state": [states[-1] if states else "flat"],
            "prior_state": [states[-2] if len(states) > 1 else None],
            "next_state_probability_long_vol_bias": [float(probs_map.get("long_bias", 0.0))],
            "next_state_probability_neutral": [float(probs_map.get("flat", 0.0))],
            "next_state_probability_short_vol_bias": [float(probs_map.get("short_bias", 0.0))],
            "median_state_duration_hours": [median_duration],
            "quality_flag": ["ok"],
        }
    )
    return outputs, signal_rows, regime_transition


def _load_training_frame(store: DatasetStore, *, asset: str, horizon: str) -> pd.DataFrame:
    features = store.read(
        DatasetName.MICROSTRUCTURE_FEATURE_SERIES,
        filters={"asset": asset, "horizon": horizon},
    ).sort("ts")
    labels = store.read(
        DatasetName.MICROSTRUCTURE_LABEL_SERIES,
        filters={"asset": asset, "horizon": horizon},
    ).sort("ts")
    if features.is_empty() or labels.is_empty():
        return pd.DataFrame()
    frame = features.join(
        labels.select(
            [
                "ts",
                "label",
                "forward_return_bps",
                "forward_microprice_return_bps",
                pl.col("bar_size_secs").alias("label_bar_size_secs"),
                "label_horizon_steps",
                "horizon_seconds",
                pl.col("quality_flag").alias("label_quality_flag"),
            ]
        ),
        on="ts",
        how="inner",
    ).with_columns(pl.col("quality_flag").alias("feature_quality_flag")).filter(
        pl.col("feature_quality_flag").is_in(["ok", "partial_context"])
        & (pl.col("label_quality_flag") == "ok")
    )
    rows: list[dict[str, Any]] = []
    for row in frame.iter_rows(named=True):
        payload = parse_feature_payload(cast(str | None, row.get("feature_payload")))
        for key, value in row.items():
            if key in {
                "ts",
                "date",
                "asset",
                "horizon",
                "source",
                "quality_flag",
                "feature_quality_flag",
                "label_quality_flag",
                "feature_payload",
                "run_id",
                "run_ts",
                "logical_ts",
                "write_mode",
            }:
                continue
            if value is None:
                continue
            try:
                payload[key] = float(value)
            except (TypeError, ValueError):
                continue
        payload["ts"] = row["ts"]
        payload["asset"] = row["asset"]
        payload["horizon"] = row["horizon"]
        payload["label"] = row["label"]
        rows.append(payload)
    dataset = pd.DataFrame(rows)
    if dataset.empty:
        return dataset
    dataset["ts"] = pd.to_datetime(dataset["ts"], utc=True)
    numeric_cols = [
        column
        for column in dataset.columns
        if column not in {"ts", "asset", "horizon", "label"}
    ]
    dataset[numeric_cols] = dataset[numeric_cols].apply(pd.to_numeric, errors="coerce")
    dataset = dataset.replace([np.inf, -np.inf], np.nan)
    horizon_seconds = float(horizon_to_seconds(horizon))
    if "horizon_seconds" not in dataset.columns:
        dataset["horizon_seconds"] = horizon_seconds
    else:
        dataset["horizon_seconds"] = dataset["horizon_seconds"].fillna(horizon_seconds)
    if "bar_size_secs" not in dataset.columns and "label_bar_size_secs" in dataset.columns:
        dataset["bar_size_secs"] = dataset["label_bar_size_secs"]
    inferred_bar_size = None
    if "bar_size_secs" in dataset.columns:
        non_null_bar_sizes = dataset["bar_size_secs"].dropna()
        if not non_null_bar_sizes.empty:
            inferred_bar_size = float(non_null_bar_sizes.iloc[-1])
    if inferred_bar_size is None and "label_bar_size_secs" in dataset.columns:
        non_null_label_bar_sizes = dataset["label_bar_size_secs"].dropna()
        if not non_null_label_bar_sizes.empty:
            inferred_bar_size = float(non_null_label_bar_sizes.iloc[-1])
    if inferred_bar_size is None:
        inferred_bar_size = 60.0
    if "bar_size_secs" not in dataset.columns:
        dataset["bar_size_secs"] = inferred_bar_size
    else:
        dataset["bar_size_secs"] = dataset["bar_size_secs"].fillna(inferred_bar_size)
    inferred_horizon_steps = max(1.0, horizon_seconds / max(float(inferred_bar_size), 1.0))
    if "label_horizon_steps" not in dataset.columns:
        dataset["label_horizon_steps"] = inferred_horizon_steps
    else:
        dataset["label_horizon_steps"] = dataset["label_horizon_steps"].fillna(
            inferred_horizon_steps
        )
    required_cols = [
        "mid_price",
        "return_1",
        "return_5",
        "realized_vol",
        "book_imbalance",
        "microprice_tilt",
        "order_flow_imbalance",
        "cum_delta_10",
        "cum_delta_30",
        "cum_delta_50",
        "cross_venue_spread",
        "spread_rel",
        "depth_total",
        "volume_delta",
        "ofi_mean_12",
        "ofi_std_12",
        "imbalance_mean_12",
        "spread_rel_mean_12",
        "depth_change_1",
        "microprice_tilt_change_5",
        "microprice_gap",
        "forward_microprice_return_bps",
        "label_horizon_steps",
        "horizon_seconds",
    ]
    available_required = [column for column in required_cols if column in dataset.columns]
    dataset = dataset.dropna(subset=available_required)
    available_optional = [
        column for column in OPTIONAL_FEATURE_COLUMNS if column in dataset.columns
    ]
    for column in available_optional:
        dataset[column] = dataset[column].fillna(0.0)
    return dataset.sort_values("ts").reset_index(drop=True)


def _build_training_batch(
    dataset: pd.DataFrame,
    *,
    feature_columns: list[str] | None = None,
    sequence_columns: list[str] | None = None,
) -> SequenceBatch:
    tabular_columns = feature_columns or [
        column for column in TABULAR_FEATURE_COLUMNS if column in dataset.columns
    ]
    seq_columns = sequence_columns or [
        column for column in SEQUENCE_FEATURE_COLUMNS if column in dataset.columns
    ]
    if not tabular_columns or not seq_columns:
        raise ValueError("No trainable microstructure feature columns available")
    tabular_rows: list[np.ndarray] = []
    sequence_rows: list[np.ndarray] = []
    labels: list[int] = []
    forward_returns_bps: list[float] = []
    timestamps: list[pd.Timestamp] = []
    label_horizon_steps = int(dataset["label_horizon_steps"].iloc[-1])
    for end_idx in range(SEQUENCE_LENGTH - 1, len(dataset)):
        window = dataset.iloc[end_idx - SEQUENCE_LENGTH + 1 : end_idx + 1]
        tabular_rows.append(dataset.iloc[end_idx][tabular_columns].to_numpy(dtype=np.float32))
        sequence_rows.append(window[seq_columns].to_numpy(dtype=np.float32))
        labels.append(int(dataset.iloc[end_idx]["label"]))
        forward_returns_bps.append(float(dataset.iloc[end_idx]["forward_microprice_return_bps"]))
        timestamps.append(cast(pd.Timestamp, dataset.iloc[end_idx]["ts"]))
    label_series = dataset["label"].astype(int)
    move_series = dataset.get("forward_microprice_return_bps")
    mean_moves = {
        label: float(move_series[label_series == label].mean())
        if move_series is not None and not move_series[label_series == label].empty
        else 0.0
        for label in [0, 1, 2]
    }
    return SequenceBatch(
        tabular=np.asarray(tabular_rows, dtype=np.float32),
        sequence=np.asarray(sequence_rows, dtype=np.float32),
        labels=np.asarray(labels, dtype=np.int64),
        forward_returns_bps=np.asarray(forward_returns_bps, dtype=np.float32),
        timestamps=timestamps,
        feature_columns=tabular_columns,
        sequence_columns=seq_columns,
        mean_moves=mean_moves,
        label_horizon_steps=label_horizon_steps,
        feature_schema_version=FEATURE_SCHEMA_VERSION,
        label_schema_version=LABEL_SCHEMA_VERSION,
    )


def _validation_split(count: int, label_horizon_steps: int) -> tuple[np.ndarray, np.ndarray] | None:
    if count < MIN_TRAIN_ROWS:
        return None
    cv = PurgedTimeSeriesCV(
        n_splits=1,
        embargo_bars=max(SEQUENCE_LENGTH - 1, label_horizon_steps),
    )
    for train_idx, val_idx in cv.split(pd.DataFrame(index=np.arange(count))):
        return train_idx, val_idx
    return None


def _train_xgb(
    x_train: np.ndarray,
    y_train: np.ndarray,
    x_val: np.ndarray,
    y_val: np.ndarray,
) -> tuple[xgb.XGBClassifier, int]:
    model = xgb.XGBClassifier(
        objective="multi:softprob",
        num_class=3,
        n_estimators=300,
        early_stopping_rounds=50,
        max_depth=5,
        learning_rate=0.05,
        subsample=0.85,
        colsample_bytree=0.85,
        min_child_weight=5.0,
        reg_lambda=2.0,
        tree_method="hist",
        eval_metric="mlogloss",
        random_state=42,
    )
    model.fit(
        x_train,
        y_train,
        eval_set=[(x_val, y_val)],
        verbose=False,
    )
    best_iteration_raw = getattr(model, "best_iteration", None)
    default_iteration = int(getattr(model, "n_estimators", 300) or 300)
    best_iteration = (
        int(best_iteration_raw) if best_iteration_raw is not None else default_iteration
    )
    return model, best_iteration


def _train_tcn(
    x_train: np.ndarray,
    y_train: np.ndarray,
    x_val: np.ndarray,
    y_val: np.ndarray,
) -> TcnClassifier:
    device = torch.device("cpu")
    model = TcnClassifier(x_train.shape[2]).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)
    class_weights = torch.tensor(
        _class_weights(y_train),
        dtype=torch.float32,
        device=device,
    )
    loss_fn = FocalLoss(gamma=2.0, alpha=class_weights)
    train_x = torch.tensor(np.transpose(x_train, (0, 2, 1)), dtype=torch.float32, device=device)
    train_y = torch.tensor(y_train, dtype=torch.long, device=device)
    val_x = torch.tensor(np.transpose(x_val, (0, 2, 1)), dtype=torch.float32, device=device)
    val_y = torch.tensor(y_val, dtype=torch.long, device=device)
    train_loader = DataLoader(
        TensorDataset(train_x, train_y),
        batch_size=min(256, len(train_x)),
        shuffle=False,
    )
    best_state = model.state_dict()
    best_loss = float("inf")
    patience = 10
    stale_epochs = 0
    for _ in range(50):
        model.train()
        for batch_x, batch_y in train_loader:
            optimizer.zero_grad()
            logits, _ = model(batch_x)
            loss = loss_fn(logits, batch_y)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
        model.eval()
        with torch.no_grad():
            val_logits, _ = model(val_x)
            val_loss = float(loss_fn(val_logits, val_y))
        scheduler.step()
        if val_loss < best_loss:
            best_loss = val_loss
            best_state = {
                key: value.detach().clone() for key, value in model.state_dict().items()
            }
            stale_epochs = 0
        else:
            stale_epochs += 1
            if stale_epochs >= patience:
                break
    model.load_state_dict(best_state)
    model.eval()
    return model


def _predict_tcn(model: TcnClassifier, sequence: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    tensor = torch.tensor(np.transpose(sequence, (0, 2, 1)), dtype=torch.float32)
    with torch.no_grad():
        logits, embedding = model(tensor)
    probs = torch.softmax(logits, dim=1).cpu().numpy()
    return probs, embedding.cpu().numpy()


def _embed_tcn(model: TcnClassifier, sequence: np.ndarray) -> np.ndarray:
    _, embedding = _predict_tcn(model, sequence)
    return embedding


def _fit_calibrator(probs: np.ndarray, labels: np.ndarray) -> dict[str, dict[str, list[float]]]:
    calibration: dict[str, dict[str, list[float]]] = {}
    for klass in range(probs.shape[1]):
        iso = IsotonicRegression(out_of_bounds="clip")
        iso.fit(probs[:, klass], (labels == klass).astype(float))
        calibration[str(klass)] = {
            "x": iso.X_thresholds_.tolist(),
            "y": iso.y_thresholds_.tolist(),
        }
    return calibration


def _apply_calibrator(
    calibration: dict[str, dict[str, list[float]]],
    probs: np.ndarray,
) -> np.ndarray:
    calibrated = np.zeros_like(probs)
    for klass in range(probs.shape[1]):
        curve = calibration.get(str(klass))
        if not curve:
            calibrated[:, klass] = probs[:, klass]
            continue
        calibrated[:, klass] = np.interp(
            probs[:, klass],
            np.asarray(curve["x"], dtype=np.float64),
            np.asarray(curve["y"], dtype=np.float64),
        )
    row_sum = calibrated.sum(axis=1, keepdims=True)
    row_sum[row_sum == 0] = 1.0
    return calibrated / row_sum


def _predict_fusion(model: dict[str, Any], batch: SequenceBatch) -> np.ndarray:
    fusion_model = cast(xgb.XGBClassifier, model["fusion_model"])
    embeddings = _embed_tcn(cast(TcnClassifier, model["tcn_model"]), batch.sequence)
    features = np.concatenate([batch.tabular, embeddings], axis=1)
    raw = fusion_model.predict_proba(features)
    return _apply_calibrator(model["calibration"], raw)


def _load_exported_model(artifact: dict[str, Any]) -> dict[str, Any]:
    if artifact.get("model_type") != MODEL_VERSION:
        raise ValueError(f"Incompatible model artifact: {artifact.get('model_type')}")
    tabular = xgb.XGBClassifier()
    _load_xgb_model_bytes(tabular, base64.b64decode(artifact["tabular_model_b64"]))
    fusion = xgb.XGBClassifier()
    _load_xgb_model_bytes(fusion, base64.b64decode(artifact["fusion_model_b64"]))
    tcn = TcnClassifier(len(artifact["sequence_feature_columns"]))
    state_dict = torch.load(
        io.BytesIO(base64.b64decode(artifact["tcn_state_b64"])),
        map_location="cpu",
    )
    tcn.load_state_dict(state_dict)
    tcn.eval()
    return {
        "tabular_model": tabular,
        "fusion_model": fusion,
        "tcn_model": tcn,
        "calibration": artifact["calibration"],
        "tabular_feature_columns": artifact["tabular_feature_columns"],
        "sequence_feature_columns": artifact["sequence_feature_columns"],
        "class_move_bps": artifact["class_move_bps"],
        "profitable": artifact.get("profitable", False),
        "feature_schema_version": artifact.get("feature_schema_version", FEATURE_SCHEMA_VERSION),
        "label_schema_version": artifact.get("label_schema_version", LABEL_SCHEMA_VERSION),
        "validation_embargo_bars": artifact.get("validation_embargo_bars", 0),
    }


def _dump_xgb_model_bytes(model: xgb.XGBClassifier) -> bytes:
    booster = model.get_booster()
    try:
        return bytes(booster.save_raw(raw_format="json"))
    except TypeError:
        with tempfile.NamedTemporaryFile(suffix=".json") as handle:
            booster.save_model(handle.name)
            handle.flush()
            return Path(handle.name).read_bytes()


def _load_xgb_model_bytes(model: xgb.XGBClassifier, payload: bytes) -> None:
    try:
        model.load_model(bytearray(payload))
        return
    except TypeError:
        pass
    with tempfile.NamedTemporaryFile(suffix=".json") as handle:
        handle.write(payload)
        handle.flush()
        model.load_model(handle.name)


def _classification_metrics(
    labels: np.ndarray,
    probs: np.ndarray,
    realized_moves_bps: np.ndarray,
) -> dict[str, float]:
    preds = probs.argmax(axis=1)
    accuracy = float(np.mean(preds == labels))
    directions = np.where(preds == 2, 1.0, np.where(preds == 0, -1.0, 0.0))
    gross_returns = directions * realized_moves_bps
    costs = np.where(directions != 0.0, ROUND_TRIP_COST_BPS, 0.0)
    net_returns = gross_returns - costs
    paper_edge = float(np.mean(net_returns))
    paper_std = float(np.std(net_returns, ddof=0))
    paper_sharpe = 0.0 if paper_std == 0.0 else float(np.mean(net_returns) / paper_std)
    return {
        "mean_accuracy": accuracy,
        "mean_brier_score": _multiclass_brier(labels, probs),
        "paper_sharpe": paper_sharpe,
        "paper_net_edge_bps": paper_edge,
        "directional_f1": float(f1_score(labels, preds, average="macro")),
    }


def _multiclass_brier(y_true: np.ndarray, probs: np.ndarray) -> float:
    classes = np.arange(probs.shape[1])
    scores = [
        brier_score_loss((y_true == klass).astype(float), probs[:, klass])
        for klass in classes
    ]
    return float(np.mean(scores))


def _class_weights(labels: np.ndarray) -> np.ndarray:
    counts = np.bincount(labels, minlength=3).astype(np.float32)
    counts[counts == 0] = 1.0
    return counts.sum() / counts
