//! ComputeBrick: Token-centric, self-verifying compute units.
//!
//! Per CBTOP-SPEC-001 and SHOWCASE-BRICK-001, every compute operation is a brick with:
//! - Token budget (tok/sec performance target)
//! - Assertions (falsifiable correctness claims)
//! - Verification (self-checking via baseline comparison)
//!
//! # Toyota Way Integration
//! - **Jidoka**: Stop-the-line on budget violation
//! - **Poka-Yoke**: Type-safe brick composition
//! - **Genchi Genbutsu**: Real metrics from hardware
//! - **Mieruka**: Visual control via cbtop TUI
//!
//! # References
//! - Popper, K. (1959). "The Logic of Scientific Discovery."
//! - Ohno, T. (1988). "Toyota Production System."
//! - Little, J.D.C. (1961). "A Proof for the Queuing Formula: L = λW."

use std::fmt;
use std::time::Instant;

use crate::quantize::Q8_0Block;

// PMAT-802: Extracted modules
#[cfg(feature = "cuda")]
mod fused;
#[cfg(feature = "cuda")]
pub use fused::{CoalescedDp4aBrick, FusedFfnBrick};

// Phase 14: BrickTracer for GPU/CPU parity debugging
pub mod tracer;
pub use tracer::{BrickTracer, TraceComparison, TraceDiff, TraceEvent};

// PMAT-112: BrickProfiler for real-time inference telemetry
pub mod profiler;
pub use profiler::{BrickProfiler, OpStats, ProfileReport};

// ============================================================================
// Core Types
// ============================================================================

/// Performance budget expressed in token terms.
/// Aligns compute costs with LLM inference metrics.
///
/// Per Little's Law (1961): throughput = 1 / latency
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct TokenBudget {
    /// Latency budget per token (microseconds)
    pub us_per_token: f64,
    /// Throughput target (tokens/second)
    pub tokens_per_sec: f64,
    /// Batch size for amortization
    pub batch_size: usize,
}

impl TokenBudget {
    /// Create budget from latency target.
    /// 50µs/token = 20,000 tokens/sec
    #[must_use]
    pub fn from_latency(us_per_token: f64) -> Self {
        Self {
            us_per_token,
            tokens_per_sec: 1_000_000.0 / us_per_token,
            batch_size: 1,
        }
    }

    /// Create budget from throughput target.
    /// 20,000 tokens/sec = 50µs/token
    #[must_use]
    pub fn from_throughput(tokens_per_sec: f64) -> Self {
        Self {
            us_per_token: 1_000_000.0 / tokens_per_sec,
            tokens_per_sec,
            batch_size: 1,
        }
    }

    /// Create budget with batch size.
    #[must_use]
    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }

    /// Check if actual performance meets budget.
    #[must_use]
    pub fn is_met(&self, actual_us_per_token: f64) -> bool {
        actual_us_per_token <= self.us_per_token
    }

    /// Calculate gap factor (actual / budget).
    /// < 1.0 means under budget (good), > 1.0 means over budget (bad).
    #[must_use]
    pub fn gap_factor(&self, actual_us_per_token: f64) -> f64 {
        actual_us_per_token / self.us_per_token
    }
}

impl Default for TokenBudget {
    fn default() -> Self {
        Self::from_latency(100.0) // 100µs = 10k tok/s default
    }
}

/// Result of ComputeBrick execution with token metrics.
#[derive(Debug, Clone)]
pub struct TokenResult<T> {
    /// Computed output
    pub output: T,
    /// Number of tokens processed
    pub tokens_processed: usize,
    /// Actual latency (microseconds/token)
    pub us_per_token: f64,
    /// Actual throughput (tokens/second)
    pub tokens_per_sec: f64,
    /// Did we meet the budget?
    pub budget_met: bool,
}

impl<T: Default> Default for TokenResult<T> {
    fn default() -> Self {
        Self {
            output: T::default(),
            tokens_processed: 0,
            us_per_token: 0.0,
            tokens_per_sec: 0.0,
            budget_met: true,
        }
    }
}

impl<T> TokenResult<T> {
    /// Create a new token result.
    pub fn new(output: T, tokens: usize, elapsed_us: f64, budget: &TokenBudget) -> Self {
        let us_per_token = elapsed_us / tokens.max(1) as f64;
        let tokens_per_sec = if us_per_token > 0.0 {
            1_000_000.0 / us_per_token
        } else {
            0.0
        };

        Self {
            output,
            tokens_processed: tokens,
            us_per_token,
            tokens_per_sec,
            budget_met: budget.is_met(us_per_token),
        }
    }
}

/// Errors from ComputeBrick execution.
/// Tells you exactly what failed (Jidoka: stop and signal).
#[derive(Debug)]
pub enum BrickError {
    /// Assertion failed during verification.
    AssertionFailed {
        /// Assertion name
        name: String,
        /// Expected value
        expected: String,
        /// Actual value
        actual: String,
    },
    /// Performance budget exceeded.
    BudgetExceeded {
        /// Budget limit in µs
        limit_us: f64,
        /// Actual time in µs
        actual_us: f64,
    },
    /// Compute operation failed.
    ComputeError(String),
    /// Invalid input.
    InvalidInput(String),
}

impl fmt::Display for BrickError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::AssertionFailed {
                name,
                expected,
                actual,
            } => {
                write!(
                    f,
                    "Assertion failed: {name} - expected {expected}, got {actual}"
                )
            },
            Self::BudgetExceeded {
                limit_us,
                actual_us,
            } => {
                write!(
                    f,
                    "Budget exceeded: {limit_us:.1}µs/tok limit, {actual_us:.1}µs/tok actual"
                )
            },
            Self::ComputeError(msg) => write!(f, "Compute error: {msg}"),
            Self::InvalidInput(msg) => write!(f, "Invalid input: {msg}"),
        }
    }
}

impl std::error::Error for BrickError {}

// ============================================================================
// Brick Assertions
// ============================================================================

/// Falsifiable assertion for a brick (Popper criterion).
#[derive(Debug, Clone)]
pub struct BrickAssertion {
    /// Assertion name (for error messages)
    pub name: String,
    /// Description of what is being asserted
    pub description: String,
    /// Assertion type
    pub kind: AssertionKind,
}

/// Types of assertions a brick can make.
#[derive(Debug, Clone)]
pub enum AssertionKind {
    /// Output matches scalar baseline within tolerance.
    EquivScalar {
        /// Tolerance for comparison
        tolerance: f64,
    },
    /// Output contains no NaN values.
    NoNaN,
    /// Output contains no Inf values.
    NoInf,
    /// Output values within bounds.
    Bounds {
        /// Minimum allowed value
        min: f64,
        /// Maximum allowed value
        max: f64,
    },
    /// Budget must be met.
    BudgetMet,
    /// Custom assertion (returns true if passed).
    Custom {
        /// Name of the custom check
        check_name: String,
    },
}

