fidc-backtest-engine/crates/fidc-core/src/universe.rs

use chrono::NaiveDate;
use serde::Serialize;

use crate::data::{BenchmarkSnapshot, DataSet, EligibleUniverseSnapshot};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BandRegime {
    Bullish,
    Neutral,
    Defensive,
}

#[derive(Debug, Clone, Serialize)]
pub struct UniverseCandidate {
    pub symbol: String,
    pub market_cap_bn: f64,
    pub free_float_cap_bn: f64,
    pub band_low: f64,
    pub band_high: f64,
}

#[derive(Debug, Clone, Serialize)]
pub struct SelectionDiagnostics {
    pub decision_date: NaiveDate,
    pub reference_level: f64,
    pub band_low: f64,
    pub band_high: f64,
    pub factor_total: usize,
    pub market_cap_missing_count: usize,
    pub candidate_missing_count: usize,
    pub market_missing_count: usize,
    pub not_eligible_count: usize,
    pub paused_count: usize,
    pub out_of_band_count: usize,
    pub selected_before_limit: usize,
    pub selected_after_limit: usize,
    pub missing_market_cap_symbols: Vec<String>,
    pub selected_symbols: Vec<String>,
    pub rejection_examples: Vec<String>,
}

pub struct SelectionContext<'a> {
    pub decision_date: NaiveDate,
    pub benchmark: &'a BenchmarkSnapshot,
    pub reference_level: f64,
    pub data: &'a DataSet,
}

pub trait UniverseSelector {
    fn select(&self, ctx: &SelectionContext<'_>) -> Vec<UniverseCandidate>;
    fn select_with_diagnostics(
        &self,
        ctx: &SelectionContext<'_>,
    ) -> (Vec<UniverseCandidate>, SelectionDiagnostics);
}

#[derive(Debug, Clone)]
pub struct DynamicMarketCapBandSelector {
    pub base_index_level: f64,
    pub base_cap_floor: f64,
    pub cap_span: f64,
    pub xs: f64,
    pub top_n: usize,
}

impl DynamicMarketCapBandSelector {
    pub fn new(
        base_index_level: f64,
        base_cap_floor: f64,
        cap_span: f64,
        xs: f64,
        top_n: usize,
    ) -> Self {
        Self {
            base_index_level,
            base_cap_floor,
            cap_span,
            xs,
            top_n,
        }
    }

    pub fn demo(top_n: usize) -> Self {
        Self::new(2000.0, 7.0, 10.0, 4.0 / 500.0, top_n)
    }

    pub fn regime(&self, benchmark_level: f64) -> BandRegime {
        if benchmark_level >= self.base_index_level + 400.0 {
            BandRegime::Bullish
        } else if benchmark_level >= self.base_index_level {
            BandRegime::Neutral
        } else {
            BandRegime::Defensive
        }
    }

    pub fn band_for_level(&self, benchmark_level: f64) -> (f64, f64) {
        let start = ((benchmark_level - self.base_index_level) * self.xs) + self.base_cap_floor;
        let low = start.round();
        (low, low + self.cap_span)
    }
}

impl UniverseSelector for DynamicMarketCapBandSelector {
    fn select(&self, ctx: &SelectionContext<'_>) -> Vec<UniverseCandidate> {
        self.select_with_diagnostics(ctx).0
    }

    fn select_with_diagnostics(
        &self,
        ctx: &SelectionContext<'_>,
    ) -> (Vec<UniverseCandidate>, SelectionDiagnostics) {
        let _regime = self.regime(ctx.reference_level);
        let (min_cap, max_cap) = self.band_for_level(ctx.reference_level);
        let mut diagnostics = SelectionDiagnostics {
            decision_date: ctx.decision_date,
            reference_level: ctx.reference_level,
            band_low: min_cap,
            band_high: max_cap,
            factor_total: 0,
            market_cap_missing_count: 0,
            candidate_missing_count: 0,
            market_missing_count: 0,
            not_eligible_count: 0,
            paused_count: 0,
            out_of_band_count: 0,
            selected_before_limit: 0,
            selected_after_limit: 0,
            missing_market_cap_symbols: Vec::new(),
            selected_symbols: Vec::new(),
            rejection_examples: Vec::new(),
        };

        diagnostics.factor_total = ctx.data.factor_snapshots_on(ctx.decision_date).len();
        diagnostics.market_cap_missing_count = diagnostics
            .factor_total
            .saturating_sub(ctx.data.eligible_universe_on(ctx.decision_date).len());

        let eligible = ctx.data.eligible_universe_on(ctx.decision_date);
        let start_idx = lower_bound_by_market_cap(eligible, min_cap);
        let mut selected = Vec::new();

        for factor in eligible.iter().skip(start_idx) {
            if factor.market_cap_bn > max_cap {
                break;
            }
            selected.push(to_universe_candidate(factor, min_cap, max_cap));
        }

        diagnostics.out_of_band_count = eligible.len().saturating_sub(selected.len());

        diagnostics.selected_before_limit = selected.len();
        if selected.len() > self.top_n {
            selected.truncate(self.top_n);
        }
        diagnostics.selected_symbols = selected.iter().map(|item| item.symbol.clone()).collect();
        diagnostics.selected_after_limit = diagnostics.selected_symbols.len();
        (selected, diagnostics)
    }
}

fn lower_bound_by_market_cap(rows: &[EligibleUniverseSnapshot], target: f64) -> usize {
    let mut left = 0usize;
    let mut right = rows.len();
    while left < right {
        let mid = left + (right - left) / 2;
        if rows[mid].market_cap_bn < target {
            left = mid + 1;
        } else {
            right = mid;
        }
    }
    left
}

fn to_universe_candidate(
    factor: &EligibleUniverseSnapshot,
    band_low: f64,
    band_high: f64,
) -> UniverseCandidate {
    UniverseCandidate {
        symbol: factor.symbol.clone(),
        market_cap_bn: factor.market_cap_bn,
        free_float_cap_bn: factor.free_float_cap_bn,
        band_low,
        band_high,
    }
}