use argmin::core::observers::Observe;
use argmin::core::{Error, State};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

/// Observer that enforces strict trial budget limits
///
/// Argmin's PSO does not reliably respect max_iters when set to small values (e.g., 1).
/// This observer shares the same `Arc<AtomicUsize>` counter that the cost function and
/// LHS evaluation phase use, so every evaluation (LHS + PSO) counts toward the budget.
#[derive(Clone, Debug)]
pub struct TrialBudgetObserver {
    max_trials: usize,
    trial_counter: Arc<AtomicUsize>,
}

impl TrialBudgetObserver {
    /// Create a new observer that shares the caller-owned trial counter.
    ///
    /// The same `trial_counter` must be passed to the cost function and incremented
    /// on every evaluation (LHS and PSO) so the budget is accurately tracked.
    pub fn new(max_trials: usize, trial_counter: Arc<AtomicUsize>) -> Self {
        Self {
            max_trials,
            trial_counter,
        }
    }

    pub fn should_terminate(&self) -> bool {
        self.trial_counter.load(Ordering::SeqCst) >= self.max_trials
    }

    pub fn get_trials_used(&self) -> usize {
        self.trial_counter.load(Ordering::SeqCst)
    }
}

impl<I> Observe<I> for TrialBudgetObserver
where
    I: State,
{
    fn observe_iter(&mut self, _state: &I, _kv: &argmin::core::KV) -> Result<(), Error> {
        let used = self.trial_counter.load(Ordering::SeqCst);
        if used >= self.max_trials {
            tracing::warn!(
                "Trial budget exhausted: {}/{} trials used. Terminating optimization.",
                used,
                self.max_trials
            );
            return Err(Error::msg(format!(
                "Trial budget exhausted: {used}/{}",
                self.max_trials
            )));
        }

        Ok(())
    }
}