impl BrickAssertion {
    /// Create assertion that output matches scalar baseline.
    pub fn equiv_scalar(tolerance: f64) -> Self {
        Self {
            name: "equiv_scalar".to_string(),
            description: format!("Output matches scalar baseline within {tolerance}"),
            kind: AssertionKind::EquivScalar { tolerance },
        }
    }

    /// Create assertion that output contains no NaN.
    pub fn no_nan() -> Self {
        Self {
            name: "no_nan".to_string(),
            description: "Output contains no NaN values".to_string(),
            kind: AssertionKind::NoNaN,
        }
    }

    /// Create assertion that output contains no Inf.
    pub fn no_inf() -> Self {
        Self {
            name: "no_inf".to_string(),
            description: "Output contains no Inf values".to_string(),
            kind: AssertionKind::NoInf,
        }
    }

    /// Create assertion that values are within bounds.
    pub fn bounds(min: f64, max: f64) -> Self {
        Self {
            name: "bounds".to_string(),
            description: format!("Output values in [{min}, {max}]"),
            kind: AssertionKind::Bounds { min, max },
        }
    }

    /// Create assertion that budget is met.
    pub fn budget_met() -> Self {
        Self {
            name: "budget_met".to_string(),
            description: "Performance budget is met".to_string(),
            kind: AssertionKind::BudgetMet,
        }
    }

    /// Check assertion against f32 slice output.
    pub fn check_f32(&self, output: &[f32], budget_met: bool) -> Result<(), BrickError> {
        match &self.kind {
            AssertionKind::NoNaN => {
                if let Some(idx) = output.iter().position(|x| x.is_nan()) {
                    return Err(BrickError::AssertionFailed {
                        name: self.name.clone(),
                        expected: "no NaN".to_string(),
                        actual: format!("NaN at index {idx}"),
                    });
                }
            },
            AssertionKind::NoInf => {
                if let Some(idx) = output.iter().position(|x| x.is_infinite()) {
                    return Err(BrickError::AssertionFailed {
                        name: self.name.clone(),
                        expected: "no Inf".to_string(),
                        actual: format!("Inf at index {idx}"),
                    });
                }
            },
            AssertionKind::Bounds { min, max } => {
                for (idx, &val) in output.iter().enumerate() {
                    if (val as f64) < *min || (val as f64) > *max {
                        return Err(BrickError::AssertionFailed {
                            name: self.name.clone(),
                            expected: format!("value in [{min}, {max}]"),
                            actual: format!("value {val} at index {idx}"),
                        });
                    }
                }
            },
            AssertionKind::BudgetMet => {
                if !budget_met {
                    return Err(BrickError::AssertionFailed {
                        name: self.name.clone(),
                        expected: "budget met".to_string(),
                        actual: "budget exceeded".to_string(),
                    });
                }
            },
            AssertionKind::EquivScalar { .. } | AssertionKind::Custom { .. } => {
                // These require external comparison, skip for basic check
            },
        }
        Ok(())
    }
}

// ============================================================================
// Brick Verification
// ============================================================================

/// Result of brick verification.
#[derive(Debug, Clone)]
pub struct BrickVerification {
    /// All assertions passed?
    pub is_valid: bool,
    /// Individual assertion results
    pub results: Vec<(String, bool, String)>,
}

impl BrickVerification {
    /// Create a passing verification.
    pub fn pass() -> Self {
        Self {
            is_valid: true,
            results: vec![],
        }
    }

    /// Create a failing verification.
    pub fn fail(name: &str, reason: &str) -> Self {
        Self {
            is_valid: false,
            results: vec![(name.to_string(), false, reason.to_string())],
        }
    }

    /// Add an assertion result.
    pub fn add(&mut self, name: &str, passed: bool, message: &str) {
        self.results
            .push((name.to_string(), passed, message.to_string()));
        if !passed {
            self.is_valid = false;
        }
    }
}

// ============================================================================
// ComputeBrick Trait
// ============================================================================

/// Core trait for self-verifying, token-centric compute units.
///
/// Every brick must:
/// 1. Have at least one assertion (Popper criterion)
/// 2. Have a non-zero budget (accountability)
/// 3. Be verifiable against baseline
///
/// # Invariants (PROBAR-SPEC-009 §3)
/// - `assertions().len() > 0` (at least one falsifiable claim)
/// - `budget().us_per_token > 0` (performance accountability)
/// - `verify()` checks ALL assertions
pub trait ComputeBrick: Send + Sync {
    /// Output type of this brick.
    type Output;

    /// Brick name for identification.
    fn name(&self) -> &'static str;

    /// Token throughput budget.
    fn budget(&self) -> TokenBudget;

    /// Falsifiable assertions for this brick.
    fn assertions(&self) -> Vec<BrickAssertion>;

    /// Verify brick state without running.
    fn verify(&self) -> BrickVerification {
        let assertions = self.assertions();
        if assertions.is_empty() {
            return BrickVerification::fail(
                self.name(),
                "No assertions defined (Popper violation)",
            );
        }

        let budget = self.budget();
        if budget.us_per_token <= 0.0 {
            return BrickVerification::fail(self.name(), "Zero or negative budget");
        }

        BrickVerification::pass()
    }

    /// Can this brick run? (Jidoka gate)
    fn can_run(&self) -> bool {
        self.verify().is_valid
    }
}

// ============================================================================
// Transformer Brick Implementations
// ============================================================================

/// RMSNorm brick - layer normalization.
#[derive(Debug)]
pub struct RmsNormBrick {
    /// Weight vector
    pub weight: Vec<f32>,
    /// Epsilon for numerical stability
    pub eps: f32,
    /// Budget
    budget: TokenBudget,
}

impl RmsNormBrick {
    /// Create a new RMSNorm brick.
    pub fn new(weight: Vec<f32>, eps: f32) -> Self {
        Self {
            weight,
            eps,
            budget: TokenBudget::from_latency(1.5), // 1.5µs target
        }
    }

    /// Set custom budget.
    #[must_use]
    pub fn with_budget(mut self, budget: TokenBudget) -> Self {
        self.budget = budget;
        self
    }

