foxhunt/crates/ml-alpha/examples/alpha_mamba_baseline.rs

//! Phase 1d.1 — Stateful Mamba2 sequence encoder smoke.
//!
//! Trains the from-scratch GPU-pure Mamba2 block on the snapshot-stream
//! fxcache. Compares val AUC against the MLP baseline (Phase 1c snapshot:
//! AUC=0.685 at K=100 stateless).
//!
//! Falsification gate (per the implementation plan):
//!   - AUC > 0.72 → sequence-state lifts beyond stateless MLP. Proceed
//!     to Phase 1d.2 (multi-minute horizon).
//!   - AUC ≤ 0.72 → stateless features have a hard ceiling; richer
//!     features (not richer model) are the next move.
//!
//! Sequence semantics: each training example is a contiguous block of
//! `seq_len` snapshots ending at bar `t`; the model predicts the binary
//! direction label at bar `t` (sign of `mid_price[t+H] - mid_price[t]`,
//! computed by `prepare_phase1a_data`).
//!
//! Memory model: feature matrix + labels stay CPU-resident (640 MB
//! f32 + 1.5 MB labels for the 1.97 M-row snapshot fxcache). Per batch,
//! we gather contiguous seq_len-bar slices into a pinned host buffer,
//! DMA to GPU once, then do forward + backward + AdamW step.

use anyhow::{Context, Result};
use clap::Parser;
use cudarc::driver::CudaContext;
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;
use std::sync::Arc;
use tracing_subscriber::EnvFilter;

use ml_alpha::eval::{accuracy_from_logits, auc_from_logits};
use ml_alpha::mamba2_block::{
    Mamba2AdamW, Mamba2AdamWConfig, Mamba2Block, Mamba2BlockConfig,
};
use ml_alpha::fxcache_reader::FxCacheReader;
use ml_alpha::purged_split::PurgedSplit;
use ml_alpha::training::{prepare_phase1a_data, Phase1aConfig};

use ml_core::cuda_autograd::gpu_tensor::GpuTensor;

#[derive(Parser, Debug)]
#[command(name = "alpha_mamba_baseline", about = "FoxhuntQ-Δ Phase 1d.1 Mamba2 sequence smoke")]
struct Cli {
    #[arg(long)] fxcache_path: String,
    #[arg(long, default_value_t = 100)] horizon: usize,
    #[arg(long, default_value_t = 32)] seq_len: usize,
    #[arg(long, default_value_t = 64)] hidden_dim: usize,
    #[arg(long, default_value_t = 16)] state_dim: usize,
    #[arg(long, default_value_t = 5)] epochs: usize,
    #[arg(long, default_value_t = 256)] batch_size: usize,
    #[arg(long, default_value_t = 1e-3)] lr: f32,
    #[arg(long, default_value_t = 42)] seed: u64,
    /// Subsample stride for train sequences (1 = use every starting position;
    /// larger = subsample). Memory + time both scale linearly with this.
    #[arg(long, default_value_t = 4)] train_stride: usize,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")))
        .init();
    let cli = Cli::parse();
    let ctx = CudaContext::new(0).context("init CUDA")?;
    let stream = ctx.default_stream();

    // ── Load fxcache, build features+labels (mirrors phase1a) ──────────
    let mut cfg_p1a = Phase1aConfig::default();
    cfg_p1a.fxcache_path = cli.fxcache_path.clone();
    cfg_p1a.horizon = cli.horizon;
    let reader = FxCacheReader::open(&cfg_p1a.fxcache_path)?;
    let alpha_dim = reader.alpha_feature_dim()
        .ok_or_else(|| anyhow::anyhow!("fxcache lacks alpha column"))?;
    cfg_p1a.mlp.in_dim = alpha_dim;
    let split = PurgedSplit::new(reader.bar_count(), cfg_p1a.train_frac,
        cfg_p1a.horizon, cfg_p1a.embargo_bars)?.split();
    let data = prepare_phase1a_data(&reader, &split, &cfg_p1a)?;
    println!(
        "fxcache: {} rows, alpha_dim={}, n_train={}, n_val={}",
        reader.bar_count(), alpha_dim, data.train_indices.len(), data.val_indices.len()
    );

