//! Performance metrics for speculative decoding

use std::collections::VecDeque;
use std::time::Duration;

// ---------------------------------------------------------------------------
// WindowedAcceptanceTracker
// ---------------------------------------------------------------------------

/// Sliding-window acceptance rate over the last `window_size` decode steps.
///
/// Each element in the window is the per-step acceptance rate
/// (`accepted_tokens / draft_tokens`) for that step, clamped to `[0.0, 1.0]`.
/// Oldest entries are evicted automatically once the window is full.
///
/// # Example
/// ```
/// use rtx_inference::speculative::WindowedAcceptanceTracker;
///
/// let mut tracker = WindowedAcceptanceTracker::new(3);
/// tracker.push(0.5);
/// tracker.push(0.8);
/// tracker.push(0.7);
/// let rate = tracker.windowed_rate();
/// assert!((rate - (0.5_f32 + 0.8 + 0.7) / 3.0).abs() < 1e-5);
/// ```
#[derive(Debug, Clone)]
pub struct WindowedAcceptanceTracker {
    window_size: usize,
    buffer: VecDeque<f32>,
    sum: f32,
}

impl WindowedAcceptanceTracker {
    /// Create a new tracker with the given sliding-window capacity.
    ///
    /// # Panics
    /// Panics if `window_size` is 0.
    pub fn new(window_size: usize) -> Self {
        assert!(window_size > 0, "window_size must be > 0");
        Self {
            window_size,
            buffer: VecDeque::with_capacity(window_size),
            sum: 0.0,
        }
    }

    /// Push the per-step acceptance rate for one decode step.
    ///
    /// `step_rate` is clamped to `[0.0, 1.0]` before storage.
    /// If the window is full the oldest value is evicted and its contribution
    /// subtracted from the running sum before the new value is inserted.
    pub fn push(&mut self, step_rate: f32) {
        let clamped = step_rate.clamp(0.0, 1.0);
        if self.buffer.len() == self.window_size {
            // Evict the oldest entry.
            if let Some(oldest) = self.buffer.pop_front() {
                self.sum -= oldest;
            }
        }
        self.buffer.push_back(clamped);
        self.sum += clamped;
    }

    /// Mean acceptance rate over all values currently in the window.
    ///
    /// Returns `0.0` when the window is empty.
    pub fn windowed_rate(&self) -> f32 {
        if self.buffer.is_empty() {
            0.0
        } else {
            // Guard against accumulated floating-point drift going negative.
            (self.sum / self.buffer.len() as f32).max(0.0)
        }
    }

    /// The `p`-th percentile of values currently in the window (0.0 – 1.0).
    ///
    /// Uses nearest-rank method: index = `(p * (len - 1)).round()`.
    /// Returns `0.0` when the window is empty.
    pub fn percentile(&self, p: f32) -> f32 {
        if self.buffer.is_empty() {
            return 0.0;
        }
        let mut sorted: Vec<f32> = self.buffer.iter().copied().collect();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        let idx = (p * (sorted.len() as f32 - 1.0)).round() as usize;
        sorted[idx.min(sorted.len() - 1)]
    }

    /// 50th-percentile (median) of the current window.
    pub fn p50(&self) -> f32 {
        self.percentile(0.5)
    }

    /// 95th-percentile of the current window.
    pub fn p95(&self) -> f32 {
        self.percentile(0.95)
    }

    /// Minimum value in the current window, or `0.0` if empty.
    pub fn min(&self) -> f32 {
        if self.buffer.is_empty() {
            return 0.0;
        }
        self.buffer.iter().copied().fold(f32::INFINITY, f32::min)
    }

    /// Maximum value in the current window, or `0.0` if empty.
    pub fn max(&self) -> f32 {
        if self.buffer.is_empty() {
            return 0.0;
        }
        self.buffer
            .iter()
            .copied()
            .fold(f32::NEG_INFINITY, f32::max)
    }