    /// Run RMSNorm on input.
    pub fn run(&self, input: &[f32]) -> Result<TokenResult<Vec<f32>>, BrickError> {
        if input.len() != self.weight.len() {
            return Err(BrickError::InvalidInput(format!(
                "Input len {} != weight len {}",
                input.len(),
                self.weight.len()
            )));
        }

        let start = Instant::now();

        // Compute RMS
        let rms = (input.iter().map(|x| x * x).sum::<f32>() / input.len() as f32 + self.eps).sqrt();

        // Normalize and scale
        let output: Vec<f32> = input
            .iter()
            .zip(self.weight.iter())
            .map(|(x, w)| (x / rms) * w)
            .collect();

        let elapsed_us = start.elapsed().as_micros() as f64;
        let result = TokenResult::new(output, 1, elapsed_us, &self.budget);

        // Check assertions
        for assertion in self.assertions() {
            assertion.check_f32(&result.output, result.budget_met)?;
        }

        Ok(result)
    }
}

impl ComputeBrick for RmsNormBrick {
    type Output = Vec<f32>;

    fn name(&self) -> &'static str {
        "rms_norm"
    }

    fn budget(&self) -> TokenBudget {
        self.budget
    }

    fn assertions(&self) -> Vec<BrickAssertion> {
        vec![
            BrickAssertion::no_nan(),
            BrickAssertion::no_inf(),
            BrickAssertion::budget_met(),
        ]
    }
}

/// QKV projection brick.
#[derive(Debug)]
pub struct QkvBrick {
    /// Hidden dimension
    pub hidden_dim: usize,
    /// Q output dimension
    pub q_dim: usize,
    /// K output dimension
    pub k_dim: usize,
    /// V output dimension
    pub v_dim: usize,
    /// Budget
    budget: TokenBudget,
    /// Has bias (Qwen2 has large biases)
    pub has_bias: bool,
}

impl QkvBrick {
    /// Create a new QKV brick.
    pub fn new(hidden_dim: usize, q_dim: usize, k_dim: usize, v_dim: usize) -> Self {
        Self {
            hidden_dim,
            q_dim,
            k_dim,
            v_dim,
            budget: TokenBudget::from_latency(6.0), // 6µs target
            has_bias: false,
        }
    }

    /// Set custom budget.
    #[must_use]
    pub fn with_budget(mut self, budget: TokenBudget) -> Self {
        self.budget = budget;
        self
    }

    /// Mark as having bias.
    #[must_use]
    pub fn with_bias(mut self) -> Self {
        self.has_bias = true;
        self
    }

    /// Total output dimension.
    pub fn total_out_dim(&self) -> usize {
        self.q_dim + self.k_dim + self.v_dim
    }
}

impl ComputeBrick for QkvBrick {
    type Output = (Vec<f32>, Vec<f32>, Vec<f32>);

    fn name(&self) -> &'static str {
        "qkv_proj"
    }

    fn budget(&self) -> TokenBudget {
        self.budget
    }

    fn assertions(&self) -> Vec<BrickAssertion> {
        vec![
            BrickAssertion::no_nan(),
            BrickAssertion::no_inf(),
            BrickAssertion::budget_met(),
        ]
    }
}

/// RoPE brick - rotary position embedding.
#[derive(Debug)]
pub struct RopeBrick {
    /// Head dimension
    pub head_dim: usize,
    /// Number of heads
    pub num_heads: usize,
    /// Base theta
    pub theta: f32,
    /// RoPE type (0=NORM, 2=NEOX)
    pub rope_type: u32,
    /// Budget
    budget: TokenBudget,
}

impl RopeBrick {
    /// Create a new RoPE brick.
    pub fn new(head_dim: usize, num_heads: usize, theta: f32, rope_type: u32) -> Self {
        Self {
            head_dim,
            num_heads,
            theta,
            rope_type,
            budget: TokenBudget::from_latency(1.0), // 1µs target
        }
    }

    /// Set custom budget.
    #[must_use]
    pub fn with_budget(mut self, budget: TokenBudget) -> Self {
        self.budget = budget;
        self
    }
}

impl ComputeBrick for RopeBrick {
    type Output = Vec<f32>;

    fn name(&self) -> &'static str {
        "rope"
    }

    fn budget(&self) -> TokenBudget {
        self.budget
    }

    fn assertions(&self) -> Vec<BrickAssertion> {
        vec![
            BrickAssertion::no_nan(),
            BrickAssertion::no_inf(),
            BrickAssertion::budget_met(),
        ]
    }
}

/// Attention brick.
#[derive(Debug)]
pub struct AttentionBrick {
    /// Number of query heads
    pub num_heads: usize,
    /// Number of KV heads (for GQA)
    pub num_kv_heads: usize,
    /// Head dimension
    pub head_dim: usize,
    /// Budget
    budget: TokenBudget,
}

impl AttentionBrick {
    /// Create a new attention brick.
    pub fn new(num_heads: usize, num_kv_heads: usize, head_dim: usize) -> Self {
        Self {
            num_heads,
            num_kv_heads,
            head_dim,
            budget: TokenBudget::from_latency(10.0), // 10µs target
        }
    }

    /// Set custom budget.
    #[must_use]
    pub fn with_budget(mut self, budget: TokenBudget) -> Self {
        self.budget = budget;
        self
    }

    /// GQA group size.
    pub fn group_size(&self) -> usize {
        self.num_heads / self.num_kv_heads.max(1)
    }
}

impl ComputeBrick for AttentionBrick {
    type Output = Vec<f32>;

    fn name(&self) -> &'static str {
        "attention"
    }

    fn budget(&self) -> TokenBudget {
        self.budget
    }

    fn assertions(&self) -> Vec<BrickAssertion> {
        vec![
            BrickAssertion::no_nan(),
            BrickAssertion::no_inf(),
            BrickAssertion::budget_met(),
            // Attention outputs should be bounded
            BrickAssertion::bounds(-100.0, 100.0),
        ]
    }
}

