#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! Walk-forward supervised training binary for all non-RL models.
//!
//! Trains models using the `UnifiedTrainable` trait on real OHLCV data loaded
//! from Databento DBN files. Supports walk-forward windows, early stopping,
//! checkpoint saving, and z-score normalization.
//!
//! # Supported Models
//!
//! tft, mamba2, liquid, tggn, tlob, kan, xlstm, diffusion
//!
//! # Usage
//!
//! ```bash
//! SQLX_OFFLINE=true cargo run -p ml --example train_baseline_supervised --release -- \
//!     --model kan --epochs 50 --batch-size 128 \
//!     --data-dir test_data/futures-baseline \
//!     --output-dir ml/trained_models
//! ```

#![allow(unused_crate_dependencies, clippy::cognitive_complexity)]
#![deny(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::indexing_slicing
)]

use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Instant;

use anyhow::{Context, Result};
use clap::Parser;
use serde_json::Value;
use tracing::{error, info, warn};

use ml::features::extraction::extract_ml_features;
use ml::training::unified_trainer::UnifiedTrainable;
use ml::types::OHLCVBar;
use ml::walk_forward::{generate_walk_forward_windows, NormStats, WalkForwardConfig};

#[allow(unreachable_pub)]
mod baseline_common;
use baseline_common::completion::{write_failure_marker, write_success_marker, CompletionMetrics};
use baseline_common::{load_all_bars, spread_cost_bps};
use common::metrics::{server as metrics_server, training_metrics as metrics};

// Model adapter imports -- using verified paths from existing examples
use ml::diffusion::{DiffusionConfig, DiffusionTrainableAdapter};
use ml::kan::{KANConfig, KANTrainableAdapter};
use ml::liquid::{CfCTrainConfig, DeviceConfig, LiquidTrainableAdapter};
use ml::mamba::{Mamba2Config, trainable_adapter::Mamba2TrainableAdapter};
use ml::tft::{TFTConfig, TrainableTFT};
use ml::tgnn::trainable_adapter::TGGNTrainableAdapter;
use ml::tgnn::TGGNConfig;
use ml::tlob::{TLOBAdapterConfig, TLOBTrainableAdapter};
use ml::xlstm::{XLSTMConfig, XLSTMTrainableAdapter};

// ---------------------------------------------------------------------------
// CLI Arguments
// ---------------------------------------------------------------------------

/// Walk-forward supervised training binary for non-RL models.
#[derive(Parser, Debug)]
#[command(name = "train_baseline_supervised", about = "Train supervised models with walk-forward windows")]
struct Args {
    /// Model to train: tft, mamba2, liquid, tggn, tlob, kan, xlstm, diffusion, all
    #[arg(long)]
    model: String,

    /// Path to directory containing .dbn.zst files (with symbol subdirectories, env: FOXHUNT_DATA_DIR)
    #[arg(long, env = "FOXHUNT_DATA_DIR")]
    data_dir: PathBuf,

    /// Symbol subdirectory to load (e.g. "ES.FUT", "NQ.FUT")
    #[arg(long, default_value = "ES.FUT")]
    symbol: String,

    /// Maximum training epochs per fold
    #[arg(long, default_value_t = 50)]
    epochs: usize,

    /// Training batch size
    #[arg(long, default_value_t = 128)]
    batch_size: usize,

    /// Learning rate
    #[arg(long, default_value_t = 1e-3)]
    learning_rate: f64,

    /// Feature dimension (must match `extract_ml_features` output)
    #[arg(long, default_value_t = 51)]
    feature_dim: usize,

    /// Output directory for trained model checkpoints
    #[arg(long, default_value = "ml/trained_models")]
    output_dir: PathBuf,

    /// Early stopping patience (epochs without improvement)
    #[arg(long, default_value_t = 10)]
    patience: usize,

    /// Walk-forward: initial training window in months
    #[arg(long, default_value_t = 12)]
    train_months: u32,

    /// Walk-forward: validation window in months
    #[arg(long, default_value_t = 3)]
    val_months: u32,

    /// Walk-forward: test window in months
    #[arg(long, default_value_t = 3)]
    test_months: u32,

    /// Walk-forward: step size in months between folds
    #[arg(long, default_value_t = 3)]
    step_months: u32,

