foxhunt/crates/ml/examples/alpha_fit_fill_model.rs

//! Phase E.0 Task 5 — fit FillModel coefficients from historical MBP-10 + trade tape.
//!
//! ## What this does
//!
//! Streams MBP-10 snapshots concurrently with the time-sorted trade tape. For
//! each snapshot, determines whether a hypothetical posted limit at the L1
//! bid (resp. ask) would have been filled within `--window-seconds` (default
//! 60s). Accumulates (FillFeatures, fill_outcome) pairs per side, then calls
//! `fit_poisson` from `ml::env::fill_model` (cloglog Bernoulli likelihood, see
//! [`pearl_glm_fitter_link_must_match_inference`]). Writes the 6 fitted
//! coefficient sets to JSON.
//!
//! ## L1-only limitation (2026-05-15)
//!
//! `DbnParser::parse_mbp10_streaming` currently tracks only the L1 best
//! bid/ask via `update_level(0, ...)` and **ignores** the `levels:
//! [BidAskPair; 10]` array carried by every `Mbp10Msg`. The full top-10 is
//! in the message but is not copied. Until that parser is fixed to copy
//! the full top-10, this binary fits ONLY the L1 distribution and
//! replicates L1 coefficients across the L2/L3 slots with an attenuated
//! intercept (`β_0 -= ln(L+1)`) to keep monotonically lower fill rates at
//! deeper levels.
//!
//! ## Run
//!
//! ```bash
//! SQLX_OFFLINE=true cargo run --release --example alpha_fit_fill_model -- \
//!   --mbp10-dir  /home/jgrusewski/Work/foxhunt/test_data/futures-baseline-mbp10/ES.FUT \
//!   --trades-dir /home/jgrusewski/Work/foxhunt/test_data/futures-baseline-trades/ES.FUT \
//!   --window-seconds 60 \
//!   --out-path alpha_fill_coeffs.json
//! ```

use std::fs::File;
use std::io::Write;
use std::path::PathBuf;

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

use data::providers::databento::{
    dbn_parser::DbnParser,
    mbp10::Mbp10Snapshot,
};

/// Convert raw `Mbp10Msg.price` (i64 fixed-point at 1e9 scale per the
/// databento DBN format) to f32 in price units. Bypasses
/// `BidAskPair::price_to_f64`, which divides by 1e12 (different convention)
/// and would give bid/ask values 1000× too small. The parser stores raw
/// `mbp10.price` directly into `bid_px`/`ask_px` without re-scaling, so
/// readers must apply the 1e-9 themselves.
#[inline]
fn raw_price_to_f32(fixed: i64) -> f32 {
    (fixed as f64 * 1e-9) as f32
}
use ml::env::fill_model::{fit_poisson_l2, FillCoeffs, FillFeatures};
use ml::features::{filter_front_month, load_trades_sync, DbnTrade};
use ml::trainers::dqn::collect_dbn_files_recursive;

#[derive(Debug, Parser)]
#[command(
    name = "alpha_fit_fill_model",
    about = "Phase E.0 Task 5 — fit FillModel coefficients from historical MBP-10 + trades"
)]
struct Cli {
    /// Directory containing MBP-10 `.dbn[.zst]` files.
    #[arg(long)]
    mbp10_dir: PathBuf,
    /// Directory containing trade `.dbn[.zst]` files.
    #[arg(long)]
    trades_dir: PathBuf,
    /// Forward-looking fill window in seconds.
    #[arg(long, default_value_t = 60.0)]
    window_seconds: f32,
    /// Output JSON path.
    #[arg(long, default_value = "config/ml/alpha_fill_coeffs.json")]
    out_path: PathBuf,
    /// MBP-10 events between snapshot emissions (denser = larger sample).
    #[arg(long, default_value_t = 10)]
    snapshot_interval: usize,
    /// SGD iterations for `fit_poisson`.
    #[arg(long, default_value_t = 1000)]
    fit_iters: usize,
    /// SGD learning rate for `fit_poisson`.
    #[arg(long, default_value_t = 1e-2)]
    fit_lr: f32,
    /// L2 regularization strength on slope coefficients (intercept is
    /// NOT penalized). Default 0.01 — small but non-zero, prevents the
    /// β_spread runaway (-40 in the unregularized canonical run) while
    /// staying close to the data-fit minimum. Use 0.0 for the
    /// (deprecated) unregularized fit.
    #[arg(long, default_value_t = 0.01)]
    l2_lambda: f32,
    /// Optional cap on snapshots processed (for quick smokes).
    #[arg(long)]
    max_snapshots: Option<usize>,
}

fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .init();
    let cli = Cli::parse();
    info!("Phase E.0 Task 5 — FillModel calibration starting");
    info!("  mbp10_dir  = {}", cli.mbp10_dir.display());
    info!("  trades_dir = {}", cli.trades_dir.display());
    info!("  window_seconds   = {:.1}", cli.window_seconds);
    info!("  snapshot_interval = {}", cli.snapshot_interval);
    info!("  fit_iters = {}, fit_lr = {}, l2_lambda = {}",
          cli.fit_iters, cli.fit_lr, cli.l2_lambda);
    if let Some(cap) = cli.max_snapshots {
        info!("  max_snapshots cap = {}", cap);
    }

    // --- 1. Load trades, front-month-filtered per file, time-sorted ---
    let trade_files = collect_dbn_files_recursive(&cli.trades_dir);
    if trade_files.is_empty() {
        anyhow::bail!("no .dbn[.zst] files found in {:?}", cli.trades_dir);
    }
    info!("Found {} trade file(s)", trade_files.len());
    let mut trades: Vec<DbnTrade> = Vec::new();
    for file in &trade_files {
        let raw = load_trades_sync(file)
            .with_context(|| format!("load_trades_sync({})", file.display()))?;
        let filtered = filter_front_month(&raw);
        info!(
            "  {} -> {} raw, {} front-month",
            file.file_name().unwrap_or_default().to_string_lossy(),
            raw.len(),
            filtered.len()
        );
        trades.extend(filtered);
    }
    trades.sort_by_key(|t| t.timestamp);
    info!("Total trades loaded (front-month, time-sorted): {}", trades.len());
    if trades.is_empty() {
        anyhow::bail!("trades vec empty after filtering — check --trades-dir");
    }

    // --- 2. Stream snapshots, accumulate per-(side, L1) observations ---
    let window_ns: u64 = (cli.window_seconds * 1_000_000_000.0) as u64;
    let parser = DbnParser::new().context("DbnParser::new")?;
    let mbp10_files = collect_dbn_files_recursive(&cli.mbp10_dir);
    if mbp10_files.is_empty() {
        anyhow::bail!("no MBP-10 .dbn[.zst] files found in {:?}", cli.mbp10_dir);
    }
    info!("Found {} MBP-10 file(s)", mbp10_files.len());

    let mut bid_feat: Vec<FillFeatures> = Vec::new();
    let mut bid_outcomes: Vec<bool> = Vec::new();
    let mut ask_feat: Vec<FillFeatures> = Vec::new();
    let mut ask_outcomes: Vec<bool> = Vec::new();
    let mut n_processed: usize = 0;
    let mut n_degenerate_skipped: usize = 0;
    let limit = cli.max_snapshots.unwrap_or(usize::MAX);

    'files: for file in &mbp10_files {
        let mut hit_limit = false;
        let result = parser.parse_mbp10_streaming(
            file,
            cli.snapshot_interval,
            |snap: &Mbp10Snapshot| {
                if n_processed >= limit {
                    hit_limit = true;
                    return;
                }
                if snap.levels.is_empty() {
                    return;
                }
                let l1 = &snap.levels[0];
                let bid_l1 = raw_price_to_f32(l1.bid_px);
                let ask_l1 = raw_price_to_f32(l1.ask_px);
                if bid_l1 <= 0.0 || ask_l1 <= 0.0 || bid_l1 >= ask_l1 {
                    // Skip pre-quote / one-sided / crossed snapshots — the parser
                    // updates one side per message so the early file region can
                    // legitimately have a stale opposite side.
                    n_degenerate_skipped += 1;
                    return;
                }
                let mid = 0.5 * (bid_l1 + ask_l1);
                let spread = ask_l1 - bid_l1;
                let spread_bps = 10_000.0 * spread / mid;
                let bid_sz = l1.bid_sz as f32;
                let ask_sz = l1.ask_sz as f32;
                let l1_imbalance = if bid_sz + ask_sz > 0.0 {
                    bid_sz / (bid_sz + ask_sz)
                } else {
                    0.5
                };
                // OFI proxy uses only L1 because L2-L9 aren't populated by the
                // current streaming parser. Documented limitation; OK because
                // we're fitting L1 distributions only.
                let ofi_sum_5 = bid_sz - ask_sz;
                let ts = snap.timestamp;
                let window_end = ts.saturating_add(window_ns);
                let i_end = trades.partition_point(|t| t.timestamp < window_end);
                let i_start = trades.partition_point(|t| t.timestamp <= ts);
                let window_trades: &[DbnTrade] = &trades[i_start..i_end];

                // time_since_trade: prior trade strictly before snapshot
                let tau_s = if i_start > 0 {
                    let prior_ts = trades[i_start - 1].timestamp;
                    if ts >= prior_ts {
                        ((ts - prior_ts) as f32) * 1e-9
                    } else {
                        0.0
                    }
                } else {
                    0.0
                };

                let feat = FillFeatures {
                    spread_bps,
                    l1_imbalance,
                    ofi_sum_5,
                    time_since_trade_s: tau_s,
                };

                // L1 bid fill: any sell-aggressor (is_buy=false) at price ≤ bid_l1
                //   ⇒ a posted-bid limit would have been hit by an aggressive seller
                // L1 ask fill: any buy-aggressor (is_buy=true) at price ≥ ask_l1
                let bid_fill = window_trades
                    .iter()
                    .any(|t| !t.is_buy && (t.price as f32) <= bid_l1);
                let ask_fill = window_trades
                    .iter()
                    .any(|t| t.is_buy && (t.price as f32) >= ask_l1);

                bid_feat.push(feat);
                bid_outcomes.push(bid_fill);
                ask_feat.push(feat);
                ask_outcomes.push(ask_fill);

                n_processed += 1;
                if n_processed % 100_000 == 0 {
                    info!("  processed {} snapshots", n_processed);
                }
            },
        );
        match result {
            Ok(snap_count) => info!(
                "  {} → {} streamed snapshots from parser",
                file.file_name().unwrap_or_default().to_string_lossy(),
                snap_count
            ),
            Err(e) => warn!("parser error on {}: {}", file.display(), e),
        }
        if hit_limit {
            info!("max_snapshots cap reached; stopping iteration");
            break 'files;
        }
    }
    info!(
        "Snapshots accumulated: {}, degenerate-skipped: {}",
        n_processed, n_degenerate_skipped
    );
    if n_processed == 0 {
        anyhow::bail!("No usable snapshots accumulated — check data paths and parser output");
    }
    let bid_rate = bid_outcomes.iter().filter(|b| **b).count() as f64
        / bid_outcomes.len().max(1) as f64;
    let ask_rate = ask_outcomes.iter().filter(|b| **b).count() as f64
        / ask_outcomes.len().max(1) as f64;
    info!(
        "Empirical fill rates within {:.0}s window: bid_l1 = {:.4}, ask_l1 = {:.4}",
        cli.window_seconds, bid_rate, ask_rate
    );

    // --- 3. Fit cloglog Bernoulli (L1 only) ---
    info!(
        "Fitting bid L1 (n={}, iters={}, lr={})...",
        bid_feat.len(),
        cli.fit_iters,
        cli.fit_lr
    );
    let bid_l1_coeffs = fit_poisson_l2(&bid_feat, &bid_outcomes, cli.fit_iters, cli.fit_lr, cli.l2_lambda)
        .map_err(|e| anyhow::anyhow!("fit bid L1: {}", e))?;
    info!(
        "Fitting ask L1 (n={}, iters={}, lr={})...",
        ask_feat.len(),
        cli.fit_iters,
        cli.fit_lr
    );
    let ask_l1_coeffs = fit_poisson_l2(&ask_feat, &ask_outcomes, cli.fit_iters, cli.fit_lr, cli.l2_lambda)
        .map_err(|e| anyhow::anyhow!("fit ask L1: {}", e))?;

    // --- 4. Synthesize L2/L3 via β_0 attenuation ---
    //
    // λ_L = λ_L1 / (L+1) corresponds to β_0 -= ln(L+1):
    //   L1 (slot 0): β_0 unchanged
    //   L2 (slot 1): β_0 -= ln(2) ≈ -0.693
    //   L3 (slot 2): β_0 -= ln(3) ≈ -1.099
    //
    // This gives a monotonically lower expected fill rate at deeper levels,
    // which matches qualitative reality even though the L2/L3 SLOPES are
    // L1-equivalent (a documented coarse approximation).
    let make_levels = |c: &FillCoeffs| -> [FillCoeffs; 3] {
        let mut levels = [*c, *c, *c];
        levels[1].beta[0] -= 2.0_f32.ln();
        levels[2].beta[0] -= 3.0_f32.ln();
        levels
    };
    let bid_levels = make_levels(&bid_l1_coeffs);
    let ask_levels = make_levels(&ask_l1_coeffs);

    // --- 5. Write JSON ---
    let json = serde_json::json!({
        "phase": "E.0 Task 5",
        "n_snapshots": n_processed,
        "n_trades": trades.len(),
        "n_degenerate_skipped": n_degenerate_skipped,
        "window_seconds": cli.window_seconds,
        "snapshot_interval": cli.snapshot_interval,
        "fit_iters": cli.fit_iters,
        "fit_lr": cli.fit_lr,
        "l2_lambda": cli.l2_lambda,
        "empirical_bid_l1_rate": bid_rate,
        "empirical_ask_l1_rate": ask_rate,
        "l1_only_limitation": true,
        "l1_only_limitation_reason":
            "parse_mbp10_streaming ignores Mbp10Msg.levels[1..10]; \
             L2/L3 coefficients are L1 with attenuated β_0 (-ln(L+1))",
        "bid_coeffs": [
            bid_levels[0].beta.to_vec(),
            bid_levels[1].beta.to_vec(),
            bid_levels[2].beta.to_vec(),
        ],
        "ask_coeffs": [
            ask_levels[0].beta.to_vec(),
            ask_levels[1].beta.to_vec(),
            ask_levels[2].beta.to_vec(),
        ],
    });
    let mut f = File::create(&cli.out_path).context("create output file")?;
    write!(f, "{}", serde_json::to_string_pretty(&json)?)?;
    info!("Wrote coefficients to {}", cli.out_path.display());

    // --- 6. Sanity-check signs + report ---
    info!("=== Fitted L1 coefficient summary ===");
    info!(
        "BID L1: β_0={:+.4} β_spread={:+.4} β_imbal={:+.4} β_ofi={:+.4} β_logτ={:+.4}",
        bid_l1_coeffs.beta[0],
        bid_l1_coeffs.beta[1],
        bid_l1_coeffs.beta[2],
        bid_l1_coeffs.beta[3],
        bid_l1_coeffs.beta[4]
    );
    info!(
        "ASK L1: β_0={:+.4} β_spread={:+.4} β_imbal={:+.4} β_ofi={:+.4} β_logτ={:+.4}",
        ask_l1_coeffs.beta[0],
        ask_l1_coeffs.beta[1],
        ask_l1_coeffs.beta[2],
        ask_l1_coeffs.beta[3],
        ask_l1_coeffs.beta[4]
    );
    // Wider spread should reduce fill probability — β_spread should be negative
    // (or at worst slightly positive if the spread-fill correlation is masked
    // by another regime variable).
    if bid_l1_coeffs.beta[1] > 0.0 {
        warn!(
            "bid β_spread is positive ({:.4}) — counter-intuitive; expected negative",
            bid_l1_coeffs.beta[1]
        );
    }
    if ask_l1_coeffs.beta[1] > 0.0 {
        warn!(
            "ask β_spread is positive ({:.4}) — counter-intuitive",
            ask_l1_coeffs.beta[1]
        );
    }

    Ok(())
}