/// FlashAttentionBrick - incremental flash attention for decode (P1 optimization).
///
/// **Algorithm** (FlashAttention-2, Dao et al. 2023):
/// ```text
/// For decode (single query token):
///   Q: [1, H, D]     (single query)
///   K: [S, H_kv, D]  (full KV cache)
///   V: [S, H_kv, D]  (full KV cache)
///
///   Online softmax (no full attention matrix materialization):
///   for tile in KV_tiles(TILE_SIZE=128):
///       S_tile = Q @ K_tile^T / sqrt(D)    # [1, H, TILE_SIZE]
///       m_new = max(m_old, max(S_tile))    # Running max
///       P_tile = exp(S_tile - m_new)       # Stable softmax numerator
///       O = O * exp(m_old - m_new) + P_tile @ V_tile  # Accumulate
///       l = l * exp(m_old - m_new) + sum(P_tile)      # Running denominator
///   O = O / l  # Final output
/// ```
///
/// **Performance vs naive**:
/// - Naive: O(S) memory for attention matrix
/// - Flash: O(TILE_SIZE) memory, 2x speedup from better cache locality
///
/// **Reference**: Dao, T., et al. (2023). "FlashAttention-2: Faster Attention
/// with Better Parallelism and Work Partitioning." arXiv:2307.08691.
#[derive(Debug, Clone)]
pub struct FlashAttentionBrick {
    /// Number of query heads
    pub num_heads: usize,
    /// Number of KV heads (for GQA)
    pub num_kv_heads: usize,
    /// Head dimension
    pub head_dim: usize,
    /// Tile size for KV cache (default: 128 for L2 cache fit)
    pub tile_size: usize,
    /// Budget (target: 5.0µs for 2x improvement over naive)
    budget: TokenBudget,
    /// Use online softmax (FlashAttention algorithm)
    pub use_online_softmax: bool,
}

impl FlashAttentionBrick {
    /// Create new flash attention brick with default tile size.
    #[must_use]
    pub fn new(num_heads: usize, num_kv_heads: usize, head_dim: usize) -> Self {
        Self {
            num_heads,
            num_kv_heads,
            head_dim,
            tile_size: 128, // Optimal for L2 cache on most GPUs
            budget: TokenBudget::from_latency(5.0), // 5µs target (2x vs 10µs naive)
            use_online_softmax: true,
        }
    }

    /// Create with custom tile size (for tuning).
    #[must_use]
    pub fn with_tile_size(
        num_heads: usize,
        num_kv_heads: usize,
        head_dim: usize,
        tile_size: usize,
    ) -> Self {
        Self {
            num_heads,
            num_kv_heads,
            head_dim,
            tile_size,
            budget: TokenBudget::from_latency(5.0),
            use_online_softmax: true,
        }
    }

    /// Set custom budget.
    #[must_use]
    pub fn with_budget(mut self, budget: TokenBudget) -> Self {
        self.budget = budget;
        self
    }

    /// GQA group size (query heads per KV head).
    #[must_use]
    pub fn group_size(&self) -> usize {
        self.num_heads / self.num_kv_heads.max(1)
    }

    /// Compute FLOPs per token for attention at given sequence length.
    ///
    /// FLOPs = 2 * H * D * S (Q @ K^T) + 2 * H * S * D (attn @ V)
    ///       = 4 * H * D * S
    #[must_use]
    pub fn flops(&self, seq_len: usize) -> u64 {
        4 * self.num_heads as u64 * self.head_dim as u64 * seq_len as u64
    }

    /// Compute memory bandwidth for naive vs flash attention.
    ///
    /// Naive: Reads full KV cache for each head
    /// Flash: Reads KV cache in tiles, better cache reuse
    #[must_use]
    pub fn memory_bytes(&self, seq_len: usize) -> (u64, u64) {
        let kv_bytes = 2 * self.num_kv_heads as u64 * self.head_dim as u64 * seq_len as u64 * 4; // K + V, f32
        let naive = kv_bytes + self.num_heads as u64 * seq_len as u64 * 4; // + attention matrix
        let flash = kv_bytes; // No attention matrix materialized
        (naive, flash)
    }

    /// Compute arithmetic intensity (FLOPs / bytes).
    #[must_use]
    pub fn arithmetic_intensity(&self, seq_len: usize) -> f64 {
        let (_, flash_bytes) = self.memory_bytes(seq_len);
        self.flops(seq_len) as f64 / flash_bytes as f64
    }

    /// Number of tiles needed for given sequence length.
    #[must_use]
    pub fn num_tiles(&self, seq_len: usize) -> usize {
        seq_len.div_ceil(self.tile_size)
    }

    /// Compute flash attention with online softmax algorithm.
    ///
    /// **REAL IMPLEMENTATION** - FlashAttention-2 (Dao et al. 2023)
    ///
    /// # Arguments
    /// * `query` - Query tensor [num_heads, head_dim]
    /// * `keys` - Key cache [seq_len, num_kv_heads, head_dim]
    /// * `values` - Value cache [seq_len, num_kv_heads, head_dim]
    ///
    /// # Returns
    /// * Output tensor [num_heads, head_dim]
    pub fn forward(
        &self,
        query: &[f32],  // [num_heads * head_dim]
        keys: &[f32],   // [seq_len * num_kv_heads * head_dim]
        values: &[f32], // [seq_len * num_kv_heads * head_dim]
        seq_len: usize,
    ) -> Result<Vec<f32>, BrickError> {
        if self.num_heads == 0 || self.head_dim == 0 {
            return Err(BrickError::InvalidInput("Zero dimension".to_string()));
        }
        if query.len() != self.num_heads * self.head_dim {
            return Err(BrickError::InvalidInput(format!(
                "Query length {} != num_heads * head_dim = {}",
                query.len(),
                self.num_heads * self.head_dim
            )));
        }
        let expected_kv_len = seq_len * self.num_kv_heads * self.head_dim;
        if keys.len() != expected_kv_len || values.len() != expected_kv_len {
            return Err(BrickError::InvalidInput(format!(
                "KV length {} != seq_len * num_kv_heads * head_dim = {}",
                keys.len(),
                expected_kv_len
            )));
        }

        let scale = 1.0 / (self.head_dim as f32).sqrt();
        let group_size = self.group_size();
        let mut output = vec![0.0f32; self.num_heads * self.head_dim];

        // Process each query head
        for h in 0..self.num_heads {
            let kv_head = h / group_size; // GQA: map query head to KV head
            let q_start = h * self.head_dim;

            // Online softmax variables (FlashAttention-2)
            let mut m = f32::NEG_INFINITY; // Running max
            let mut l = 0.0f32; // Running sum of exp
            let mut o = vec![0.0f32; self.head_dim]; // Running output

            // Process in tiles for cache efficiency
            for tile_start in (0..seq_len).step_by(self.tile_size) {
                let tile_end = (tile_start + self.tile_size).min(seq_len);

                for s in tile_start..tile_end {
                    // Compute Q @ K^T for this position
                    let k_start = (s * self.num_kv_heads + kv_head) * self.head_dim;
                    let mut score = 0.0f32;
                    for d in 0..self.head_dim {
                        score += query[q_start + d] * keys[k_start + d];
                    }
                    score *= scale;

                    // Online softmax update (Milakov & Gimelshein, 2018)
                    let m_new = m.max(score);
                    let exp_old = (m - m_new).exp();
                    let exp_score = (score - m_new).exp();

                    // Update running sum
                    l = l * exp_old + exp_score;

                    // Update running output: O = O * exp(m_old - m_new) + exp(score - m_new) * V
                    let v_start = (s * self.num_kv_heads + kv_head) * self.head_dim;
                    for d in 0..self.head_dim {
                        o[d] = o[d] * exp_old + exp_score * values[v_start + d];
                    }

                    m = m_new;
                }
            }

            // Normalize output: O = O / l
            if l > 0.0 {
                for d in 0..self.head_dim {
                    output[h * self.head_dim + d] = o[d] / l;
                }
            }
        }

        Ok(output)
    }