    /// Number of samples currently in the window.
    pub fn count(&self) -> usize {
        self.buffer.len()
    }

    /// `true` once `window_size` samples have been pushed (or more).
    pub fn is_full(&self) -> bool {
        self.buffer.len() == self.window_size
    }
}

// ---------------------------------------------------------------------------
// AcceptanceTrend
// ---------------------------------------------------------------------------

/// Direction of acceptance rate movement relative to the global average.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AcceptanceTrend {
    /// Windowed rate is more than 5 percentage points above the global rate.
    Rising,
    /// Windowed rate is more than 5 percentage points below the global rate.
    Falling,
    /// Windowed rate is within ±5 percentage points of the global rate.
    Stable,
}

impl AcceptanceTrend {
    fn from_rates(windowed: f32, global: f32) -> Self {
        const THRESHOLD: f32 = 0.05;
        let delta = windowed - global;
        if delta > THRESHOLD {
            Self::Rising
        } else if delta < -THRESHOLD {
            Self::Falling
        } else {
            Self::Stable
        }
    }
}

// ---------------------------------------------------------------------------
// AcceptanceDashboard
// ---------------------------------------------------------------------------

/// Snapshot of acceptance rate statistics for real-time monitoring.
///
/// Obtained via [`PerformanceMetrics::dashboard`].
#[derive(Debug, Clone)]
pub struct AcceptanceDashboard {
    /// Global running acceptance rate since the decoder started.
    pub global_rate: f32,
    /// Mean acceptance rate over the last `window_size` steps.
    pub windowed_rate: f32,
    /// Median (p50) of the windowed samples.
    pub p50: f32,
    /// 95th-percentile of the windowed samples.
    pub p95: f32,
    /// Minimum acceptance rate observed in the current window.
    pub min: f32,
    /// Maximum acceptance rate observed in the current window.
    pub max: f32,
    /// Total decode steps recorded (lifetime of the decoder).
    pub total_steps: u64,
    /// Configured window size.
    pub window_size: usize,
    /// `true` once `window_size` samples have accumulated.
    pub window_full: bool,
    /// Trend of the windowed rate relative to the global rate.
    pub trend: AcceptanceTrend,
}

// ---------------------------------------------------------------------------
// PerformanceMetrics (extended)
// ---------------------------------------------------------------------------

/// Performance metrics for speculative decoding
#[derive(Debug, Clone, Default)]
pub struct PerformanceMetrics {
    /// Total number of decoding steps
    pub total_steps: u64,
    /// Total tokens generated by draft model
    pub tokens_generated: u64,
    /// Total tokens accepted by target model
    pub tokens_accepted: u64,
    /// Current acceptance rate
    pub acceptance_rate: f32,
    /// Speedup ratio compared to sequential generation
    pub speedup_ratio: f32,
    /// Average time for draft generation
    pub draft_generation_time: Duration,
    /// Average time for verification
    pub verification_time: Duration,
    /// Total time spent in speculative decoding
    pub total_time: Duration,
    /// Sliding-window acceptance tracker (default window = 100 steps).
    pub(crate) windowed: WindowedAcceptanceTracker,
}

impl Default for WindowedAcceptanceTracker {
    fn default() -> Self {
        Self::new(100)
    }
}

impl PerformanceMetrics {
    pub(crate) fn new() -> Self {
        Self {
            total_steps: 0,
            tokens_generated: 0,
            tokens_accepted: 0,
            acceptance_rate: 0.0,
            speedup_ratio: 0.0,
            draft_generation_time: Duration::ZERO,
            verification_time: Duration::ZERO,
            total_time: Duration::ZERO,
            windowed: WindowedAcceptanceTracker::new(100),
        }
    }