    /// Maximum absolute per-bar return; larger moves are clamped
    #[arg(long, default_value_t = 0.01)]
    max_bar_return: f64,

    /// Round-trip commission cost in basis points
    #[arg(long, default_value_t = 1.0)]
    tx_cost_bps: f64,

    /// Instrument tick size in price units (ES=0.25)
    #[arg(long, default_value_t = 0.25)]
    tick_size: f64,

    /// Typical bid-ask spread in ticks
    #[arg(long, default_value_t = 1.0)]
    spread_ticks: f64,

    /// Optional path to hyperopt results JSON -- overrides matching config fields
    #[arg(long)]
    hyperopt_params: Option<PathBuf>,

    /// Training profile name (e.g. "supervised-production", "supervised-dev")
    #[arg(long, default_value = "supervised-production")]
    training_profile: String,
}

// ---------------------------------------------------------------------------
// Supported models
// ---------------------------------------------------------------------------

const ALL_MODELS: &[&str] = &[
    "tft", "mamba2", "liquid", "tggn", "tlob", "kan", "xlstm", "diffusion",
];

fn validate_model(name: &str) -> Result<()> {
    if name == "all" || ALL_MODELS.contains(&name) {
        Ok(())
    } else {
        anyhow::bail!(
            "Unknown model '{}'. Valid: {} or 'all'",
            name,
            ALL_MODELS.join(", ")
        );
    }
}

// ---------------------------------------------------------------------------
// Hyperopt parameter loading
// ---------------------------------------------------------------------------

/// Load `best_params` from a hyperopt results JSON file.
///
/// Expected format: `{ "model_key": { "best_params": { ... }, ... } }`
/// Returns `None` if the file doesn't exist or can't be parsed.
fn load_hyperopt_params(hp_path: &Option<PathBuf>, model_key: &str) -> Option<Value> {
    let file_path = hp_path.as_ref()?;
    if !file_path.exists() {
        info!("Hyperopt params file not found: {}, using defaults", file_path.display());
        return None;
    }
    let contents = match std::fs::read_to_string(file_path) {
        Ok(c) => c,
        Err(e) => {
            warn!("Failed to read hyperopt params {}: {}", file_path.display(), e);
            return None;
        }
    };
    let json: Value = match serde_json::from_str(&contents) {
        Ok(v) => v,
        Err(e) => {
            warn!("Failed to parse hyperopt params JSON: {}", e);
            return None;
        }
    };
    let params = json.get(model_key)
        .and_then(|m| m.get("best_params"))
        .cloned();
    if params.is_some() {
        info!("Loaded hyperopt params for '{}' from {}", model_key, file_path.display());
    } else {
        warn!("No best_params found for '{}' in {}", model_key, file_path.display());
    }
    params
}

fn hp_f64(params: &Option<Value>, key: &str) -> Option<f64> {
    params.as_ref()?.get(key)?.as_f64()
}

fn hp_usize(params: &Option<Value>, key: &str) -> Option<usize> {
    params.as_ref()?.get(key)?.as_u64().map(|v| v as usize)
}

// ---------------------------------------------------------------------------
// Model Factory
// ---------------------------------------------------------------------------