    /// Execute flash attention with timing (for benchmarking).
    pub fn forward_timed(
        &self,
        query: &[f32],
        keys: &[f32],
        values: &[f32],
        seq_len: usize,
    ) -> Result<TokenResult<Vec<f32>>, BrickError> {
        let start = Instant::now();
        let output = self.forward(query, keys, values, seq_len)?;
        let elapsed_us = start.elapsed().as_secs_f64() * 1_000_000.0;

        Ok(TokenResult {
            output,
            tokens_processed: 1,
            us_per_token: elapsed_us,
            tokens_per_sec: 1_000_000.0 / elapsed_us,
            budget_met: elapsed_us <= self.budget.us_per_token,
        })
    }

    /// Legacy stub for backward compatibility (prefer `forward()`)
    #[deprecated(note = "Use forward() for real implementation")]
    pub fn execute(&self, _seq_len: usize) -> Result<Vec<f32>, BrickError> {
        if self.num_heads == 0 || self.head_dim == 0 {
            return Err(BrickError::InvalidInput("Zero dimension".to_string()));
        }
        Ok(vec![0.0; self.num_heads * self.head_dim])
    }
}

impl ComputeBrick for FlashAttentionBrick {
    type Output = Vec<f32>;

    fn name(&self) -> &'static str {
        "flash_attention"
    }

    fn budget(&self) -> TokenBudget {
        self.budget
    }

    fn assertions(&self) -> Vec<BrickAssertion> {
        vec![
            BrickAssertion::no_nan(),
            BrickAssertion::no_inf(),
            BrickAssertion::budget_met(),
            // Attention outputs should be bounded
            BrickAssertion::bounds(-100.0, 100.0),
            // Custom assertions for flash attention
            BrickAssertion {
                name: "online_softmax".to_string(),
                description: "Uses online softmax (no full attention matrix)".to_string(),
                kind: AssertionKind::Custom {
                    check_name: "online_softmax".to_string(),
                },
            },
            BrickAssertion {
                name: "tiled_kv_access".to_string(),
                description: "KV cache accessed in tiles for cache locality".to_string(),
                kind: AssertionKind::Custom {
                    check_name: "tiled_kv_access".to_string(),
                },
            },
        ]
    }

    fn can_run(&self) -> bool {
        self.num_heads > 0 && self.head_dim > 0 && self.tile_size > 0
    }
}

/// FFN brick (SwiGLU).
#[derive(Debug)]
pub struct FfnBrick {
    /// Hidden dimension
    pub hidden_dim: usize,
    /// Intermediate dimension
    pub intermediate_dim: usize,
    /// Budget
    budget: TokenBudget,
}

impl FfnBrick {
    /// Create a new FFN brick.
    pub fn new(hidden_dim: usize, intermediate_dim: usize) -> Self {
        Self {
            hidden_dim,
            intermediate_dim,
            budget: TokenBudget::from_latency(12.2), // 12.2µs target
        }
    }

    /// Set custom budget.
    #[must_use]
    pub fn with_budget(mut self, budget: TokenBudget) -> Self {
        self.budget = budget;
        self
    }
}

impl ComputeBrick for FfnBrick {
    type Output = Vec<f32>;

    fn name(&self) -> &'static str {
        "ffn"
    }

    fn budget(&self) -> TokenBudget {
        self.budget
    }

    fn assertions(&self) -> Vec<BrickAssertion> {
        vec![
            BrickAssertion::no_nan(),
            BrickAssertion::no_inf(),
            BrickAssertion::budget_met(),
        ]
    }
}

/// Output projection brick.
#[derive(Debug)]
pub struct OProjBrick {
    /// Input dimension (num_heads * head_dim)
    pub in_dim: usize,
    /// Output dimension (hidden_dim)
    pub out_dim: usize,
    /// Budget
    budget: TokenBudget,
}

impl OProjBrick {
    /// Create a new O projection brick.
    pub fn new(in_dim: usize, out_dim: usize) -> Self {
        Self {
            in_dim,
            out_dim,
            budget: TokenBudget::from_latency(3.5), // 3.5µs target
        }
    }

    /// Set custom budget.
    #[must_use]
    pub fn with_budget(mut self, budget: TokenBudget) -> Self {
        self.budget = budget;
        self
    }
}

impl ComputeBrick for OProjBrick {
    type Output = Vec<f32>;

    fn name(&self) -> &'static str {
        "o_proj"
    }

    fn budget(&self) -> TokenBudget {
        self.budget
    }

    fn assertions(&self) -> Vec<BrickAssertion> {
        vec![
            BrickAssertion::no_nan(),
            BrickAssertion::no_inf(),
            BrickAssertion::budget_met(),
        ]
    }
}

/// ActivationQuantBrick - Q8 activation quantization for memory bandwidth reduction (P2).
///
/// **Purpose**: Quantize intermediate activations from f32 to int8 to reduce
/// memory bandwidth by 4x during inter-layer communication.
///
/// **Pipeline**:
/// ```text
/// Layer N output (f32) → Q8 quantize → transfer → Q8 dequantize → Layer N+1 input (f32)
///
/// Memory bandwidth reduction:
///   f32: 4 bytes/element
///   int8: 1 byte/element + 2 floats (scale, zero_point)
///   Effective: ~4x reduction for large activations
/// ```
///
/// **Algorithm** (per-tensor affine quantization):
/// ```text
/// Quantize:
///   scale = (max - min) / 255
///   zero_point = round(-min / scale)
///   q[i] = clamp(round(x[i] / scale + zero_point), 0, 255)
///
/// Dequantize:
///   x[i] = (q[i] - zero_point) * scale
/// ```
///
/// **Performance**: 2x memory BW improvement with ~0.1% accuracy loss
///
/// **Reference**: Jacob, B., et al. (2018). "Quantization and Training of
/// Neural Networks for Efficient Integer-Arithmetic-Only Inference." CVPR '18.
#[derive(Debug, Clone)]
pub struct ActivationQuantBrick {
    /// Activation dimension (e.g., hidden_dim or intermediate_dim)
    pub dim: usize,
    /// Budget (target: 0.5µs for quant+dequant overhead)
    budget: TokenBudget,
    /// Use per-channel quantization (more accurate but slower)
    pub per_channel: bool,
}