    pub(crate) fn update(&mut self, step_result: &StepMetrics) {
        self.total_steps += 1;
        self.tokens_generated += step_result.draft_tokens as u64;
        self.tokens_accepted += step_result.accepted_tokens as u64;

        // Update running averages
        let weighted_draft = self.draft_generation_time.as_nanos()
            * u128::from(self.total_steps - 1)
            + step_result.draft_time.as_nanos();
        self.draft_generation_time =
            Duration::from_nanos((weighted_draft / u128::from(self.total_steps)) as u64);

        let weighted_verification = self.verification_time.as_nanos()
            * u128::from(self.total_steps - 1)
            + step_result.verification_time.as_nanos();
        self.verification_time =
            Duration::from_nanos((weighted_verification / u128::from(self.total_steps)) as u64);
        self.total_time += step_result.total_time;

        // Update acceptance rate
        self.acceptance_rate = if self.tokens_generated > 0 {
            self.tokens_accepted as f32 / self.tokens_generated as f32
        } else {
            0.0
        };

        // Push per-step acceptance rate into the sliding window.
        let step_rate = step_result.accepted_tokens as f32 / step_result.draft_tokens.max(1) as f32;
        self.windowed.push(step_rate);

        // Estimate speedup ratio
        // This is simplified - in reality would compare against baseline sequential generation
        let draft_efficiency =
            step_result.draft_tokens as f32 / step_result.draft_time.as_secs_f32();
        let acceptance_benefit = self.acceptance_rate * step_result.draft_tokens as f32;
        self.speedup_ratio = 1.0 + (acceptance_benefit / draft_efficiency).min(10.0);
    }

    /// Build a real-time [`AcceptanceDashboard`] from the current state.
    pub fn dashboard(&self) -> AcceptanceDashboard {
        let windowed_rate = self.windowed.windowed_rate();
        let trend = AcceptanceTrend::from_rates(windowed_rate, self.acceptance_rate);
        AcceptanceDashboard {
            global_rate: self.acceptance_rate,
            windowed_rate,
            p50: self.windowed.p50(),
            p95: self.windowed.p95(),
            min: self.windowed.min(),
            max: self.windowed.max(),
            total_steps: self.total_steps,
            window_size: self.windowed.window_size,
            window_full: self.windowed.is_full(),
            trend,
        }
    }
}

/// Metrics for a single decoding step
pub(crate) struct StepMetrics {
    pub draft_tokens: usize,
    pub accepted_tokens: usize,
    pub draft_time: Duration,
    pub verification_time: Duration,
    pub total_time: Duration,
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    // ------------------------------------------------------------------
    // WindowedAcceptanceTracker tests
    // ------------------------------------------------------------------

    #[test]
    fn test_windowed_tracker_empty() {
        let tracker = WindowedAcceptanceTracker::new(10);
        assert_eq!(tracker.windowed_rate(), 0.0);
        assert_eq!(tracker.count(), 0);
        assert!(!tracker.is_full());
    }

    #[test]
    fn test_windowed_tracker_push_and_evict() {
        let mut tracker = WindowedAcceptanceTracker::new(3);
        tracker.push(0.1); // [0.1]
        tracker.push(0.2); // [0.1, 0.2]
        tracker.push(0.3); // [0.1, 0.2, 0.3]  — full
        tracker.push(0.4); // [0.2, 0.3, 0.4]  — 0.1 evicted
        assert_eq!(tracker.count(), 3);
        assert!(tracker.is_full());
        // Sum should be 0.2 + 0.3 + 0.4 = 0.9 → mean ≈ 0.3
        let rate = tracker.windowed_rate();
        assert!((rate - 0.3).abs() < 1e-5, "expected ~0.3, got {rate}");
    }

    #[test]
    fn test_windowed_rate_accuracy() {
        let mut tracker = WindowedAcceptanceTracker::new(10);
        tracker.push(0.5);
        tracker.push(0.8);
        tracker.push(0.7);
        let expected = (0.5_f32 + 0.8 + 0.7) / 3.0;
        let got = tracker.windowed_rate();
        assert!(
            (got - expected).abs() < 1e-5,
            "expected {expected}, got {got}"
        );
    }

