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

//! Phase 1d.2 — Multi-minute label smoke (K=6000 snapshots ≈ 1-5 min forward).
//!
//! THE decisive gate for the two-head architecture. The K-sweep (Phase 1c,
//! commit `db874b184`) showed alpha at stateless single-snapshot resolution
//! decays from K=50 peak to gone by K=500. Mamba2's job is to amplify
//! short-horizon evidence into a long-horizon prediction via SSM state
//! accumulation across `seq_len` snapshots.
//!
//! This test: feed a trailing window of 32 snapshots into Mamba2; predict
//! the binary direction `sign(mid[t+6000] - mid[t])`. Mamba's state must
//! integrate the per-step microstructure signal into a multi-minute call.
//!
//! Gate (per implementation plan):
//!   - AUC > 0.55 at K=6000 → multi-minute alpha confirmed; design works.
//!   - AUC < 0.52 → DECISIVE FAIL. The two-head architecture cannot turn
//!     short-window snapshot context into multi-minute prediction with
//!     the current model; need different inputs, model class, or both.
//!
//! Why this is decisive: if we *can* predict long horizons from short
//! tick context plus sequence state, the existing 81-dim snapshot stack
//! is the right foundation and the rest of FoxhuntQ-Δ proceeds. If we
//! can't, the whole architecture is wrong and needs a redesign before
//! any production work.

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::fxcache_reader::{FxCacheReader, COL_RAW_CLOSE, FEAT_DIM};
use ml_alpha::mamba2_block::{
    Mamba2AdamW, Mamba2AdamWConfig, Mamba2Block, Mamba2BlockConfig,
};
use ml_alpha::multi_horizon_labels::generate_labels;

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