impl ActivationQuantBrick {
    /// Create new activation quantization brick.
    #[must_use]
    pub fn new(dim: usize) -> Self {
        Self {
            dim,
            budget: TokenBudget::from_latency(0.5), // 0.5µs overhead target
            per_channel: false,
        }
    }

    /// Create with per-channel quantization.
    #[must_use]
    pub fn with_per_channel(dim: usize) -> Self {
        Self {
            dim,
            budget: TokenBudget::from_latency(1.0), // 1.0µs for per-channel
            per_channel: true,
        }
    }

    /// Set custom budget.
    #[must_use]
    pub fn with_budget(mut self, budget: TokenBudget) -> Self {
        self.budget = budget;
        self
    }

    /// Compute memory bandwidth reduction factor.
    ///
    /// f32 (4 bytes) → int8 (1 byte) + scale/zero_point = ~4x reduction
    #[must_use]
    pub fn bandwidth_reduction(&self) -> f64 {
        // Original: dim * 4 bytes (f32)
        // Quantized: dim * 1 byte (int8) + 8 bytes (scale + zero_point)
        let original_bytes = self.dim * 4;
        let quantized_bytes = self.dim + 8; // +8 for scale and zero_point (f32 each)
        original_bytes as f64 / quantized_bytes as f64
    }

    /// Compute quantization error estimate (typical for 8-bit).
    ///
    /// Per Jacob et al. 2018, typical Q8 error is ~0.1% for activations.
    #[must_use]
    pub fn estimated_error(&self) -> f64 {
        if self.per_channel {
            0.0005 // 0.05% for per-channel
        } else {
            0.001 // 0.1% for per-tensor
        }
    }

    /// Compute bytes saved per token.
    #[must_use]
    pub fn bytes_saved(&self) -> usize {
        // f32 (4 bytes) → int8 (1 byte) = 3 bytes saved per element
        self.dim * 3
    }

    /// Quantize f32 activations to int8 using Q8_0 block format.
    ///
    /// **REAL IMPLEMENTATION** - Not a stub.
    /// Uses symmetric quantization: scale = max(abs(values)) / 127.0
    ///
    /// # Arguments
    /// * `input` - f32 activations to quantize (must be length == self.dim)
    ///
    /// # Returns
    /// * Quantized int8 values and scale factors
    ///
    /// # Example
    /// ```ignore
    /// let brick = ActivationQuantBrick::new(64);
    /// let input = vec![1.0f32; 64];
    /// let (quants, scales) = brick.quantize(&input)?;
    /// assert_eq!(quants.len(), 64);
    /// ```
    pub fn quantize(&self, input: &[f32]) -> Result<(Vec<i8>, Vec<f32>), BrickError> {
        if input.len() != self.dim {
            return Err(BrickError::InvalidInput(format!(
                "Input length {} != dim {}",
                input.len(),
                self.dim
            )));
        }
        if self.dim == 0 {
            return Err(BrickError::InvalidInput("Zero dimension".to_string()));
        }

        // Quantize in blocks of 32 (Q8_0 block size)
        let num_blocks = self.dim.div_ceil(32);
        let mut quants = Vec::with_capacity(self.dim);
        let mut scales = Vec::with_capacity(num_blocks);

        for block_idx in 0..num_blocks {
            let start = block_idx * 32;
            let end = (start + 32).min(self.dim);

            // Pad to 32 if needed
            let mut block_data = [0.0f32; 32];
            for (i, &v) in input[start..end].iter().enumerate() {
                block_data[i] = v;
            }

            let block = Q8_0Block::quantize(&block_data);
            scales.push(block.scale);

            // Only take the actual values (not padding)
            for &q in &block.quants[0..(end - start)] {
                quants.push(q);
            }
        }

        Ok((quants, scales))
    }

    /// Dequantize int8 back to f32 using stored scales.
    ///
    /// **REAL IMPLEMENTATION** - Not a stub.
    pub fn dequantize(&self, quants: &[i8], scales: &[f32]) -> Result<Vec<f32>, BrickError> {
        if quants.len() != self.dim {
            return Err(BrickError::InvalidInput(format!(
                "Quants length {} != dim {}",
                quants.len(),
                self.dim
            )));
        }

        let mut output = Vec::with_capacity(self.dim);
        for (block_idx, &scale) in scales.iter().enumerate() {
            let start = block_idx * 32;
            let end = (start + 32).min(self.dim);
            for &q in &quants[start..end] {
                output.push(q as f32 * scale);
            }
        }

        Ok(output)
    }

    /// Compute quantization error vs original input.
    ///
    /// **REAL IMPLEMENTATION** - Measures actual error, not estimates.
    pub fn measure_error(
        &self,
        original: &[f32],
        quants: &[i8],
        scales: &[f32],
    ) -> Result<f64, BrickError> {
        let dequantized = self.dequantize(quants, scales)?;

        let max_error = original
            .iter()
            .zip(dequantized.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);

        let max_val = original.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
        if max_val < 1e-10 {
            return Ok(0.0);
        }

        Ok((max_error / max_val) as f64)
    }

    /// Execute quantization with timing (for benchmarking).
    #[allow(clippy::type_complexity)]
    pub fn execute_timed(
        &self,
        input: &[f32],
    ) -> Result<TokenResult<(Vec<i8>, Vec<f32>)>, BrickError> {
        let start = Instant::now();
        let (quants, scales) = self.quantize(input)?;
        let elapsed_us = start.elapsed().as_secs_f64() * 1_000_000.0;

        Ok(TokenResult {
            output: (quants, scales),
            tokens_processed: 1,
            us_per_token: elapsed_us,
            tokens_per_sec: 1_000_000.0 / elapsed_us,
            budget_met: elapsed_us <= self.budget.us_per_token,
        })
    }

    /// Legacy stub for backward compatibility (prefer `quantize()`)
    #[deprecated(note = "Use quantize() for real implementation")]
    pub fn execute(&self) -> Result<Vec<u8>, BrickError> {
        if self.dim == 0 {
            return Err(BrickError::InvalidInput("Zero dimension".to_string()));
        }
        // Return zeros for backward compat - use quantize() for real output
        Ok(vec![128u8; self.dim])
    }
}

impl ComputeBrick for ActivationQuantBrick {
    type Output = Vec<u8>;

