//! Parallel Scan Algorithms for Mamba-2 State Space Models
//!
//! Implementation of efficient parallel scan algorithms for computing
//! state space model sequences with linear time complexity.
//!
//! All scan operations run on GPU using `StreamTensor` ops --
//! no bulk host downloads in production paths.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;

use cudarc::driver::CudaStream;
use tracing::{debug, instrument};

use crate::gpu_tensor::{
    gpu_add, gpu_mul, gpu_narrow_2d, gpu_cat_dim0, gpu_select_dim1,
    GpuTensor,
};
use crate::liquid::FixedPoint;
use ml_core::MLError;

/// Scan operations for parallel prefix scan
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ScanOperator {
    Add,
    Mul,
    Max,
    Min,
    SSMScan,
}

/// Performance benchmark result
#[derive(Debug, Clone)]
pub struct ScanBenchmark {
    pub sequence_length: usize,
    pub duration_nanos: u64,
    pub throughput_elements_per_sec: FixedPoint,
    pub memory_bandwidth_gb_per_sec: FixedPoint,
    pub cache_efficiency: FixedPoint,
}

/// Parallel scan engine with hardware optimizations
#[derive(Debug)]
pub struct ParallelScanEngine {
    stream: Arc<CudaStream>,
    pub block_size: usize,
    parallel_threshold: usize,
    scan_operations: std::sync::atomic::AtomicU64,
    total_latency_ns: std::sync::atomic::AtomicU64,
    memory_transfers: std::sync::atomic::AtomicU64,
}

impl ParallelScanEngine {
    /// Create new parallel scan engine
    #[allow(clippy::missing_const_for_fn)]
    pub fn new(stream: Arc<CudaStream>, parallel_threshold: usize) -> Self {
        Self {
            stream,
            block_size: 1024,
            parallel_threshold,
            scan_operations: std::sync::atomic::AtomicU64::new(0),
            total_latency_ns: std::sync::atomic::AtomicU64::new(0),
            memory_transfers: std::sync::atomic::AtomicU64::new(0),
        }
    }

    /// Parallel prefix scan - main entry point
    #[instrument(skip(self, input))]
    pub fn parallel_prefix_scan(
        &self,
        input: &GpuTensor,
        op: ScanOperator,
    ) -> Result<GpuTensor, MLError> {
        let start = Instant::now();

        // Input: [batch, seq_len] or [batch, seq_len, features]
        if input.shape.len() < 2 {
            return Err(MLError::InvalidInput(format!(
                "parallel_prefix_scan requires at least 2D, got {:?}",
                input.shape
            )));
        }

        let seq_len = input.dim(1)?;
        let result = self.sequential_scan(input, op)?;

        let elapsed = start.elapsed();
        self.scan_operations
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        self.total_latency_ns.fetch_add(
            elapsed.as_nanos() as u64,
            std::sync::atomic::Ordering::Relaxed,
        );
        self.memory_transfers
            .fetch_add(seq_len as u64, std::sync::atomic::Ordering::Relaxed);

        debug!(
            "Parallel scan completed in {}us for sequence length {}",
            elapsed.as_micros(),
            seq_len
        );

        Ok(result)
    }

    /// Sequential prefix scan using GPU ops (no host download)
    ///
    /// For 2D: processes `[batch, seq_len]` by iterating timesteps on GPU.
    /// For 3D: processes `[batch, seq_len, features]` similarly.
    pub fn sequential_scan(&self, input: &GpuTensor, op: ScanOperator) -> Result<GpuTensor, MLError> {
        if input.shape.len() == 2 {
            self.sequential_scan_2d(input, op)
        } else if input.shape.len() == 3 {
            self.sequential_scan_3d(input, op)
        } else {
            Err(MLError::InvalidInput(format!(
                "sequential_scan requires 2D or 3D, got {:?}",
                input.shape
            )))
        }
    }

