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

//! Durable intent progress, deliberately separate from actual-fill holding
//! protection. A published target starts no holding/protection timer.
use std::collections::{BTreeMap, BTreeSet};

use chrono::NaiveDate;
use rust_decimal::Decimal;
use serde::{Deserialize, Serialize};

use crate::stock_pool_execution::{
    Position, StockPoolMemberSpec, StockPoolPlan, normalize_stock_symbol,
};

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct StockPoolEntryProgress {
    pub pending: bool,
    pub observed_holding: bool,
    pub first_decision_date: NaiveDate,
    pub latest_generation: String,
    pub latest_target_value: Decimal,
    /// Fully funded entry goal, fixed at the last plan. Reconcile against
    /// actual holdings before repricing, never against today's market value.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub completion_quantity: Option<Decimal>,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct StockPoolExecutionState {
    pub schema_version: u32,
    pub last_execution_date: Option<NaiveDate>,
    pub entries: BTreeMap<String, StockPoolEntryProgress>,
    #[serde(default)]
    pub last_target_weights: BTreeMap<String, i32>,
    /// First signal excluding an actually held member; not an acquisition date.
    pub removed_since: BTreeMap<String, NaiveDate>,
}

pub struct StockPoolGoalObservation<'a> {
    pub symbol: &'a str,
    pub target_weight_bps: i32,
    pub target_value: Decimal,
    pub current_quantity: Decimal,
    pub target_quantity: Decimal,
    pub status: &'a str,
}

impl Default for StockPoolExecutionState {
    fn default() -> Self {
        Self {
            schema_version: 1,
            last_execution_date: None,
            entries: BTreeMap::new(),
            last_target_weights: BTreeMap::new(),
            removed_since: BTreeMap::new(),
        }
    }
}

impl StockPoolExecutionState {
    pub fn validate(&self) -> Result<(), String> {
        if self.schema_version != 1
            || self.entries.len() > 10000
            || self.removed_since.len() > 10000
        {
            return Err("stock_pool_execution_state_invalid_schema_or_size".into());
        }
        for symbol in self
            .entries
            .keys()
            .chain(self.removed_since.keys())
            .chain(self.last_target_weights.keys())
        {
            if normalize_stock_symbol(symbol).as_ref() != Some(symbol) {
                return Err("stock_pool_execution_state_invalid_symbol".into());
            }
        }
        if self.last_target_weights.len() > 10000
            || self
                .last_target_weights
                .values()
                .any(|value| !(0..=10000).contains(value))
        {
            return Err("stock_pool_execution_state_invalid_weights".into());
        }
        if self.entries.values().any(|entry| {
            entry.latest_target_value < Decimal::ZERO
                || entry.completion_quantity.is_some_and(|quantity| quantity <= Decimal::ZERO)
                || entry.latest_generation.is_empty()
                || self
                    .last_execution_date
                    .is_none_or(|last| entry.first_decision_date > last)
        }) || self
            .removed_since
            .values()
            .any(|day| self.last_execution_date.is_none_or(|last| *day > last))
        {
            return Err("stock_pool_execution_state_invalid_goal_or_clock".into());
        }
        Ok(())
    }

    pub fn observe(
        &self,
        decision_date: NaiveDate,
        execution_date: NaiveDate,
        official_dates: &[NaiveDate],
        members: &[StockPoolMemberSpec],
        positions: &[Position],
    ) -> Result<Self, String> {
        self.validate()?;
        if decision_date > execution_date
            || !official_dates.contains(&execution_date)
            || !official_dates.contains(&decision_date)
            || official_dates.windows(2).any(|pair| pair[0] >= pair[1])
            || self
                .last_execution_date
                .is_some_and(|last| last > execution_date)
        {
            return Err("stock_pool_execution_state_requires_monotone_official_clock".into());
        }
        let mut next = self.clone();
        next.last_execution_date = Some(execution_date);
        let members = members
            .iter()
            .map(|member| member.symbol.clone())
            .collect::<BTreeSet<_>>();
        let held = positions
            .iter()
            .filter(|position| position.quantity > Decimal::ZERO)
            .map(|position| position.symbol.clone())
            .collect::<BTreeSet<_>>();
        next.entries.retain(|symbol, entry| {
            // Confirmed flat starts a new cycle. A still-unfilled fresh target
            // may remain pending while the latest pool still requests it.
            !(entry.observed_holding && !held.contains(symbol))
                && (members.contains(symbol) || held.contains(symbol))
        });
        next.last_target_weights
            .retain(|symbol, _| members.contains(symbol) || held.contains(symbol));
        for (symbol, entry) in &mut next.entries {
            entry.observed_holding |= held.contains(symbol);
            if entry.pending
                && entry.completion_quantity.is_some_and(|goal| {
                    positions.iter().any(|position| {
                        &position.symbol == symbol && position.quantity >= goal
                    })
                })
            {
                entry.pending = false;
            }
        }
        next.removed_since
            .retain(|symbol, _| held.contains(symbol) && !members.contains(symbol));
        for symbol in held.difference(&members) {
            next.removed_since
                .entry(symbol.clone())
                .or_insert(decision_date);
        }
        next.validate()?;
        Ok(next)
    }

