foxhunt/crates/ml-backtesting/src/harness.rs

//! BacktestHarness — orchestrator wiring ml-alpha's MultiHorizonLoader
//! and CfcTrunk into the LobSimCuda. See spec §7.
//!
//! Design: the harness reuses the trainer's loader + captured graph
//! verbatim. Train-vs-deploy skew is structurally impossible because the
//! input struct (`Mbp10RawInput`), feature-assembly cubin
//! (`snap_feature_assemble`), and graph capture (`capture_graph_a` →
//! `perception_forward_captured`) are the same code paths used in
//! training. The harness's only job is to walk the chronological
//! snapshot stream and drive the sim's decision loop.

use anyhow::{Context, Result};
use ml_alpha::cfc::trunk::CfcTrunk;
use ml_alpha::data::loader::{
    discover_mbp10_files_sorted, MultiHorizonLoader, MultiHorizonLoaderConfig,
};
use ml_core::device::MlDevice;
use std::path::PathBuf;

use crate::sim::LobSimCuda;

#[derive(Clone, Debug)]
pub struct BacktestHarnessConfig {
    pub data_root: PathBuf,
    pub predecoded_dir: PathBuf,
    pub n_parallel: usize,
    pub decision_stride: usize,
    pub target_annual_vol_units: f32,
    pub annualisation_factor: f32,
    pub max_lots: u16,
    /// Max number of inference inputs to consume. 0 = exhaust loader.
    pub max_events: u64,
    /// Per-backtest Strategy compositions. Empty = every cell uses the
    /// hardcoded default policy (WeightedByRealizedSharpe across all 5
    /// horizons). Non-empty: len MUST equal n_parallel; each entry is
    /// flattened to bytecode and uploaded to its backtest slot.
    pub strategies: Vec<crate::policy::Strategy>,
    /// Total submission→fill latency in nanoseconds. 0 = immediate
    /// (legacy path). 100_000_000 = 100ms IBKR + Scaleway baseline.
    /// When non-zero, the orchestrator submits aggressive IOC limits
    /// with `active=2`/`arrival_ts_ns = current + latency_ns` so the
    /// book can move under them during the in-flight window.
    pub latency_ns: u32,
}

pub struct BacktestHarness {
    cfg: BacktestHarnessConfig,
    loader: MultiHorizonLoader,
    trunk: CfcTrunk,
    sim: LobSimCuda,
    decision_count: u64,
    event_count: u64,
    /// Per-cell cumulative P&L curve in USD, sampled once per event.
    /// `pnl_curves[b][i]` = realized_pnl USD for backtest b after event i.
    pnl_curves: Vec<Vec<f32>>,
}

impl BacktestHarness {
    /// Construct with an externally-built MlDevice + trunk. v1 ships
    /// with random-init trunk; loading from a checkpoint file is a
    /// follow-up (no checkpoint format pinned in ml-alpha yet).
    pub fn new(
        cfg: BacktestHarnessConfig,
        dev: &MlDevice,
        mut trunk: CfcTrunk,
    ) -> Result<Self> {
        if !cfg.strategies.is_empty() {
            anyhow::ensure!(
                cfg.strategies.len() == cfg.n_parallel,
                "strategies len {} ≠ n_parallel {}",
                cfg.strategies.len(),
                cfg.n_parallel
            );
        }
        let files = discover_mbp10_files_sorted(&cfg.data_root)?;
        let loader_cfg = MultiHorizonLoaderConfig {
            files,
            predecoded_dir: cfg.predecoded_dir.clone(),
            seq_len: 1,
            horizons: [30, 100, 300, 1000, 6000],
            n_max_sequences: 0,
            seed: 0,
            decision_stride: cfg.decision_stride,
            inference_only: true,
        };
        let loader = MultiHorizonLoader::new(&loader_cfg)?;

        // Capture the trunk's perception graph once using the first snapshot
        // as a template — required before any perception_forward_captured call.
        let first = loader.peek_first().context("peek_first")?;
        trunk.capture_graph_a(&first).context("capture_graph_a")?;

        let mut sim = LobSimCuda::new(cfg.n_parallel, dev)?;

        // Upload per-cell bytecode programs if the caller provided any.
        // Empty Vec means every cell uses the hardcoded default policy.
        for (b, strat) in cfg.strategies.iter().enumerate() {
            let prog = strat.flatten();
            sim.upload_program(b, &prog)?;
        }

        let pnl_curves = (0..cfg.n_parallel).map(|_| Vec::with_capacity(1024)).collect();
        Ok(Self {
            cfg,
            loader,
            trunk,
            sim,
            decision_count: 0,
            event_count: 0,
            pnl_curves,
        })
    }