    /// 2D scan: `[batch, seq_len]` -> `[batch, seq_len]`
    fn sequential_scan_2d(&self, input: &GpuTensor, op: ScanOperator) -> Result<GpuTensor, MLError> {
        let _batch_size = input.dim(0)?;
        let seq_len = input.dim(1)?;

        if seq_len == 0 {
            return Ok(input.clone());
        }

        // Extract first timestep: column 0 -> [batch, 1]
        let mut acc = gpu_narrow_2d(input, 1, 0, 1)?; // [batch, 1]
        let mut columns: Vec<GpuTensor> = Vec::with_capacity(seq_len);
        columns.push(acc.clone());

        for t in 1..seq_len {
            let col = gpu_narrow_2d(input, 1, t, 1)?; // [batch, 1]
            acc = self.apply_op_tensor(&acc, &col, op)?;
            columns.push(acc.clone());
        }

        // Concatenate columns along dim1: [batch, seq_len]
        let mut result = columns.first().ok_or_else(|| {
            MLError::InvalidInput("empty columns".to_owned())
        })?.clone();
        for col in columns.iter().skip(1) {
            result = crate::gpu_tensor::gpu_cat_dim1(&result, col)?;
        }

        Ok(result)
    }

    /// 3D scan: `[batch, seq_len, features]` -> `[batch, seq_len, features]`
    fn sequential_scan_3d(&self, input: &GpuTensor, op: ScanOperator) -> Result<GpuTensor, MLError> {
        let batch_size = input.dim(0)?;
        let seq_len = input.dim(1)?;
        let _features = input.dim(2)?;

        if seq_len == 0 {
            return Ok(input.clone());
        }

        // Extract first timestep: [batch, features]
        let mut acc = gpu_select_dim1(input, 0)?; // [batch, features]
        let mut steps: Vec<GpuTensor> = Vec::with_capacity(seq_len);
        steps.push(acc.clone());

        for t in 1..seq_len {
            let step = gpu_select_dim1(input, t)?; // [batch, features]
            acc = self.apply_op_tensor(&acc, &step, op)?;
            steps.push(acc.clone());
        }

        // Stack steps: each [batch, features] -> [seq_len, batch, features]
        // Then rearrange to [batch, seq_len, features]
        let mut batch_sequences: Vec<GpuTensor> = Vec::with_capacity(batch_size);
        for b in 0..batch_size {
            let mut step_rows: Vec<GpuTensor> = Vec::with_capacity(seq_len);
            for step in &steps {
                // step: [batch, features] -> extract row b: [1, features]
                let row = gpu_narrow_2d(step, 0, b, 1)?;
                step_rows.push(row);
            }
            // cat [seq_len] items of [1, features] -> [seq_len, features]
            let batch_seq = gpu_cat_dim0(&step_rows, &self.stream)?;
            batch_sequences.push(batch_seq);
        }
        // Stack [batch] items of [seq_len, features] -> [batch, seq_len, features]
        crate::gpu_tensor::gpu_stack_2d(&batch_sequences)
    }

    /// Apply scan operator element-wise on two GPU tensors
    fn apply_op_tensor(&self, left: &GpuTensor, right: &GpuTensor, op: ScanOperator) -> Result<GpuTensor, MLError> {
        match op {
            ScanOperator::Add => gpu_add(left, right),
            ScanOperator::Mul => gpu_mul(left, right),
            ScanOperator::Max => {
                // max(a, b) on GPU: no dedicated op, use (a + b + |a - b|) / 2
                let sum = gpu_add(left, right)?;
                let diff = crate::gpu_tensor::gpu_sub(left, right)?;
                let abs_diff = crate::gpu_tensor::gpu_abs(&diff)?;
                let numerator = gpu_add(&sum, &abs_diff)?;
                crate::gpu_tensor::gpu_scale(&numerator, 0.5)
            },
            ScanOperator::Min => {
                // min(a, b) on GPU: (a + b - |a - b|) / 2
                let sum = gpu_add(left, right)?;
                let diff = crate::gpu_tensor::gpu_sub(left, right)?;
                let abs_diff = crate::gpu_tensor::gpu_abs(&diff)?;
                let numerator = crate::gpu_tensor::gpu_sub(&sum, &abs_diff)?;
                crate::gpu_tensor::gpu_scale(&numerator, 0.5)
            },
            ScanOperator::SSMScan => {
                // SSM scan: new = 0.9 * old + 0.1 * input
                let scaled_left = crate::gpu_tensor::gpu_scale(left, 0.9)?;
                let scaled_right = crate::gpu_tensor::gpu_scale(right, 0.1)?;
                gpu_add(&scaled_left, &scaled_right)
            },
        }
    }