/// Create a model adapter by name with production-default configs.
///
/// All OHLCV models receive `feature_dim` as input dimension.
/// The adapter's internal projection layers handle mapping to model-native dims.
fn create_model(
    name: &str,
    feature_dim: usize,
    learning_rate: f64,
    hp: &Option<Value>,
) -> Result<Box<dyn UnifiedTrainable>> {
    let lr = hp_f64(hp, "learning_rate").unwrap_or(learning_rate);
    match name {
        "tft" => {
            let config = TFTConfig {
                input_dim: feature_dim,
                hidden_dim: hp_usize(hp, "hidden_size").unwrap_or(128),
                num_heads: hp_usize(hp, "num_heads").unwrap_or(4),
                num_layers: 2,
                num_quantiles: 3,
                num_static_features: 0,
                num_known_features: 0,
                num_unknown_features: feature_dim,
                // Each sample is a single-timestep 51-dim feature vector with
                // a 1-step-ahead return target — no windowed sequences.
                sequence_length: 1,
                prediction_horizon: 1,
                dropout_rate: hp_f64(hp, "dropout").unwrap_or(0.1),
                ..TFTConfig::default()
            };
            let mut adapter = TrainableTFT::new(config)
                .map_err(|e| anyhow::anyhow!("Failed to create TFT: {}", e))?;
            adapter
                .set_learning_rate(lr)
                .map_err(|e| anyhow::anyhow!("Failed to set TFT learning rate: {}", e))?;
            Ok(Box::new(adapter))
        }
        "mamba2" => {
            let config = Mamba2Config {
                d_model: 128,
                num_layers: 4,
                d_state: 16,
                max_seq_len: 60,
                ..Mamba2Config::default()
            };
            let native_device = ml_core::native_types::NativeDevice::Cuda(0);
            let mut adapter = Mamba2TrainableAdapter::from_native_device(config, &native_device)
                .map_err(|e| anyhow::anyhow!("Failed to create Mamba2: {}", e))?;
            adapter
                .set_learning_rate(lr)
                .map_err(|e| anyhow::anyhow!("Failed to set Mamba2 learning rate: {}", e))?;
            Ok(Box::new(adapter))
        }
        "liquid" => {
            let config = CfCTrainConfig {
                input_size: feature_dim,
                hidden_size: hp_usize(hp, "hidden_size").unwrap_or(128),
                output_size: 1,
                backbone_hidden_sizes: vec![128, 64],
                learning_rate: lr,
                device: DeviceConfig::Auto,
                ..CfCTrainConfig::default()
            };
            let adapter = LiquidTrainableAdapter::new(config)
                .map_err(|e| anyhow::anyhow!("Failed to create Liquid: {}", e))?;
            Ok(Box::new(adapter))
        }
        "tggn" => {
            let config = TGGNConfig {
                node_dim: feature_dim,
                hidden_dim: hp_usize(hp, "hidden_dim").unwrap_or(32),
                num_layers: hp_usize(hp, "num_layers").unwrap_or(2),
                max_nodes: 64,
                max_edges: 128,
                edge_dim: 4,
                temporal_decay: hp_f64(hp, "temporal_decay").unwrap_or(0.99),
                update_frequency_ns: 1_000_000,
                use_simd: false,
            };
            let mut adapter = TGGNTrainableAdapter::new(config)
                .map_err(|e| anyhow::anyhow!("Failed to create TGGN: {}", e))?;
            adapter
                .set_learning_rate(lr)
                .map_err(|e| anyhow::anyhow!("Failed to set TGGN learning rate: {}", e))?;
            Ok(Box::new(adapter))
        }
        "tlob" => {
            let config = TLOBAdapterConfig {
                d_model: hp_usize(hp, "d_model").unwrap_or(128),
                num_heads: hp_usize(hp, "num_heads").unwrap_or(4),
                num_layers: hp_usize(hp, "num_layers").unwrap_or(2),
                seq_len: hp_usize(hp, "seq_len").unwrap_or(1),
                feature_dim,
            };
            let mut adapter = TLOBTrainableAdapter::new(config)
                .map_err(|e| anyhow::anyhow!("Failed to create TLOB: {}", e))?;
            adapter
                .set_learning_rate(lr)
                .map_err(|e| anyhow::anyhow!("Failed to set TLOB learning rate: {}", e))?;
            Ok(Box::new(adapter))
        }
        "kan" => {
            let config = KANConfig {
                grid_size: hp_usize(hp, "grid_size").unwrap_or(5),
                spline_order: hp_usize(hp, "spline_order").unwrap_or(4),
                layer_widths: vec![feature_dim, 32, 16, 1],
                learning_rate: lr,
                weight_decay: hp_f64(hp, "weight_decay").unwrap_or(1e-4),
                grad_clip: hp_f64(hp, "grad_clip").unwrap_or(1.0),
            };
            let adapter = KANTrainableAdapter::new(config)
                .map_err(|e| anyhow::anyhow!("Failed to create KAN: {}", e))?;
            Ok(Box::new(adapter))
        }
        "xlstm" => {
            let config = XLSTMConfig {
                input_dim: feature_dim,
                hidden_dim: hp_usize(hp, "hidden_dim").unwrap_or(128),
                ..XLSTMConfig::default()
            };
            let mut adapter = XLSTMTrainableAdapter::new(config)
                .map_err(|e| anyhow::anyhow!("Failed to create xLSTM: {}", e))?;
            adapter
                .set_learning_rate(lr)
                .map_err(|e| anyhow::anyhow!("Failed to set xLSTM learning rate: {}", e))?;
            Ok(Box::new(adapter))
        }
        "diffusion" => {
            let config = DiffusionConfig {
                feature_dim,
                hidden_dim: hp_usize(hp, "hidden_dim").unwrap_or(128),
                ..DiffusionConfig::default()
            };
            let ctx = cudarc::driver::CudaContext::new(0)
                .map_err(|e| anyhow::anyhow!("CUDA context for Diffusion: {}", e))?;
            let stream = ctx
                .new_stream()
                .map_err(|e| anyhow::anyhow!("CUDA stream for Diffusion: {}", e))?;
            let adapter = DiffusionTrainableAdapter::new(config, &stream)
                .map_err(|e| anyhow::anyhow!("Failed to create Diffusion: {}", e))?;
            Ok(Box::new(adapter))
        }
        _ => anyhow::bail!("Unknown model: {}", name),
    }
}