    fn name(&self) -> &'static str {
        "activation_quant"
    }

    fn budget(&self) -> TokenBudget {
        self.budget
    }

    fn assertions(&self) -> Vec<BrickAssertion> {
        vec![
            BrickAssertion::budget_met(),
            BrickAssertion {
                name: "symmetric_range".to_string(),
                description: "Q8 values centered around 128 (zero_point)".to_string(),
                kind: AssertionKind::Custom {
                    check_name: "symmetric_range".to_string(),
                },
            },
            BrickAssertion {
                name: "error_bound".to_string(),
                description: "Quantization error < 0.1% (per-tensor) or 0.05% (per-channel)"
                    .to_string(),
                kind: AssertionKind::Custom {
                    check_name: "error_bound".to_string(),
                },
            },
        ]
    }

    fn can_run(&self) -> bool {
        self.dim > 0
    }
}

// ============================================================================
// Transformer Layer Brick
// ============================================================================

/// Full transformer layer as a composed brick.
#[derive(Debug)]
pub struct TransformerLayerBrick {
    /// Layer index
    pub layer_idx: usize,
    /// Attention layer normalization brick
    pub attn_norm: RmsNormBrick,
    /// QKV projection brick
    pub qkv: QkvBrick,
    /// Rotary position embedding brick
    pub rope: RopeBrick,
    /// Attention computation brick
    pub attention: AttentionBrick,
    /// Output projection brick
    pub o_proj: OProjBrick,
    /// FFN layer normalization brick
    pub ffn_norm: RmsNormBrick,
    /// Feed-forward network brick
    pub ffn: FfnBrick,
    /// Timing metrics (updated after each run)
    pub last_timing: Option<LayerTiming>,
}

/// Timing breakdown for a layer.
#[derive(Debug, Clone, Default)]
pub struct LayerTiming {
    /// Attention normalization time (µs)
    pub attn_norm_us: f64,
    /// QKV projection time (µs)
    pub qkv_us: f64,
    /// RoPE application time (µs)
    pub rope_us: f64,
    /// Attention computation time (µs)
    pub attention_us: f64,
    /// Output projection time (µs)
    pub o_proj_us: f64,
    /// FFN normalization time (µs)
    pub ffn_norm_us: f64,
    /// FFN computation time (µs)
    pub ffn_us: f64,
    /// Total layer time (µs)
    pub total_us: f64,
}

impl LayerTiming {
    /// Find the bottleneck brick.
    pub fn bottleneck(&self) -> (&'static str, f64) {
        let bricks = [
            ("attn_norm", self.attn_norm_us),
            ("qkv", self.qkv_us),
            ("rope", self.rope_us),
            ("attention", self.attention_us),
            ("o_proj", self.o_proj_us),
            ("ffn_norm", self.ffn_norm_us),
            ("ffn", self.ffn_us),
        ];

        bricks
            .into_iter()
            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
            .unwrap_or(("unknown", 0.0))
    }
}

impl TransformerLayerBrick {
    /// Create from configuration.
    pub fn from_config(
        layer_idx: usize,
        hidden_dim: usize,
        num_heads: usize,
        num_kv_heads: usize,
        intermediate_dim: usize,
        eps: f32,
        rope_theta: f32,
        rope_type: u32,
    ) -> Self {
        let head_dim = hidden_dim / num_heads;
        let q_dim = num_heads * head_dim;
        let kv_dim = num_kv_heads * head_dim;

        Self {
            layer_idx,
            attn_norm: RmsNormBrick::new(vec![1.0; hidden_dim], eps),
            qkv: QkvBrick::new(hidden_dim, q_dim, kv_dim, kv_dim),
            rope: RopeBrick::new(head_dim, num_heads, rope_theta, rope_type),
            attention: AttentionBrick::new(num_heads, num_kv_heads, head_dim),
            o_proj: OProjBrick::new(q_dim, hidden_dim),
            ffn_norm: RmsNormBrick::new(vec![1.0; hidden_dim], eps),
            ffn: FfnBrick::new(hidden_dim, intermediate_dim),
            last_timing: None,
        }
    }

    /// Get total budget for this layer.
    pub fn total_budget_us(&self) -> f64 {
        self.attn_norm.budget().us_per_token
            + self.qkv.budget().us_per_token
            + self.rope.budget().us_per_token
            + self.attention.budget().us_per_token
            + self.o_proj.budget().us_per_token
            + self.ffn_norm.budget().us_per_token
            + self.ffn.budget().us_per_token
    }
}

impl ComputeBrick for TransformerLayerBrick {
    type Output = Vec<f32>;

    fn name(&self) -> &'static str {
        "transformer_layer"
    }

    fn budget(&self) -> TokenBudget {
        TokenBudget::from_latency(self.total_budget_us())
    }

    fn assertions(&self) -> Vec<BrickAssertion> {
        vec![
            BrickAssertion::no_nan(),
            BrickAssertion::no_inf(),
            BrickAssertion::budget_met(),
        ]
    }

    fn verify(&self) -> BrickVerification {
        // Verify all component bricks
        let mut result = BrickVerification::pass();

        for brick in [
            &self.attn_norm as &dyn ComputeBrick<Output = Vec<f32>>,
            &self.ffn_norm as &dyn ComputeBrick<Output = Vec<f32>>,
        ] {
            let v = brick.verify();
            if !v.is_valid {
                result.is_valid = false;
                result.results.extend(v.results);
            }
        }

        result
    }
}

// ============================================================================
// Bottleneck Report
// ============================================================================

/// Report identifying pipeline bottleneck.
#[derive(Debug, Clone)]
pub struct BottleneckReport {
    /// Layer index containing bottleneck
    pub layer_idx: usize,
    /// Brick name
    pub brick_name: &'static str,
    /// Actual latency (µs)
    pub actual_us: f64,
    /// Budget latency (µs)
    pub budget_us: f64,
    /// Gap factor (actual / budget)
    pub gap_factor: f64,
}

impl fmt::Display for BottleneckReport {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "Bottleneck: {} (layer {}) - {:.1}µs actual vs {:.1}µs budget ({:.2}x)",
            self.brick_name, self.layer_idx, self.actual_us, self.budget_us, self.gap_factor
        )
    }
}

// ============================================================================
// Benchmark Brick
// ============================================================================

/// Configuration for benchmark runs.
#[derive(Debug, Clone)]
pub struct BenchmarkConfig {
    /// Number of warmup iterations
    pub warmup: usize,
    /// Number of sample iterations
    pub samples: usize,
    /// Maximum allowed CV (coefficient of variation)
    pub max_cv: f64,
}