    pub fn pending_symbols(&self) -> BTreeSet<String> {
        self.entries
            .iter()
            .filter(|(_, entry)| entry.pending)
            .map(|(symbol, _)| symbol.clone())
            .collect()
    }

    pub fn next_day_exit_symbols(&self, execution_date: NaiveDate) -> BTreeSet<String> {
        self.removed_since
            .iter()
            .filter(|(_, removed)| **removed < execution_date)
            .map(|(symbol, _)| symbol.clone())
            .collect()
    }

    pub fn record_plan(
        &self,
        decision_date: NaiveDate,
        generation: &str,
        plan: &StockPoolPlan,
    ) -> Result<Self, String> {
        self.record_targets(
            decision_date,
            generation,
            plan.rows.iter().map(|row| StockPoolGoalObservation {
                symbol: &row.symbol,
                target_weight_bps: row.target_weight_bps,
                target_value: row.target_value,
                current_quantity: row.current_quantity,
                target_quantity: row.target_quantity,
                status: &row.status,
            }),
        )
    }

    pub fn record_targets<'a>(
        &self,
        decision_date: NaiveDate,
        generation: &str,
        rows: impl IntoIterator<Item = StockPoolGoalObservation<'a>>,
    ) -> Result<Self, String> {
        self.validate()?;
        if generation.is_empty()
            || self
                .last_execution_date
                .is_none_or(|date| decision_date > date)
        {
            return Err("stock_pool_execution_state_plan_clock_invalid".into());
        }
        let mut next = self.clone();
        for row in rows {
            if row.target_weight_bps > 0 {
                next.last_target_weights
                    .insert(row.symbol.into(), row.target_weight_bps);
            }
            let eligible = row.target_weight_bps > 0 && row.target_value > Decimal::ZERO;
            let completion_quantity = (row.status == "READY"
                && row.target_quantity > row.current_quantity)
                .then_some(row.target_quantity);
            let satisfied = matches!(
                row.status,
                "ALREADY_SATISFIED"
                    | "ENTRY_TARGET_ALREADY_SATISFIED"
                    | "BELOW_MINIMUM_TRADE_UNIT_ALREADY_SATISFIED"
            );
            if row.current_quantity == Decimal::ZERO && eligible && satisfied {
                next.entries.remove(row.symbol);
                continue;
            }
            if let Some(entry) = next.entries.get_mut(row.symbol) {
                entry.latest_generation = generation.into();
                entry.latest_target_value = row.target_value;
                if entry.pending && completion_quantity.is_some() {
                    entry.completion_quantity = completion_quantity;
                }
                entry.observed_holding |= row.current_quantity > Decimal::ZERO;
                if entry.pending && eligible && satisfied {
                    entry.pending = false;
                }
            } else if eligible && row.current_quantity == Decimal::ZERO && !satisfied {
                next.entries.insert(
                    row.symbol.into(),
                    StockPoolEntryProgress {
                        pending: true,
                        observed_holding: false,
                        first_decision_date: decision_date,
                        latest_generation: generation.into(),
                        latest_target_value: row.target_value,
                        completion_quantity,
                    },
                );
            }
        }
        next.validate()?;
        Ok(next)
    }
}