// ---------------------------------------------------------------------------
// Data preparation
// ---------------------------------------------------------------------------

/// Build (input, target) f32 vector pairs from OHLCV bars for regression training.
///
/// Features are extracted via `extract_ml_features` (51-dim OHLCV market features),
/// z-score normalized, and the target is the next-bar return in basis points, clipped.
#[allow(clippy::type_complexity)]
fn prepare_fold_data(
    train_bars: &[OHLCVBar],
    val_bars: &[OHLCVBar],
    args: &Args,
) -> Result<(Vec<(Vec<f32>, Vec<f32>)>, Vec<(Vec<f32>, Vec<f32>)>, NormStats)> {
    let train_features = extract_ml_features(train_bars)
        .context("Failed to extract training features")?;
    let val_features = extract_ml_features(val_bars)
        .context("Failed to extract validation features")?;

    if train_features.len() < 2 {
        anyhow::bail!(
            "Insufficient training features: {}",
            train_features.len()
        );
    }
    if val_features.len() < 2 {
        anyhow::bail!(
            "Insufficient validation features: {}",
            val_features.len()
        );
    }

    let norm_stats = NormStats::from_features(&train_features);
    let norm_train = norm_stats.normalize_batch(&train_features);
    let norm_val = norm_stats.normalize_batch(&val_features);

    let train_bar_offset = train_bars.len().saturating_sub(train_features.len());
    let val_bar_offset = val_bars.len().saturating_sub(val_features.len());

    let train_pairs = build_vector_pairs(&norm_train, train_bars, train_bar_offset, args)?;
    let val_pairs = build_vector_pairs(&norm_val, val_bars, val_bar_offset, args)?;

    Ok((train_pairs, val_pairs, norm_stats))
}

/// Convert normalized features + bars into (input, target) f32 vector pairs.
fn build_vector_pairs(
    norm_features: &[[f64; 42]],
    bars: &[OHLCVBar],
    bar_offset: usize,
    args: &Args,
) -> Result<Vec<(Vec<f32>, Vec<f32>)>> {
    let mut pairs = Vec::new();
    let n = norm_features.len();
    let step_limit = n.saturating_sub(1);

    for i in 0..step_limit {
        let Some(feat) = norm_features.get(i) else {
            continue;
        };

        let bar_idx = i + bar_offset;
        let close_cur = bars.get(bar_idx).map(|b| b.close).unwrap_or(0.0);
        let close_next = bars
            .get(bar_idx + 1)
            .map(|b| b.close)
            .unwrap_or(close_cur);

        let return_bps = if close_cur.abs() > 1e-10 {
            (close_next - close_cur) / close_cur * 10_000.0
        } else {
            0.0
        };
        let max_bps = args.max_bar_return * 10_000.0;
        let clipped = return_bps.clamp(-max_bps, max_bps);
        let spread = spread_cost_bps(close_cur, args.tick_size, args.spread_ticks);
        let net_return = clipped - (args.tx_cost_bps + spread);

        let input_f32: Vec<f32> = feat.iter().map(|&v| v as f32).collect();
        let target_f32 = vec![net_return as f32];

        pairs.push((input_f32, target_f32));
    }

    Ok(pairs)
}

// ---------------------------------------------------------------------------
// Generic training loop
// ---------------------------------------------------------------------------