    // ── Build sequence start indices ──────────────────────────────────
    // A sequence ending at bar t requires t-seq_len+1 ≥ 0 and t < bar_count.
    // We filter train/val to bars where the preceding seq_len-1 bars are
    // within the same train (or val) range. For simplicity we use bars
    // where `idx ≥ seq_len - 1`.
    let seq_len = cli.seq_len;
    let make_seq_starts = |indices: &[usize], stride: usize| -> Vec<usize> {
        indices.iter()
            .filter(|&&i| i + 1 >= seq_len)
            .step_by(stride)
            .map(|&i| i + 1 - seq_len)  // start bar of the seq_len-bar window
            .collect()
    };
    let train_starts = make_seq_starts(&data.train_indices, cli.train_stride);
    let val_starts = make_seq_starts(&data.val_indices, 1);
    println!("train sequences (stride={}): {}", cli.train_stride, train_starts.len());
    println!("val sequences:               {}", val_starts.len());

    // ── Build Mamba2 + AdamW ──────────────────────────────────────────
    let block_cfg = Mamba2BlockConfig {
        in_dim: alpha_dim,
        hidden_dim: cli.hidden_dim,
        state_dim: cli.state_dim,
        seq_len,
    };
    let mut block = Mamba2Block::new(block_cfg.clone(), Arc::clone(&stream))?;
    let mut opt = Mamba2AdamW::new(&block, Mamba2AdamWConfig {
        lr: cli.lr,
        ..Default::default()
    })?;
    println!("Mamba2Block params: {}", block.param_count());

    // ── Training loop ─────────────────────────────────────────────────
    let mut rng = ChaCha8Rng::seed_from_u64(cli.seed);
    let train_len = train_starts.len();
    if train_len < cli.batch_size {
        anyhow::bail!("n_train_sequences ({}) < batch_size ({})", train_len, cli.batch_size);
    }
    let batches_per_epoch = train_len / cli.batch_size;

    let bce_loss = |logits: &[f32], labels: &[f32]| -> f32 {
        let mut s = 0.0_f32;
        let eps = 1e-7_f32;
        for (&z, &y) in logits.iter().zip(labels.iter()) {
            let z = z.clamp(-50.0, 50.0);
            let p = (1.0 / (1.0 + (-z).exp())).clamp(eps, 1.0 - eps);
            s += -(y * p.ln() + (1.0 - y) * (1.0 - p).ln());
        }
        s / logits.len() as f32
    };

    for epoch in 0..cli.epochs {
        let mut perm: Vec<usize> = (0..train_len).collect();
        for i in (1..train_len).rev() {
            let j = rng.gen_range(0..=i);
            perm.swap(i, j);
        }
        let mut loss_sum = 0.0_f32;
        for batch_idx in 0..batches_per_epoch {
            let sel = &perm[batch_idx * cli.batch_size .. (batch_idx + 1) * cli.batch_size];
            // Pull labels from train_labels via end-bar lookup. We rebuilt gather_batch
            // to read train_labels for sequence END bars (sequence labels = last-position).
            let (input, labels) = gather_batch_for_train(
                &train_starts, sel, seq_len, alpha_dim,
                &data.feature_matrix, &data.train_labels, &data.train_indices, &stream
            )?;
            let (logit, cache) = block.forward_train(&input)?;
            let logit_host = logit.to_host(&stream)?;
            let loss = bce_loss(&logit_host, &labels);
            loss_sum += loss;

            // d_logit = (sigmoid(z) - y) / N for mean-BCE-with-logits.
            let n_b = labels.len() as f32;
            let d_logit_host: Vec<f32> = logit_host.iter().zip(labels.iter())
                .map(|(&z, &y)| {
                    let p = 1.0 / (1.0 + (-z.clamp(-50.0, 50.0)).exp());
                    (p - y) / n_b
                })
                .collect();
            let d_logit_dev = stream.clone_htod(&d_logit_host)?;
            let d_logit = GpuTensor::new(d_logit_dev, vec![labels.len(), 1])?;
            let grads = block.backward(&cache, &d_logit)?;
            opt.step(&mut block, &grads)?;
        }
        let mean_loss = loss_sum / batches_per_epoch as f32;
        println!("epoch {epoch:2}  mean_train_bce={mean_loss:.5}  n_batches={batches_per_epoch}");
    }