    #[test]
    fn test_percentile_p50() {
        let mut tracker = WindowedAcceptanceTracker::new(10);
        tracker.push(0.4);
        tracker.push(0.6);
        tracker.push(0.8);
        // sorted: [0.4, 0.6, 0.8]; idx = round(0.5 * 2) = 1 → 0.6
        assert!((tracker.p50() - 0.6).abs() < 1e-5);
    }

    #[test]
    fn test_percentile_p95() {
        let mut tracker = WindowedAcceptanceTracker::new(100);
        // Push 20 evenly spaced values: 0.0, 1/19, 2/19 ... 1.0
        for i in 0..20_u32 {
            tracker.push(i as f32 / 19.0);
        }
        // p95: idx = round(0.95 * 19) = round(18.05) = 18 → value = 18/19 ≈ 0.947
        let p95 = tracker.p95();
        assert!(p95 > 0.9, "p95 should be close to 1.0, got {p95}");
    }

    #[test]
    fn test_trend_rising() {
        let trend = AcceptanceTrend::from_rates(0.70, 0.50);
        assert_eq!(trend, AcceptanceTrend::Rising);
    }

    #[test]
    fn test_trend_falling() {
        let trend = AcceptanceTrend::from_rates(0.50, 0.70);
        assert_eq!(trend, AcceptanceTrend::Falling);
    }

    #[test]
    fn test_trend_stable() {
        let trend = AcceptanceTrend::from_rates(0.52, 0.50);
        assert_eq!(trend, AcceptanceTrend::Stable);
    }

    #[test]
    fn test_dashboard_builds_correctly() {
        let mut metrics = PerformanceMetrics::new();
        let step = StepMetrics {
            draft_tokens: 4,
            accepted_tokens: 3,
            draft_time: Duration::from_millis(5),
            verification_time: Duration::from_millis(2),
            total_time: Duration::from_millis(7),
        };
        metrics.update(&step);
        metrics.update(&step);
        metrics.update(&step);

        let dash = metrics.dashboard();
        assert_eq!(dash.total_steps, 3);
        assert_eq!(dash.window_size, 100);
        assert!(!dash.window_full); // need 100 steps for full
        // Each step: accepted=3, drafted=4 → step_rate = 0.75
        assert!((dash.windowed_rate - 0.75).abs() < 1e-4);
        assert!((dash.global_rate - 0.75).abs() < 1e-4);
        assert_eq!(dash.trend, AcceptanceTrend::Stable);
    }

    #[test]
    fn test_window_full_flag() {
        let mut tracker = WindowedAcceptanceTracker::new(5);
        for i in 0..5 {
            assert!(!tracker.is_full(), "should not be full after {i} pushes");
            tracker.push(0.5);
        }
        assert!(tracker.is_full());
    }

    #[test]
    fn test_min_max() {
        let mut tracker = WindowedAcceptanceTracker::new(10);
        tracker.push(0.3);
        tracker.push(0.7);
        tracker.push(0.9);
        assert!(
            (tracker.min() - 0.3).abs() < 1e-5,
            "min got {}",
            tracker.min()
        );
        assert!(
            (tracker.max() - 0.9).abs() < 1e-5,
            "max got {}",
            tracker.max()
        );
    }

    #[test]
    fn test_min_empty_returns_zero() {
        let tracker = WindowedAcceptanceTracker::new(5);
        assert_eq!(tracker.min(), 0.0, "min() of empty window must be 0.0");
        assert_eq!(tracker.max(), 0.0, "max() of empty window must be 0.0");
        assert_eq!(tracker.windowed_rate(), 0.0);
        assert_eq!(tracker.count(), 0);
    }

    #[test]
    fn test_clamp_out_of_range() {
        let mut tracker = WindowedAcceptanceTracker::new(5);
        tracker.push(-0.5); // clamped to 0.0
        tracker.push(1.5); // clamped to 1.0
        assert!((tracker.windowed_rate() - 0.5).abs() < 1e-5);
    }
}