/// Run one training epoch on mini-batches. Returns average loss.
///
/// UnifiedTrainable processes one sample at a time, so we iterate
/// over samples within each batch.
fn run_training_epoch(
    adapter: &mut dyn UnifiedTrainable,
    train_pairs: &[(Vec<f32>, Vec<f32>)],
    batch_size: usize,
    model_name: &str,
    fold_str: &str,
) -> Result<f64> {
    let mut epoch_loss_sum = 0.0_f64;
    let mut epoch_steps = 0_usize;
    let n_train = train_pairs.len();
    let mut batch_start = 0_usize;

    while batch_start < n_train {
        let batch_end = (batch_start + batch_size).min(n_train);
        let Some(batch_slice) = train_pairs.get(batch_start..batch_end) else {
            break;
        };

        if batch_slice.is_empty() {
            batch_start = batch_end;
            continue;
        }

        // UnifiedTrainable processes one sample at a time
        for (input, target) in batch_slice {
            adapter
                .zero_grad()
                .map_err(|e| anyhow::anyhow!("zero_grad failed: {}", e))?;
            let loss_val = adapter
                .forward_loss(input, target)
                .map_err(|e| anyhow::anyhow!("forward_loss failed: {}", e))?;

            if (loss_val as f32).is_nan() || (loss_val as f32).is_infinite() {
                metrics::record_nan_detected(model_name, fold_str);
            }

            // backward + optimizer_step: some models (TFT, xLSTM, Diffusion)
            // manage parameters internally and don't support backward via
            // UnifiedTrainable. For those, skip parameter updates — they use
            // their own native train() methods for actual training.
            if let Err(e) = adapter.backward(loss_val) {
                // Log once per epoch, not per sample
                if epoch_steps == 0 {
                    warn!("[{}] backward unsupported via UnifiedTrainable: {}", model_name, e);
                }
            } else if let Err(e) = adapter.optimizer_step() {
                if epoch_steps == 0 {
                    warn!("[{}] optimizer_step unsupported via UnifiedTrainable: {}", model_name, e);
                }
            }

            epoch_loss_sum += loss_val;
            epoch_steps += 1;
        }

        batch_start = batch_end;
    }

    metrics::set_batches_processed(model_name, fold_str, epoch_steps as f64);

    Ok(if epoch_steps > 0 {
        epoch_loss_sum / epoch_steps as f64
    } else {
        0.0
    })
}