    // ── Validation pass ───────────────────────────────────────────────
    let val_chunk = cli.batch_size;
    let mut val_logits: Vec<f32> = Vec::with_capacity(val_starts.len());
    let mut val_labels_y: Vec<f32> = Vec::with_capacity(val_starts.len());
    let mut i = 0;
    while i < val_starts.len() {
        let this = val_chunk.min(val_starts.len() - i);
        let sel: Vec<usize> = (i..i + this).collect();
        let (input, labels) = gather_batch_for_val(
            &val_starts, &sel, seq_len, alpha_dim,
            &data.feature_matrix, &data.val_labels, &data.val_indices, &stream
        )?;
        let logit = block.forward(&input)?;
        let logit_host = logit.to_host(&stream)?;
        val_logits.extend_from_slice(&logit_host);
        val_labels_y.extend_from_slice(&labels);
        i += this;
    }
    let labels_u8: Vec<u8> = val_labels_y.iter().map(|&y| if y > 0.5 { 1 } else { 0 }).collect();
    let acc = accuracy_from_logits(&val_logits, &labels_u8);
    let auc = auc_from_logits(&val_logits, &labels_u8);
    println!();
    println!("===========================================================");
    println!("PHASE 1d.1 — MAMBA2 SMOKE");
    println!("===========================================================");
    println!("val sequences: {}", val_logits.len());
    println!("val accuracy:  {:.4}", acc);
    println!("val AUC:       {:.4}", auc);
    println!("MLP baseline (Phase 1c K=100): AUC=0.6849");
    if auc > 0.72 {
        println!("GATE PASS: sequence-state lift; proceed to Phase 1d.2 (multi-minute horizon).");
    } else if auc > 0.6849 {
        println!("MARGINAL: above MLP baseline ({:.4} > 0.6849) but below gate (0.72). Inspect.", auc);
    } else {
        println!("GATE FAIL: AUC ≤ MLP baseline 0.6849; no sequence-state lift.");
    }
    Ok(())
}

/// Gather a training batch. Sequence labels come from `train_labels`, which
/// is indexed *positionally* (the i-th train_label corresponds to the i-th
/// train_indices entry). We map start → end_bar → position in train_indices.
fn gather_batch_for_train(
    starts: &[usize],
    sel: &[usize],
    seq_len: usize,
    alpha_dim: usize,
    feature_matrix: &[f32],
    train_labels: &[f32],
    train_indices: &[usize],
    stream: &Arc<cudarc::driver::CudaStream>,
) -> Result<(GpuTensor, Vec<f32>)> {
    let b = sel.len();
    let mut host = Vec::with_capacity(b * seq_len * alpha_dim);
    let mut labels = Vec::with_capacity(b);
    for &k in sel {
        let s = starts[k];
        let end_bar = s + seq_len - 1;
        for t in 0..seq_len {
            let off = (s + t) * alpha_dim;
            host.extend_from_slice(&feature_matrix[off..off + alpha_dim]);
        }
        // Find end_bar in train_indices via binary search (train_indices is
        // strictly ascending in the prepare_phase1a_data implementation).
        let pos = train_indices.binary_search(&end_bar).ok();
        labels.push(pos.map(|p| train_labels[p]).unwrap_or(0.5));
    }
    let dev = stream.clone_htod(&host)?;
    let t = GpuTensor::new(dev, vec![b, seq_len, alpha_dim])
        .map_err(|e| anyhow::anyhow!("batch tensor: {e}"))?;
    Ok((t, labels))
}

fn gather_batch_for_val(
    starts: &[usize],
    sel: &[usize],
    seq_len: usize,
    alpha_dim: usize,
    feature_matrix: &[f32],
    val_labels: &[f32],
    val_indices: &[usize],
    stream: &Arc<cudarc::driver::CudaStream>,
) -> Result<(GpuTensor, Vec<f32>)> {
    let b = sel.len();
    let mut host = Vec::with_capacity(b * seq_len * alpha_dim);
    let mut labels = Vec::with_capacity(b);
    for &k in sel {
        let s = starts[k];
        let end_bar = s + seq_len - 1;
        for t in 0..seq_len {
            let off = (s + t) * alpha_dim;
            host.extend_from_slice(&feature_matrix[off..off + alpha_dim]);
        }
        let pos = val_indices.binary_search(&end_bar).ok();
        labels.push(pos.map(|p| val_labels[p]).unwrap_or(0.5));
    }
    let dev = stream.clone_htod(&host)?;
    let t = GpuTensor::new(dev, vec![b, seq_len, alpha_dim])
        .map_err(|e| anyhow::anyhow!("batch tensor: {e}"))?;
    Ok((t, labels))
}