    /// Drive the chronological event stream → trunk inference → sim
    /// decision loop until either the stream is exhausted or
    /// cfg.max_events is reached. See spec §7 orchestrator code.
    pub fn run(&mut self) -> Result<RunStats> {
        let stride = self.cfg.decision_stride.max(1) as u64;
        let mut total_decisions = 0u64;

        while let Some(raw) = self.loader.next_inference_input()? {
            // Apply snapshot to every backtest's book.
            self.sim.apply_snapshot(&raw.bid_px, &raw.bid_sz, &raw.ask_px, &raw.ask_sz)?;

            // Process resting orders (in-flight promotion, queue decay,
            // marketability check, stop triggers, OCO). Trade-flow signal
            // is currently unavailable from Mbp10RawInput (deferred to a
            // trades-feed integration); pass 0.0 so only the price-cross
            // marketability branch fires.
            self.sim.step_resting_orders(raw.ts_ns, 0.0)?;

            // At decision-stride boundaries: run trunk inference + sim decision.
            if self.event_count % stride == 0 {
                self.trunk.update_input_buffers(&raw)?;
                let (probs, _proj) = self.trunk.perception_forward_captured()?;
                self.sim.broadcast_alpha(&probs)?;
                self.sim.step_decision_with_latency(
                    raw.ts_ns,
                    self.cfg.target_annual_vol_units,
                    self.cfg.annualisation_factor,
                    self.cfg.max_lots,
                    self.cfg.latency_ns,
                )?;
                self.decision_count += 1;
                total_decisions += 1;
            }

            // Sample per-cell realised P&L (price-units → USD via $50/index-pt).
            for b in 0..self.cfg.n_parallel {
                let pos = self.sim.read_pos(b)?;
                self.pnl_curves[b].push(pos.realized_pnl * 50.0);
            }

            self.event_count += 1;
            if self.cfg.max_events > 0 && self.event_count >= self.cfg.max_events {
                break;
            }
        }

        Ok(RunStats {
            events_processed: self.event_count,
            decisions_taken: total_decisions,
        })
    }

    /// After `run()` completes, write per-cell artifacts to
    /// `<out_dir>/cell_<b>/{summary.json,trades.csv,pnl_curve.bin}`.
    pub fn write_artifacts(&self, out_dir: &std::path::Path) -> Result<()> {
        std::fs::create_dir_all(out_dir)
            .with_context(|| format!("create out dir {}", out_dir.display()))?;
        for b in 0..self.cfg.n_parallel {
            let cell_dir = out_dir.join(format!("cell_{b:04}"));
            std::fs::create_dir_all(&cell_dir)
                .with_context(|| format!("create cell dir {}", cell_dir.display()))?;
            let records = self.sim.read_trade_records(b)?;
            let curve = &self.pnl_curves[b];
            let summary = crate::artifacts::compute_summary(&records, curve);
            crate::artifacts::write_summary(&cell_dir.join("summary.json"), &summary)?;
            crate::artifacts::write_trades_csv(&cell_dir.join("trades.csv"), &records)?;
            crate::artifacts::write_pnl_curve_bin(&cell_dir.join("pnl_curve.bin"), curve)?;
        }
        Ok(())
    }

    pub fn sim(&self) -> &LobSimCuda { &self.sim }
    pub fn event_count(&self) -> u64 { self.event_count }
    pub fn decision_count(&self) -> u64 { self.decision_count }
}

#[derive(Clone, Copy, Debug)]
pub struct RunStats {
    pub events_processed: u64,
    pub decisions_taken: u64,
}