    /// Get performance metrics
    pub fn get_performance_metrics(&self) -> HashMap<String, FixedPoint> {
        let mut metrics = HashMap::new();

        let ops = self
            .scan_operations
            .load(std::sync::atomic::Ordering::Relaxed);
        let total_latency = self
            .total_latency_ns
            .load(std::sync::atomic::Ordering::Relaxed);
        let transfers = self
            .memory_transfers
            .load(std::sync::atomic::Ordering::Relaxed);

        metrics.insert(
            "scan_operations".to_owned(),
            FixedPoint::from_f64(ops as f64),
        );
        metrics.insert(
            "total_latency_ns".to_owned(),
            FixedPoint::from_f64(total_latency as f64),
        );
        metrics.insert(
            "memory_transfers".to_owned(),
            FixedPoint::from_f64(transfers as f64),
        );

        if ops > 0 {
            let avg_latency = total_latency as f64 / ops as f64;
            metrics.insert(
                "avg_latency_ns".to_owned(),
                FixedPoint::from_f64(avg_latency),
            );

            if total_latency > 0 {
                let throughput = transfers as f64 / (total_latency as f64 / 1_000_000_000.0);
                metrics.insert(
                    "throughput_elements_per_sec".to_owned(),
                    FixedPoint::from_f64(throughput),
                );
            }
        }

        metrics.insert(
            "parallel_threshold".to_owned(),
            FixedPoint::from_f64(self.parallel_threshold as f64),
        );
        metrics.insert(
            "block_size".to_owned(),
            FixedPoint::from_f64(self.block_size as f64),
        );

        metrics
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cudarc::driver::CudaContext;

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA context required");
        ctx.new_stream().expect("Failed to create CUDA stream")
    }

    #[test]
    fn test_parallel_scan_engine_creation() {
        let stream = test_stream();
        let _engine = ParallelScanEngine::new(stream, 1_000_000);
    }

    #[test]
    fn test_sequential_scan_add() -> Result<(), MLError> {
        let stream = test_stream();
        let engine = ParallelScanEngine::new(stream.clone(), 1_000_000);
        let input = GpuTensor::from_vec(
            vec![1.0_f32, 2.0, 3.0, 4.0, 5.0],
            &[1, 5],
            &stream,
        )?;
        let result = engine.sequential_scan(&input, ScanOperator::Add)?;
        let actual = result.to_vec()?; // test-only readback
        let expected = vec![1.0, 3.0, 6.0, 10.0, 15.0];
        for (a, e) in actual.iter().zip(expected.iter()) {
            assert!((a - e).abs() < 1e-6, "Expected {}, got {}", e, a);
        }
        Ok(())
    }

    #[test]
    fn test_parallel_prefix_scan() -> Result<(), MLError> {
        let stream = test_stream();
        let engine = ParallelScanEngine::new(stream.clone(), 1_000_000);
        let input = GpuTensor::from_vec(
            vec![1.0_f32, 2.0, 3.0],
            &[1, 3],
            &stream,
        )?;
        let result = engine.parallel_prefix_scan(&input, ScanOperator::Add)?;
        let actual = result.to_vec()?; // test-only readback
        let expected = vec![1.0, 3.0, 6.0];
        for (a, e) in actual.iter().zip(expected.iter()) {
            assert!((a - e).abs() < 1e-6, "Expected {}, got {}", e, a);
        }
        Ok(())
    }
}
