//! Particle Filter implementation.
//!
//! Works by maintaining many "guesses" (particles) about the state,
//! each with a weight representing how likely that guess is.
//!
//! This is a very flexible way of estimating the state, and can
//! handle models of almost any degree of complexity, including
//! multi-modal models, non-numeric models, and models that can not be
//! expressed in closed-form expressions.

use super::*;
use rand::Rng;
use rand::distr::Distribution;

/// Describes how your system changes over time (particle filter version).
///
/// Unlike the EKF version, this samples random outcomes directly.
///
/// # Example
///
/// ```
/// use estimators::particle::ProcessModel;
/// use rand::Rng;
///
/// struct NoisyRobot {
///     speed_noise: f64,
///     turn_noise: f64,
/// }
///
/// impl ProcessModel<[f64; 3], [f64; 2]> for NoisyRobot {
///     fn sample_next<R: Rng>(
///         &self,
///         state: &[f64; 3],
///         control: &[f64; 2],
///         dt: f64,
///         rng: &mut R,
///     ) -> [f64; 3] {
///         let [x, y, heading] = *state;
///         let [speed, turn_rate] = *control;
///
///         // Add random noise to simulate real-world uncertainty
///         let noisy_speed = speed + rng.random_range(-self.speed_noise..self.speed_noise);
///         let noisy_turn = turn_rate + rng.random_range(-self.turn_noise..self.turn_noise);
///
///         [
///             x + noisy_speed * heading.cos() * dt,
///             y + noisy_speed * heading.sin() * dt,
///             heading + noisy_turn * dt,
///         ]
///     }
/// }
///
/// // Test it
/// let model = NoisyRobot { speed_noise: 0.1, turn_noise: 0.05 };
/// let mut rng = rand::rng();
/// let state = [0.0, 0.0, 0.0];
/// let control = [1.0, 0.0];
/// let next = model.sample_next(&state, &control, 0.1, &mut rng);
/// // Position should have moved roughly in the x direction
/// assert!(next[0] > 0.0);
/// ```
pub trait ProcessModel<S, C> {
    /// Sample a possible next state, including random noise.
    ///
    /// This should add appropriate random variations to simulate
    /// the uncertainty in your system's motion.
    fn sample_next<R: Rng>(&self, state: &S, control: &C, dt: f64, rng: &mut R) -> S;
}

/// Describes how likely a measurement is given a particular state.
///
/// # Example
///
/// ```
/// use estimators::particle::MeasurementModel;
///
/// struct RangeSensor {
///     beacon_x: f64,
///     beacon_y: f64,
///     noise_std: f64,
/// }
///
/// impl MeasurementModel<[f64; 3], f64> for RangeSensor {
///     fn likelihood(&self, state: &[f64; 3], measured_range: &f64) -> f64 {
///         let [x, y, _] = *state;
///
///         // Expected range to beacon
///         let dx = self.beacon_x - x;
///         let dy = self.beacon_y - y;
///         let expected_range = (dx*dx + dy*dy).sqrt();
///
///         // How likely is this measurement given the expected range?
///         // (Gaussian probability)
///         let error = measured_range - expected_range;
///         (-0.5 * (error / self.noise_std).powi(2)).exp()
///     }
/// }
///
/// // Test it
/// let sensor = RangeSensor { beacon_x: 10.0, beacon_y: 0.0, noise_std: 1.0 };
/// let state = [0.0, 0.0, 0.0];  // At origin
/// let measured = 10.0;          // Measured range = 10 (exactly right!)
/// let likelihood = sensor.likelihood(&state, &measured);
/// assert!((likelihood - 1.0).abs() < 0.01);  // Should be ~1.0
/// ```
pub trait MeasurementModel<S, M> {
    /// How likely is this measurement if the state were `state`?
    ///
    /// Return a value between 0 (impossible) and 1 (perfect match).
    /// The values don't need to be exact probabilities — they just need
    /// to be higher for better matches.
    fn likelihood(&self, state: &S, measurement: &M) -> f64;
}

