//! UCB-style count-based exploration bonus for action selection.
//!
//! Adds bonus = β × sqrt(ln(N) / (1 + `n_a`)) to Q-values during action selection,
//! where N = total actions taken this epoch, `n_a` = count of action a.
//! This gives a direct exploration incentive to under-visited actions,
//! preventing collapse onto a single exposure level across all states.

/// Count-based exploration bonus tracker for DQN action selection.
///
/// At each greedy action selection, bonuses are added to Q-values before argmax.
/// Actions taken less frequently get higher bonuses, encouraging diversity.
/// Counts are reset at epoch boundaries.
///
/// Supports both flat (N exposure actions) and branched (N exposure + M order
/// type + K urgency) tracking for Branching DQN. All dimensions are dynamic —
/// derived from `DQNConfig` at construction, not hardcoded constants.
#[derive(Debug)]
pub struct CountBonus {
    action_counts: Vec<u64>,
    order_counts: Vec<u64>,
    urgency_counts: Vec<u64>,
    total_count: u64,
    coefficient: f64,
}

impl CountBonus {
    /// Create a new count bonus tracker with the given coefficient β and branch sizes.
    ///
    /// β controls the strength of the exploration bonus.
    /// Typical values: 0.01–0.5 (tunable via hyperopt).
    ///
    /// # Arguments
    ///
    /// * `coefficient` - UCB exploration coefficient β
    /// * `num_actions` - Number of exposure actions (e.g. 5 for Short100..Long100)
    /// * `num_order_types` - Number of order type actions (e.g. 3 for Market/LimitMaker/IoC)
    /// * `num_urgency_levels` - Number of urgency actions (e.g. 3 for Patient/Normal/Aggressive)
    pub fn new(
        coefficient: f64,
        num_actions: usize,
        num_order_types: usize,
        num_urgency_levels: usize,
    ) -> Self {
        Self {
            action_counts: vec![0_u64; num_actions],
            order_counts: vec![0_u64; num_order_types],
            urgency_counts: vec![0_u64; num_urgency_levels],
            total_count: 0,
            coefficient,
        }
    }

    /// Record that an action was taken (flat / exposure-only path).
    pub fn record_action(&mut self, action_idx: usize) {
        if action_idx < self.action_counts.len() {
            self.action_counts[action_idx] += 1;
            self.total_count += 1;
        }
    }

    /// Record a branched action (exposure, `order_type`, urgency).
    ///
    /// Increments the per-branch count for each head and bumps `total_count`
    /// once (keyed to the exposure branch, consistent with the flat path).
    /// Out-of-range indices for any branch are silently ignored.
    pub fn record_branched_action(
        &mut self,
        exposure_idx: usize,
        order_idx: usize,
        urgency_idx: usize,
    ) {
        if exposure_idx < self.action_counts.len() {
            self.action_counts[exposure_idx] += 1;
            self.total_count += 1;
        }
        if order_idx < self.order_counts.len() {
            self.order_counts[order_idx] += 1;
        }
        if urgency_idx < self.urgency_counts.len() {
            self.urgency_counts[urgency_idx] += 1;
        }
    }

    /// Compute UCB-style bonus for each action.
    ///
    /// Formula: `bonus_a` = β × sqrt(ln(N) / (1 + `n_a`))
    /// - When `n_a` = 0 and N > 0: bonus is at maximum (strong incentive to try)
    /// - When `n_a` is large: bonus approaches 0 (already well-explored)
    /// - When N = 0: returns uniform bonus β (no data yet)
    pub fn bonuses(&self) -> Vec<f64> {
        let n = self.action_counts.len();
        if self.total_count == 0 {
            return vec![self.coefficient; n];
        }
        let log_n = (self.total_count as f64).ln();
        self.action_counts
            .iter()
            .map(|&count| self.coefficient * (log_n / (1.0 + count as f64)).sqrt())
            .collect()
    }

    /// Per-branch UCB bonuses for Branching DQN.
    ///
    /// Returns `(exposure, order_type, urgency)` bonus vectors using the same
    /// UCB formula as [`bonuses()`](Self::bonuses). All three branches share
    /// the same `total_count` denominator so the exploration pressure scales
    /// uniformly across heads.
    pub fn bonuses_branched(&self) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
        let exposure_bonuses = self.bonuses();