/// Train a model on a single walk-forward fold. Returns best validation loss.
fn train_fold(
    model_name: &str,
    fold: usize,
    train_pairs: &[(Vec<f32>, Vec<f32>)],
    val_pairs: &[(Vec<f32>, Vec<f32>)],
    args: &Args,
    output_dir: &Path,
    hp: &Option<Value>,
) -> Result<f64> {
    info!(
        "[{}] Fold {} -- {} train, {} val pairs",
        model_name,
        fold,
        train_pairs.len(),
        val_pairs.len()
    );

    let mut adapter = create_model(model_name, args.feature_dim, args.learning_rate, hp)?;

    let mut best_val_loss = f64::MAX;
    let mut epochs_without_improvement = 0_usize;
    let fold_str = fold.to_string();

    for epoch in 0..args.epochs {
        let epoch_start = Instant::now();

        let avg_train_loss =
            run_training_epoch(adapter.as_mut(), train_pairs, args.batch_size, model_name, &fold_str)?;

        // Validation: iterate over val_pairs using forward_loss
        let mut val_loss_sum = 0.0_f64;
        let mut val_count = 0_usize;
        for (input, target) in val_pairs {
            match adapter.forward_loss(input, target) {
                Ok(loss) => {
                    val_loss_sum += loss;
                    val_count += 1;
                }
                Err(e) => {
                    warn!("Validation forward_loss failed on sample {}: {}", val_count, e);
                }
            }
        }
        let val_loss = if val_count > 0 {
            val_loss_sum / val_count as f64
        } else {
            0.0
        };
        let elapsed = epoch_start.elapsed();

        // Prometheus metrics
        metrics::set_epoch(model_name, &fold_str, (epoch + 1) as f64);
        metrics::set_epoch_loss(model_name, &fold_str, avg_train_loss);
        metrics::set_validation_loss(model_name, &fold_str, val_loss);
        metrics::set_iteration_seconds(model_name, &fold_str, elapsed.as_secs_f64());
        metrics::set_learning_rate(model_name, &fold_str, adapter.get_learning_rate());
        metrics::set_epoch_duration(model_name, &fold_str, elapsed.as_secs_f64());

        info!(
            "  Fold {} Epoch {}/{}: train_loss={:.6}, val_loss={:.6}, lr={:.2e}, time={:.1}s",
            fold,
            epoch + 1,
            args.epochs,
            avg_train_loss,
            val_loss,
            adapter.get_learning_rate(),
            elapsed.as_secs_f64(),
        );

        if val_loss < best_val_loss {
            best_val_loss = val_loss;
            epochs_without_improvement = 0;
            let ckpt_path = output_dir.join(format!("{}_fold{}_best", model_name, fold));
            let ckpt_str = ckpt_path.to_str().unwrap_or("checkpoint");
            let ckpt_start = Instant::now();
            if let Err(e) = adapter.save_checkpoint(ckpt_str) {
                metrics::record_checkpoint_failure(model_name, &fold_str);
                warn!("  Failed to save checkpoint: {}", e);
            } else {
                let ckpt_size = ckpt_path
                    .with_extension("safetensors")
                    .metadata()
                    .map(|m| m.len() as f64)
                    .unwrap_or(0.0);
                metrics::record_checkpoint_save(
                    model_name,
                    &fold_str,
                    ckpt_start.elapsed().as_secs_f64(),
                    ckpt_size,
                );
                info!(
                    "  [{}] New best val_loss={:.6}, checkpoint saved",
                    model_name, val_loss
                );
            }
        } else {
            epochs_without_improvement += 1;
        }

        if epochs_without_improvement >= args.patience {
            info!(
                "  [{}] Early stopping at epoch {} (patience {} exhausted)",
                model_name,
                epoch + 1,
                args.patience
            );
            break;
        }
    }

    // Final checkpoint
    let final_path = output_dir.join(format!("{}_fold{}_final", model_name, fold));
    let final_str = final_path.to_str().unwrap_or("final");
    if let Err(e) = adapter.save_checkpoint(final_str) {
        warn!("  Failed to save final checkpoint: {}", e);
    }

    Ok(best_val_loss)
}

// ---------------------------------------------------------------------------
// Training orchestration
// ---------------------------------------------------------------------------

/// Container for per-model results collected during training.
struct TrainingResult {
    model_name: String,
    fold_results: Vec<(usize, f64)>,
    total_epochs: usize,
}