/// For states that can be averaged together.
///
/// Implement this for custom state types. Already implemented for
/// fixed-size float arrays like `[f64; 3]`.
pub trait WeightedAverage: Clone {
    /// Compute the weighted average of multiple states.
    fn weighted_average<'a, I>(weighted_states: I) -> Self
    where
        I: Iterator<Item = (&'a Self, f64)>,
        Self: 'a;
}

impl<const N: usize> WeightedAverage for [f64; N] {
    fn weighted_average<'a, I>(weighted_states: I) -> Self
    where
        I: Iterator<Item = (&'a Self, f64)>,
    {
        let mut result = [0.0; N];
        for (state, weight) in weighted_states {
            for (r, s) in result.iter_mut().zip(state.iter()) {
                *r += s * weight;
            }
        }
        result
    }
}

/// Configuration for the particle filter.
#[derive(Debug, Clone)]
pub struct Config {
    /// Number of particles to use. More = better accuracy but slower.
    /// Typical values: 100-10000
    pub num_particles: usize,

    /// When particle diversity drops below this fraction, resample.
    /// Default: 0.5 (resample when effective particles < 50% of total)
    pub resample_threshold: f64,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            num_particles: 1000,
            resample_threshold: 0.5,
        }
    }
}

/// Particle Filter.
///
/// Tracks state using many weighted samples ("particles").
///
/// # Example
///
/// ```
/// use estimators::particle::{ProcessModel, MeasurementModel, WeightedAverage, ParticleFilter, Config};
/// use estimators::StateEstimator;
/// use rand::{Rng, SeedableRng};
/// use rand::distr::Distribution;
/// use rand_chacha::ChaCha8Rng;
///
/// // Simple 1D state: just position
/// struct RandomWalk;
///
/// impl ProcessModel<[f64; 1], ()> for RandomWalk {
///     fn sample_next<R: Rng>(&self, state: &[f64; 1], _: &(), dt: f64, rng: &mut R) -> [f64; 1] {
///         [state[0] + rng.random_range(-0.1..0.1) * dt]
///     }
/// }
///
/// struct PositionSensor;
///
/// impl MeasurementModel<[f64; 1], f64> for PositionSensor {
///     fn likelihood(&self, state: &[f64; 1], measurement: &f64) -> f64 {
///         let error = state[0] - measurement;
///         (-0.5 * (error / 0.5).powi(2)).exp()
///     }
/// }
///
/// // Distribution to sample initial particles
/// struct Uniform { center: f64, spread: f64 }
///
/// impl Distribution<[f64; 1]> for Uniform {
///     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> [f64; 1] {
///         [self.center + (rng.random::<f64>() - 0.5) * self.spread]
///     }
/// }
///
/// // Create and run filter
/// let mut pf = ParticleFilter::new(
///     Uniform { center: 0.0, spread: 2.0 },
///     Config { num_particles: 100, ..Default::default() },
///     RandomWalk,
///     PositionSensor,
///     ChaCha8Rng::seed_from_u64(42),
/// );
///
/// // Run a few steps
/// for _ in 0..5 {
///     pf.predict(&(), 0.1);
///     pf.update(&0.0).unwrap();  // Measuring position ~0
/// }
///
/// let estimate = pf.estimate();
/// assert!(estimate[0].abs() < 1.0);  // Should be near 0
/// ```
pub struct ParticleFilter<S, M, C, PM, MM, R> {
    particles: Vec<S>,
    weights: Vec<f64>,
    config: Config,
    process_model: PM,
    measurement_model: MM,
    rng: R,
    _phantom: std::marker::PhantomData<(M, C)>,
}