        let order_bonuses = self.compute_branch_bonuses(&self.order_counts);
        let urgency_bonuses = self.compute_branch_bonuses(&self.urgency_counts);

        (exposure_bonuses, order_bonuses, urgency_bonuses)
    }

    /// Compute UCB bonuses for a single branch's count vector.
    fn compute_branch_bonuses(&self, counts: &[u64]) -> Vec<f64> {
        if self.total_count == 0 {
            return vec![self.coefficient; counts.len()];
        }
        let log_n = (self.total_count as f64).ln();
        counts
            .iter()
            .map(|&count| self.coefficient * (log_n / (1.0 + count as f64)).sqrt())
            .collect()
    }

    /// Reset all counts (call at epoch boundary).
    pub fn reset(&mut self) {
        self.action_counts.fill(0);
        self.order_counts.fill(0);
        self.urgency_counts.fill(0);
        self.total_count = 0;
    }

    /// Get total number of actions recorded since last reset.
    pub const fn total_count(&self) -> u64 {
        self.total_count
    }

    /// Get the coefficient β.
    pub const fn coefficient(&self) -> f64 {
        self.coefficient
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_new_bonus_returns_uniform_coefficient() {
        let cb = CountBonus::new(0.1, 5, 3, 3);
        let bonuses = cb.bonuses();
        for b in &bonuses {
            assert!((*b - 0.1).abs() < 1e-10, "Expected 0.1, got {}", b);
        }
    }

    #[test]
    fn test_frequent_action_gets_low_bonus() {
        let mut cb = CountBonus::new(0.1, 5, 3, 3);
        // Record action 4 (Long100) one hundred times
        for _ in 0..100 {
            cb.record_action(4);
        }
        let bonuses = cb.bonuses();
        // Action 4 should have much lower bonus than others
        let bonus_4 = bonuses[4];
        let bonus_0 = bonuses[0]; // Never taken
        assert!(
            bonus_0 > bonus_4 * 5.0,
            "Unvisited bonus ({}) should be >> frequent bonus ({})",
            bonus_0,
            bonus_4
        );
    }

    #[test]
    fn test_uniform_usage_gives_equal_bonuses() {
        let mut cb = CountBonus::new(0.1, 5, 3, 3);
        // Take each action exactly once
        for i in 0..5 {
            cb.record_action(i);
        }
        let bonuses = cb.bonuses();
        let first = bonuses[0];
        for (i, b) in bonuses.iter().enumerate() {
            assert!(
                (*b - first).abs() < 1e-10,
                "Bonus {} ({}) should equal bonus 0 ({})",
                i,
                b,
                first
            );
        }
    }

    #[test]
    fn test_reset_clears_all_counts() {
        let mut cb = CountBonus::new(0.1, 5, 3, 3);
        for _ in 0..50 {
            cb.record_action(2);
        }
        assert_eq!(cb.total_count(), 50);
        cb.reset();
        assert_eq!(cb.total_count(), 0);
        let bonuses = cb.bonuses();
        // After reset, should return uniform coefficient again
        for b in &bonuses {
            assert!((*b - 0.1).abs() < 1e-10);
        }
    }

    #[test]
    fn test_out_of_range_action_ignored() {
        let mut cb = CountBonus::new(0.1, 5, 3, 3);
        cb.record_action(5); // Out of range
        cb.record_action(100); // Out of range
        assert_eq!(cb.total_count(), 0);
    }

    #[test]
    fn test_branched_action_recording() {
        let mut cb = CountBonus::new(0.1, 5, 3, 3);
        cb.record_branched_action(0, 1, 2); // Short100, LimitMaker, Aggressive
        cb.record_branched_action(4, 0, 0); // Long100, Market, Patient
        assert_eq!(cb.total_count(), 2);
    }

    #[test]
    fn test_branched_action_out_of_range_ignored() {
        let mut cb = CountBonus::new(0.1, 5, 3, 3);
        // All indices out of range — nothing should be recorded
        cb.record_branched_action(5, 3, 3);
        assert_eq!(cb.total_count(), 0);
        // Only exposure valid
        cb.record_branched_action(0, 99, 99);
        assert_eq!(cb.total_count(), 1);
    }

    #[test]
    fn test_bonuses_branched_returns_correct_sizes() {
        let mut cb = CountBonus::new(0.1, 5, 3, 3);
        for _ in 0..10 {
            cb.record_branched_action(0, 0, 0);
        }
        let (exp, ord, urg) = cb.bonuses_branched();
        assert_eq!(exp.len(), 5);
        assert_eq!(ord.len(), 3);
        assert_eq!(urg.len(), 3);
        // Action 0 in each branch should have lower bonus than unvisited
        assert!(
            exp[1] > exp[0] * 2.0,
            "Unvisited exposure bonus ({}) should be >> visited ({})",
            exp[1],
            exp[0]
        );
        assert!(
            ord[1] > ord[0] * 2.0,
            "Unvisited order bonus ({}) should be >> visited ({})",
            ord[1],
            ord[0]
        );
        assert!(
            urg[1] > urg[0] * 2.0,
            "Unvisited urgency bonus ({}) should be >> visited ({})",
            urg[1],
            urg[0]
        );
    }

    #[test]
    fn test_reset_clears_branched_counts() {
        let mut cb = CountBonus::new(0.1, 5, 3, 3);
        cb.record_branched_action(0, 1, 2);
        cb.reset();
        assert_eq!(cb.total_count(), 0);
        let (exp, ord, urg) = cb.bonuses_branched();
        // All should return uniform coefficient after reset
        for b in &exp {
            assert!((*b - 0.1).abs() < 1e-10);
        }
        for b in &ord {
            assert!((*b - 0.1).abs() < 1e-10);
        }
        for b in &urg {
            assert!((*b - 0.1).abs() < 1e-10);
        }
    }

    #[test]
    fn test_bonuses_branched_zero_count_returns_uniform() {
        let cb = CountBonus::new(0.25, 5, 3, 3);
        let (exp, ord, urg) = cb.bonuses_branched();
        for b in &exp {
            assert!((*b - 0.25).abs() < 1e-10);
        }
        for b in &ord {
            assert!((*b - 0.25).abs() < 1e-10);
        }
        for b in &urg {
            assert!((*b - 0.25).abs() < 1e-10);
        }
    }

    #[test]
    fn test_flat_and_branched_exposure_agree() {
        // Recording via record_action vs record_branched_action (exposure part)
        // should produce the same exposure bonuses.
        let mut flat = CountBonus::new(0.1, 5, 3, 3);
        let mut branched = CountBonus::new(0.1, 5, 3, 3);
        for i in 0..5 {
            flat.record_action(i);
            branched.record_branched_action(i, i % 3, i % 3);
        }
        let flat_bonuses = flat.bonuses();
        let (branched_exp, _, _) = branched.bonuses_branched();
        for (f, b) in flat_bonuses.iter().zip(branched_exp.iter()) {
            assert!(
                (*f - *b).abs() < 1e-10,
                "Flat ({}) and branched ({}) exposure bonuses should match",
                f,
                b
            );
        }
    }

    #[test]
    fn test_dynamic_num_actions() {
        // Verify CountBonus works with any num_actions, not just 5
        let mut cb = CountBonus::new(0.1, 3, 3, 3);
        let bonuses = cb.bonuses();
        assert_eq!(bonuses.len(), 3);

        cb.record_action(0);
        cb.record_action(2);
        let bonuses = cb.bonuses();
        assert_eq!(bonuses.len(), 3);
        assert!(bonuses[1] > bonuses[0], "Unvisited should have higher bonus");
    }

    #[test]
    fn test_dynamic_branch_sizes() {
        // Non-standard branch sizes: 7 exposures, 4 order types, 2 urgency levels
        let mut cb = CountBonus::new(0.1, 7, 4, 2);
        let bonuses = cb.bonuses();
        assert_eq!(bonuses.len(), 7);

        cb.record_branched_action(6, 3, 1);
        let (exp, ord, urg) = cb.bonuses_branched();
        assert_eq!(exp.len(), 7);
        assert_eq!(ord.len(), 4);
        assert_eq!(urg.len(), 2);

        // Out-of-range for this configuration
        cb.record_branched_action(7, 4, 2);
        assert_eq!(cb.total_count(), 1); // Only first call counted
    }
}
