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 tracks evaluations across all iterations and terminates when budget is hit.
#[derive(Clone, Debug)]
pub struct TrialBudgetObserver {
    max_trials: usize,
    trials_used: Arc<AtomicUsize>,
}

impl TrialBudgetObserver {
    pub fn new(max_trials: usize) -> Self {
        Self {
            max_trials,
            trials_used: Arc::new(AtomicUsize::new(0)),
        }
    }

    pub fn increment_trial(&self) {
        self.trials_used.fetch_add(1, Ordering::Relaxed);
    }

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

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

impl<I> Observe<I> for TrialBudgetObserver
where
    I: State,
{
    fn observe_iter(&mut self, _state: &I, _kv: &argmin::core::KV) -> Result<(), Error> {
        // Check if we've exceeded budget
        if self.should_terminate() {
            tracing::warn!(
                "Trial budget exhausted: {}/{} trials used. Terminating optimization.",
                self.get_trials_used(),
                self.max_trials
            );
            // Return error to signal termination
            return Err(Error::msg(format!(
                "Trial budget exhausted: {}/{} trials",
                self.get_trials_used(),
                self.max_trials
            )));
        }

        Ok(())
    }
}