impl<S, M, C, PM, MM, R> ParticleFilter<S, M, C, PM, MM, R>
where
    S: Clone + Debug + WeightedAverage,
    PM: ProcessModel<S, C>,
    MM: MeasurementModel<S, M>,
    R: Rng,
{
    /// Create a new particle filter.
    ///
    /// # Arguments
    /// * `initial_distribution` - How to generate initial particle guesses
    /// * `config` - Filter configuration (number of particles, etc.)
    /// * `process_model` - Describes how state evolves
    /// * `measurement_model` - Describes sensor likelihood
    /// * `rng` - Random number generator
    pub fn new<D>(
        initial_distribution: D,
        config: Config,
        process_model: PM,
        measurement_model: MM,
        mut rng: R,
    ) -> Self
    where
        D: Distribution<S>,
    {
        let n = config.num_particles;
        let particles: Vec<S> = (0..n)
            .map(|_| initial_distribution.sample(&mut rng))
            .collect();
        let weights = vec![1.0 / n as f64; n];

        Self {
            particles,
            weights,
            config,
            process_model,
            measurement_model,
            rng,
            _phantom: std::marker::PhantomData,
        }
    }

    /// Create from a list of initial particles.
    pub fn from_particles(
        particles: Vec<S>,
        config: Config,
        process_model: PM,
        measurement_model: MM,
        rng: R,
    ) -> Self {
        let n = particles.len();
        let weights = vec![1.0 / n as f64; n];

        Self {
            particles,
            weights,
            config,
            process_model,
            measurement_model,
            rng,
            _phantom: std::marker::PhantomData,
        }
    }

    /// Get the most likely particle (maximum a posteriori estimate).
    ///
    /// # Errors
    ///
    /// Returns `NoParticles` if the filter has no particles.
    pub fn best_particle(&self) -> Result<&S> {
        self.weights
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .map(|(idx, _)| &self.particles[idx])
            .ok_or(EstimatorError::NoParticles)
    }

    fn normalize_weights(&mut self) -> Result<()> {
        let sum: f64 = self.weights.iter().sum();
        if sum > 0.0 && sum.is_finite() {
            for w in &mut self.weights {
                *w /= sum;
            }
            Ok(())
        } else if sum == 0.0 {
            // All weights are zero - particles don't match measurements
            Err(EstimatorError::ParticleDegeneracy)
        } else {
            Err(EstimatorError::NumericalInstability(format!(
                "weight sum = {}",
                sum
            )))
        }
    }

    fn effective_sample_size(&self) -> f64 {
        let sum_sq: f64 = self.weights.iter().map(|w| w * w).sum();
        if sum_sq > 0.0 { 1.0 / sum_sq } else { 0.0 }
    }

    fn should_resample(&self) -> bool {
        let ess = self.effective_sample_size();
        let threshold = self.config.resample_threshold * self.particles.len() as f64;
        ess < threshold
    }

    fn resample(&mut self) {
        // Systematic resampling (low variance)
        let n = self.particles.len();
        let mut cumsum = Vec::with_capacity(n);
        let mut sum = 0.0;
        for &w in &self.weights {
            sum += w;
            cumsum.push(sum);
        }

        let u0: f64 = self.rng.random::<f64>() / n as f64;
        let mut new_particles = Vec::with_capacity(n);
        let mut idx = 0;

        for i in 0..n {
            let u = u0 + i as f64 / n as f64;
            while idx < n - 1 && cumsum[idx] < u {
                idx += 1;
            }
            new_particles.push(self.particles[idx].clone());
        }

        self.particles = new_particles;
        self.weights = vec![1.0 / n as f64; n];
    }
}

impl<S, M, C, PM, MM, R> StateEstimator for ParticleFilter<S, M, C, PM, MM, R>
where
    S: Clone + Debug + WeightedAverage,
    PM: ProcessModel<S, C>,
    MM: MeasurementModel<S, M>,
    R: Rng,
{
    type State = S;
    type Measurement = M;
    type Control = C;

    fn predict(&mut self, control: &C, dt: f64) {
        for particle in &mut self.particles {
            *particle = self
                .process_model
                .sample_next(particle, control, dt, &mut self.rng);
        }
    }

    fn update(&mut self, measurement: &M) -> Result<()> {
        if self.particles.is_empty() {
            return Err(EstimatorError::NoParticles);
        }

        for (particle, weight) in self.particles.iter().zip(self.weights.iter_mut()) {
            *weight *= self.measurement_model.likelihood(particle, measurement);
        }

        self.normalize_weights()?;

        if self.should_resample() {
            self.resample();
        }

        Ok(())
    }

    fn estimate(&self) -> S {
        S::weighted_average(self.particles.iter().zip(self.weights.iter().copied()))
    }
}

impl<S, M, C, PM, MM, R> super::ParticleEstimator for ParticleFilter<S, M, C, PM, MM, R>
where
    S: Clone + Debug + WeightedAverage,
    PM: ProcessModel<S, C>,
    MM: MeasurementModel<S, M>,
    R: Rng,
{
    fn particles(&self) -> impl Iterator<Item = (&Self::State, f64)> {
        self.particles.iter().zip(self.weights.iter().copied())
    }

    fn num_particles(&self) -> usize {
        self.particles.len()
    }

    fn particle_diversity(&self) -> f64 {
        self.effective_sample_size() / self.particles.len() as f64
    }
}