/// Run the full walk-forward training pipeline.
///
/// Returns a vec of per-model results so `main()` can write completion markers.
fn run_training(args: &Args) -> Result<Vec<TrainingResult>> {
    validate_model(&args.model)?;

    info!("Training profile: {}", args.training_profile);

    let models_to_train: Vec<&str> = if args.model == "all" {
        ALL_MODELS.to_vec()
    } else {
        vec![args.model.as_str()]
    };

    // Load OHLCV bars
    info!(
        "Loading OHLCV bars for {} from {}",
        args.symbol,
        args.data_dir.display()
    );
    let data_load_start = Instant::now();
    let all_bars = load_all_bars(&args.data_dir, &args.symbol)?;
    info!("Loaded {} bars", all_bars.len());

    if all_bars.len() < 100 {
        anyhow::bail!(
            "Insufficient data: {} bars (need >= 100)",
            all_bars.len()
        );
    }

    // Walk-forward windows
    let wf_config = WalkForwardConfig {
        initial_train_months: args.train_months,
        val_months: args.val_months,
        test_months: args.test_months,
        step_months: args.step_months,
    };
    let windows = generate_walk_forward_windows(&all_bars, &wf_config);
    if windows.is_empty() {
        anyhow::bail!(
            "No walk-forward windows generated. Data too short for configured window sizes."
        );
    }
    info!("Generated {} walk-forward folds", windows.len());

    let mut all_results = Vec::new();

    // Record data loading time (shared across all models)
    for mn in &models_to_train {
        metrics::record_data_load(mn, data_load_start.elapsed().as_secs_f64());
    }

    // Train each model
    for model_name in &models_to_train {
        info!("=== Training model: {} ===", model_name);
        let hp = load_hyperopt_params(&args.hyperopt_params, model_name);
        let model_output = args.output_dir.join(model_name);
        std::fs::create_dir_all(&model_output)
            .with_context(|| format!("Failed to create {}", model_output.display()))?;

        let mut fold_results: Vec<(usize, f64)> = Vec::new();
        let mut total_epochs = 0_usize;

        for window in &windows {
            let fold = window.fold;
            info!(
                "--- Fold {}: train={}, val={}, test={} bars ---",
                fold,
                window.train.len(),
                window.val.len(),
                window.test.len()
            );

            let (train_pairs, val_pairs, norm_stats) =
                match prepare_fold_data(&window.train, &window.val, args) {
                    Ok(data) => data,
                    Err(e) => {
                        warn!("Skipping fold {} -- {}", fold, e);
                        continue;
                    }
                };

            // Persist NormStats for evaluation
            let norm_path = model_output.join(format!("norm_stats_fold{}.json", fold));
            match serde_json::to_string_pretty(&norm_stats) {
                Ok(json) => {
                    if let Err(e) = std::fs::write(&norm_path, json) {
                        warn!("Failed to save NormStats to {}: {}", norm_path.display(), e);
                    }
                }
                Err(e) => warn!("Failed to serialize NormStats: {}", e),
            }

            if train_pairs.is_empty() || val_pairs.is_empty() {
                warn!(
                    "Skipping fold {} -- empty data after feature extraction",
                    fold
                );
                continue;
            }

            match train_fold(
                model_name,
                fold,
                &train_pairs,
                &val_pairs,
                args,
                &model_output,
                &hp,
            ) {
                Ok(best_val) => {
                    total_epochs += args.epochs; // upper bound per fold
                    fold_results.push((fold, best_val));
                }
                Err(e) => error!("[{}] Fold {} failed: {}", model_name, fold, e),
            }
        }

        // Summary
        info!(
            "=== {} Results ({} folds) ===",
            model_name,
            fold_results.len()
        );
        for (fold, loss) in &fold_results {
            info!("  Fold {}: best_val_loss = {:.6}", fold, loss);
        }
        if !fold_results.is_empty() {
            let avg: f64 =
                fold_results.iter().map(|(_, l)| l).sum::<f64>() / fold_results.len() as f64;
            info!("  Average: {:.6}", avg);
        }
        info!("  Checkpoints: {}", model_output.display());

        all_results.push(TrainingResult {
            model_name: (*model_name).to_owned(),
            fold_results,
            total_epochs,
        });
    }

    Ok(all_results)
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

fn main() -> Result<()> {
    // Initialize tracing with optional OTLP export to Tempo
    let otlp_endpoint = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok();
    if let Err(e) = common::observability::init_observability(
        "train_baseline_supervised",
        otlp_endpoint.as_deref(),
    ) {
        eprintln!("Observability init failed (non-fatal): {e}");
    }

    metrics::init();
    metrics_server::start_metrics_server(9094);
    metrics::set_active_workers(1.0);

    let args = Args::parse();

    // Ensure output directory exists before training so markers can always be written.
    if let Err(e) = std::fs::create_dir_all(&args.output_dir) {
        error!("Failed to create output dir {}: {}", args.output_dir.display(), e);
    }

    let result = run_training(&args);
    metrics::set_active_workers(0.0);

    // Push final metrics to pushgateway so they persist after pod termination
    if let Err(e) = metrics_server::push_to_gateway(None, "train_baseline_supervised") {
        tracing::warn!("Failed to push metrics to gateway (non-fatal): {e}");
    }

    match result {
        Ok(results) => {
            for training_result in &results {
                let best_val = training_result
                    .fold_results
                    .iter()
                    .map(|(_, loss)| *loss)
                    .fold(f64::MAX, f64::min);

                let metrics = CompletionMetrics {
                    model: training_result.model_name.clone(),
                    symbol: args.symbol.clone(),
                    best_val_loss: (best_val < f64::MAX).then_some(best_val),
                    sharpe_ratio: None,
                    epochs_completed: training_result.total_epochs,
                    folds_completed: training_result.fold_results.len(),
                };
                write_success_marker(&args.output_dir, &metrics);
            }
            Ok(())
        }
        Err(e) => {
            let msg = format!("{:#}", e);
            error!("Training failed: {}", msg);
            write_failure_marker(&args.output_dir, &msg);
            Err(e)
        }
    }
}