#[derive(Parser, Debug)]
#[command(name = "phase1d_long_horizon", about = "FoxhuntQ-Δ Phase 1d.2 multi-minute Mamba smoke")]
struct Cli {
    #[arg(long)] fxcache_path: String,
    /// Label horizon in snapshots. K=6000 ≈ 1-5 min forward at typical rates.
    #[arg(long, default_value_t = 6000)] horizon: usize,
    /// Sequence length per training example (≤ kernel cap of 32).
    #[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 = 128)] batch_size: usize,
    #[arg(long, default_value_t = 3e-3)] lr: f32,
    /// 80/20 train/val split (purged via the embargo).
    #[arg(long, default_value_t = 0.8)] train_frac: f32,
    /// Embargo bars between train and val ranges (default = horizon for safety).
    #[arg(long)] embargo: Option<usize>,
    /// Subsample stride for train sequences (1 = use every starting position).
    #[arg(long, default_value_t = 4)] train_stride: usize,
    #[arg(long, default_value_t = 42)] seed: u64,
}

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 + extract feature matrix and prices ──────────────
    let reader = FxCacheReader::open(&cli.fxcache_path)?;
    let alpha_dim = reader.alpha_feature_dim()
        .ok_or_else(|| anyhow::anyhow!("fxcache lacks alpha column"))?;
    let n_bars = reader.bar_count();
    println!("fxcache: {} rows, alpha_dim={}, horizon K={}", n_bars, alpha_dim, cli.horizon);

    // Flat feature matrix (CPU-resident) and prices.
    let mut feature_matrix: Vec<f32> = Vec::with_capacity(n_bars * alpha_dim);
    let mut prices: Vec<f32> = Vec::with_capacity(n_bars);
    for i in 0..n_bars {
        let row = reader.alpha_features(i).context("alpha row missing")?;
        feature_matrix.extend_from_slice(row);
        let rec = reader.record(i);
        prices.push(rec.targets[COL_RAW_CLOSE - FEAT_DIM]);
    }

    // ── Generate K=6000 labels ────────────────────────────────────────
    let labels = generate_labels(&prices, cli.horizon);
    println!(
        "labels: kept={}  dropped_edge={}  dropped_invalid={}  up_frac={:.4}",
        labels.labels.len(), labels.n_dropped_edge, labels.n_dropped_invalid,
        labels.labels.iter().sum::<f32>() / labels.labels.len().max(1) as f32
    );
    if labels.labels.is_empty() {
        anyhow::bail!("no valid labels at K={} — fxcache too small or all-tied", cli.horizon);
    }

    // ── Purged train/val split over labels.valid_indices ──────────────
    // We split the VALID label list (not the raw bar list) at train_frac,
    // then apply an embargo so train sequences ending near the boundary
    // don't share forward-window prices with val.
    let embargo = cli.embargo.unwrap_or(cli.horizon);
    let n_valid = labels.valid_indices.len();
    let n_train_target = (n_valid as f32 * cli.train_frac) as usize;
    let train_split_bar = labels.valid_indices[n_train_target - 1];
    let val_start_bar = train_split_bar + embargo;
    let train_label_range: Vec<usize> = (0..n_valid)
        .filter(|&p| labels.valid_indices[p] < train_split_bar.saturating_sub(embargo))
        .collect();
    let val_label_range: Vec<usize> = (0..n_valid)
        .filter(|&p| labels.valid_indices[p] >= val_start_bar)
        .collect();
    println!(
        "split: n_train_labels={} n_val_labels={} (embargo={} bars, val_start_bar={})",
        train_label_range.len(), val_label_range.len(), embargo, val_start_bar
    );

    // Sequence start: bar `valid_indices[p] - seq_len + 1` to bar `valid_indices[p]`.
    // Filter out any p where the sequence would underflow.
    let make_seq_label_pos = |label_positions: &[usize], stride: usize| -> Vec<usize> {
        label_positions.iter()
            .filter(|&&p| labels.valid_indices[p] + 1 >= cli.seq_len)
            .copied()
            .step_by(stride)
            .collect()
    };
    let train_pos = make_seq_label_pos(&train_label_range, cli.train_stride);
    let val_pos = make_seq_label_pos(&val_label_range, 1);
    println!("train sequences (stride={}): {}  val sequences: {}",
             cli.train_stride, train_pos.len(), val_pos.len());

    // ── Build Mamba2 + AdamW ──────────────────────────────────────────
    let block_cfg = Mamba2BlockConfig {
        in_dim: alpha_dim,
        hidden_dim: cli.hidden_dim,
        state_dim: cli.state_dim,
        seq_len: cli.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());

    let gather = |sel: &[usize]| -> Result<(GpuTensor, Vec<f32>)> {
        let b = sel.len();
        let mut host = Vec::with_capacity(b * cli.seq_len * alpha_dim);
        let mut ys = Vec::with_capacity(b);
        for &p in sel {
            let end_bar = labels.valid_indices[p];
            let start_bar = end_bar + 1 - cli.seq_len;
            for t in 0..cli.seq_len {
                let off = (start_bar + t) * alpha_dim;
                host.extend_from_slice(&feature_matrix[off..off + alpha_dim]);
            }
            ys.push(labels.labels[p]);
        }
        let dev = stream.clone_htod(&host)?;
        let tensor = GpuTensor::new(dev, vec![b, cli.seq_len, alpha_dim])
            .map_err(|e| anyhow::anyhow!("batch tensor: {e}"))?;
        Ok((tensor, ys))
    };

    let bce_loss = |logits: &[f32], ys: &[f32]| -> f32 {
        let mut s = 0.0_f32;
        let eps = 1e-7_f32;
        for (&z, &y) in logits.iter().zip(ys.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
    };

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

    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 raw_sel = &perm[batch_idx * cli.batch_size .. (batch_idx + 1) * cli.batch_size];
            let sel: Vec<usize> = raw_sel.iter().map(|&i| train_pos[i]).collect();
            let (input, ys) = gather(&sel)?;
            let (logit, cache) = block.forward_train(&input)?;
            let logit_host = logit.to_host(&stream)?;
            loss_sum += bce_loss(&logit_host, &ys);

            let n_b = ys.len() as f32;
            let d_logit_host: Vec<f32> = logit_host.iter().zip(ys.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![ys.len(), 1])?;
            let grads = block.backward(&cache, &d_logit)?;
            opt.step(&mut block, &grads)?;
        }
        println!("epoch {epoch:2}  mean_train_bce={:.5}  n_batches={batches_per_epoch}",
                 loss_sum / batches_per_epoch as f32);
    }

    // ── Validation ────────────────────────────────────────────────────
    let mut val_logits: Vec<f32> = Vec::with_capacity(val_pos.len());
    let mut val_ys: Vec<f32> = Vec::with_capacity(val_pos.len());
    let mut i = 0;
    while i < val_pos.len() {
        let this = cli.batch_size.min(val_pos.len() - i);
        let sel: Vec<usize> = val_pos[i..i + this].to_vec();
        let (input, ys) = gather(&sel)?;
        let logit = block.forward(&input)?;
        let logit_host = logit.to_host(&stream)?;
        val_logits.extend_from_slice(&logit_host);
        val_ys.extend_from_slice(&ys);
        i += this;
    }
    let labels_u8: Vec<u8> = val_ys.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.2 — MAMBA2 MULTI-MINUTE SMOKE (K={})", cli.horizon);
    println!("===========================================================");
    println!("val sequences: {}", val_logits.len());
    println!("val accuracy:  {:.4}", acc);
    println!("val AUC:       {:.4}", auc);
    if auc > 0.55 {
        println!("GATE PASS: AUC > 0.55 at K={}; multi-minute alpha confirmed.", cli.horizon);
        println!("           Two-head FoxhuntQ-Δ architecture validated.");
    } else if auc < 0.52 {
        println!("GATE FAIL (decisive): AUC < 0.52; design dead in current form.");
        println!("           Need different inputs / model class / context window.");
    } else {
        println!("GATE MARGINAL: 0.52 ≤ AUC ≤ 0.55; tune hyperparameters or extend seq_len.");
    }
    Ok(())
}