impl Default for BenchmarkConfig {
    fn default() -> Self {
        Self {
            warmup: 10,
            samples: 100,
            max_cv: 0.05, // 5% per Stabilizer (Curtsinger & Berger 2013)
        }
    }
}

/// Benchmark report with statistical analysis.
#[derive(Debug, Clone)]
pub struct BenchmarkReport {
    /// Brick name
    pub brick_name: String,
    /// Mean latency (µs)
    pub mean_us: f64,
    /// Standard deviation (µs)
    pub std_us: f64,
    /// Coefficient of variation
    pub cv: f64,
    /// 50th percentile (µs)
    pub p50_us: f64,
    /// 99th percentile (µs)
    pub p99_us: f64,
    /// Throughput (tokens/sec)
    pub tokens_per_sec: f64,
    /// Budget target (µs)
    pub budget_us: f64,
    /// Budget met?
    pub budget_met: bool,
    /// Statistical validity (CV < max_cv)
    pub statistically_valid: bool,
}

impl fmt::Display for BenchmarkReport {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let status = if self.budget_met { "PASS" } else { "FAIL" };
        write!(
            f,
            "{}: {:.1}µs ± {:.1}µs (CV={:.1}%) | {:.0} tok/s | budget: {} ({})",
            self.brick_name,
            self.mean_us,
            self.std_us,
            self.cv * 100.0,
            self.tokens_per_sec,
            self.budget_us,
            status
        )
    }
}

/// Calculate percentile from sorted samples.
fn percentile(samples: &[f64], p: f64) -> f64 {
    if samples.is_empty() {
        return 0.0;
    }
    let idx = ((samples.len() as f64) * p).floor() as usize;
    samples[idx.min(samples.len() - 1)]
}

/// Run benchmark on a brick with statistical rigor.
pub fn benchmark_brick<B: ComputeBrick>(
    brick: &B,
    run_fn: impl Fn() -> f64,
    config: &BenchmarkConfig,
) -> BenchmarkReport {
    // Warmup (Jidoka: ensure stable state)
    for _ in 0..config.warmup {
        let _ = run_fn();
    }

    // Collect samples
    let mut samples: Vec<f64> = Vec::with_capacity(config.samples);
    for _ in 0..config.samples {
        samples.push(run_fn());
    }

    // Sort for percentiles
    samples.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

    // Statistical analysis
    let mean = samples.iter().sum::<f64>() / samples.len() as f64;
    let std =
        (samples.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / samples.len() as f64).sqrt();
    let cv = std / mean;

    let budget = brick.budget();

    BenchmarkReport {
        brick_name: brick.name().to_string(),
        mean_us: mean,
        std_us: std,
        cv,
        p50_us: percentile(&samples, 0.50),
        p99_us: percentile(&samples, 0.99),
        tokens_per_sec: 1_000_000.0 / mean,
        budget_us: budget.us_per_token,
        budget_met: mean <= budget.us_per_token,
        statistically_valid: cv <= config.max_cv,
    }
}

// ============================================================================
// CUDA Graph Brick (Section 5.2 - P0)
// ============================================================================

/// CUDA Graph Brick for eliminating kernel launch overhead.
///
/// Per spec: docs/specifications/qwen2.5-coder-showcase-demo.md §5.2
///
/// Uses CUDA graph capture to reduce ~280 kernel launches to single graph replay.
/// Expected impact: 5.6ms overhead → 0.02ms = 280x overhead reduction.
///
/// # Implementation
///
/// Wraps `CudaExecutor::decode_graph` and `try_graph_capture()` from cuda.rs.
/// Uses indirect kernels (KvCacheScatterIndirect, RopeIndirect) for graph compatibility.
#[derive(Debug, Clone)]
pub struct CudaGraphBrick {
    /// Number of layers captured in graph
    pub num_layers: usize,
    /// Hidden dimension
    pub hidden_dim: usize,
    /// Whether graph is currently captured
    pub captured: bool,
    /// Token budget (target: 10µs launch overhead vs 5600µs eager)
    budget: TokenBudget,
}

impl CudaGraphBrick {
    /// Create new CUDA Graph brick for model configuration.
    #[must_use]
    pub fn new(num_layers: usize, hidden_dim: usize) -> Self {
        // Graph overhead should be < 100µs (vs ~5.6ms for 280 launches)
        let budget_us = 20.0; // Conservative: 20µs for graph replay
        Self {
            num_layers,
            hidden_dim,
            captured: false,
            budget: TokenBudget::from_latency(budget_us),
        }
    }

    /// Set custom budget.
    #[must_use]
    pub fn with_budget(mut self, budget: TokenBudget) -> Self {
        self.budget = budget;
        self
    }

    /// Mark graph as captured.
    pub fn set_captured(&mut self, captured: bool) {
        self.captured = captured;
    }

    /// Check if graph can be used (captured and valid).
    #[must_use]
    pub fn can_replay(&self) -> bool {
        self.captured
    }

    /// Replay the captured graph (stub - actual execution via CudaExecutor).
    pub fn replay(&self) -> Result<(), BrickError> {
        if !self.captured {
            return Err(BrickError::ComputeError(
                "CUDA graph not captured yet".to_string(),
            ));
        }
        // Actual replay would be done via CudaExecutor::forward_graphed()
        Ok(())
    }
}

impl ComputeBrick for CudaGraphBrick {
    type Output = ();

    fn name(&self) -> &'static str {
        "cuda_graph"
    }

    fn budget(&self) -> TokenBudget {
        self.budget
    }

    fn assertions(&self) -> Vec<BrickAssertion> {
        vec![
            BrickAssertion::budget_met(),
            BrickAssertion {
                name: "graph_speedup".to_string(),
                description: "Graph replay faster than eager execution".to_string(),
                kind: AssertionKind::Custom {
                    check_name: "graph_speedup".to_string(),
                },
            },
        ]
    }

    fn can_run(&self) -> bool {
        self.num_layers > 0 && self.hidden_dim > 0
    }
}

// ============================================================================
// Tests (F001-F020)
// ============================================================================

// Tests extracted to tests.rs (PMAT-802)
#[cfg(test)]
#[path = "tests.rs"]
mod brick_tests;

// Additional tests in tests_part_02.rs
#[cfg(test)]
#[path = "tests_part_02.rs"]
mod brick_tests_part_02;

// Additional tests in tests_part_03.rs
#[cfg(test)]
#[path = "tests_part_03.rs"]
mod brick_tests_part_03;

// BrickProfiler tests (PMAT-112)
#[cfg(test)]
#[path = "profiler_tests.rs"]
mod profiler_tests;

// Fused ops tests (cuda feature)
#[cfg(all(test, feature = "cuda"))]
mod fused_tests;
