//! CUDA Graph Capture and Replay for PINN Forward Pass
//!
//! This module provides graph capture functionality that records the forward pass
//! into a CUDA graph for near-zero-overhead replay during training.
//!
//! ## Key Insight: Stable Pointers
//!
//! CUDA Graphs capture POINTERS, not VALUES. The graph records:
//! "Launch Kernel A, reading from Address X, writing to Address Y."
//!
//! This means:
//! - **Wrong**: `weights = weights - lr * grad` (allocates new memory → graph breaks)
//! - **Right**: `weights.sub_assign_(lr * grad)` (in-place update → graph valid)
//!
//! With proper in-place optimizer updates, the graph NEVER needs recapturing
//! during training.
//!
//! ## Usage
//!
//! ```ignore
//! let ctx = Arc::new(PinnStreamContext::new(0)?);
//! let mut graph = PinnGraph::new(ctx.clone());
//!
//! // Training loop - graph captured once, replayed 49,999 times
//! for epoch in 1..=50_000 {
//!     if graph.is_captured() {
//!         graph.launch()?;  // Fast: ~4µs
//!     } else {
//!         graph.capture(|| solver.forward_on_stream(&x, &mut ws, &ctx))?;
//!     }
//!
//!     // Optimizer uses in-place updates - graph remains valid!
//!     optimizer.step_inplace(&mut params)?;
//! }
//! ```
//!
//! ## When to Invalidate
//!
//! Only invalidate the graph for rare structural changes:
//! - Batch size change (workspace buffer addresses change)
//! - Model topology change (add/remove layers)
//! - Manual buffer reallocation

use rtx_tensor::{TensorError, Result};
use std::sync::Arc;

#[cfg(feature = "cuda")]
use super::cuda_stream_context::PinnStreamContext;

#[cfg(feature = "cuda")]
use cudarc::driver::safe::CudaGraph;

#[cfg(feature = "cuda")]
use cudarc::driver::sys::{CUstreamCaptureMode, CUgraphInstantiate_flags_enum};

// =============================================================================
// PINN GRAPH
// =============================================================================

/// CUDA graph for accelerated PINN forward pass.
///
/// Captures the forward pass operations into a graph that can be replayed
/// with minimal CPU overhead (~4µs vs ~44µs for 11 kernel launches).
#[cfg(feature = "cuda")]
pub struct PinnGraph {
    /// The instantiated graph (ready for replay)
    graph: Option<CudaGraph>,
    /// Stream context used for capture and launch
    ctx: Arc<PinnStreamContext>,
    /// Input shape (graphs are shape-specific)
    input_shape: Vec<usize>,
    /// Launch count for statistics
    launch_count: u64,
    /// Whether capture is in progress
    capturing: bool,
}

#[cfg(feature = "cuda")]
impl PinnGraph {
    /// Create a new (uncaptured) graph wrapper.
    ///
    /// # Arguments
    /// * `ctx` - The unified stream context for all operations
    pub fn new(ctx: Arc<PinnStreamContext>) -> Self {
        Self {
            graph: None,
            ctx,
            input_shape: Vec::new(),
            launch_count: 0,
            capturing: false,
        }
    }

    /// Check if the graph has been captured.
    pub fn is_captured(&self) -> bool {
        self.graph.is_some()
    }

    /// Get the number of times this graph has been launched.
    pub fn launch_count(&self) -> u64 {
        self.launch_count
    }

    /// Get the input shape the graph was captured with.
    pub fn input_shape(&self) -> &[usize] {
        &self.input_shape
    }

    /// Capture operations into a CUDA graph.
    ///
    /// The closure `capture_fn` should execute the forward pass operations.
    /// All operations must use the same stream (via PinnStreamContext).
    ///
    /// # Arguments
    /// * `input_shape` - Shape of the input tensor (for validation)
    /// * `capture_fn` - Closure that executes the operations to capture
    ///
    /// # Example
    /// ```ignore
    /// graph.capture(&x.shape().dims(), || {
    ///     solver.u_net.forward_on_stream(&x, &mut ws, &ctx)
    /// })?;
    /// ```
    pub fn capture<F>(&mut self, input_shape: &[usize], capture_fn: F) -> Result<()>
    where
        F: FnOnce() -> Result<()>,
    {
        // Check if shape changed
        if self.is_captured() && self.input_shape != input_shape {
            self.invalidate();
        }

        // Already captured with same shape
        if self.is_captured() {
            return Ok(());
        }

        // CRITICAL: Full device synchronization BEFORE capture.
        // This ensures all tensor allocations and operations on ALL streams are complete.
        // Without this, CUDA graph capture fails with STREAM_CAPTURE_ISOLATION because
        // the capture stream doesn't "own" the tensor data that was written by other streams.
        //
        // Note: This is a CPU-blocking operation, but it only happens once during the
        // first capture. Subsequent replays are ultra-fast (~4µs).
        self.ctx.device_synchronize()?;

        // CRITICAL: Re-establish stream dependency edge AFTER all tensor allocations.
        // The join() in PinnStreamContext::new() may have happened before the solver/model
        // was created. By joining again here, we ensure the capture stream has dependency
        // edges to ALL tensor writes, including model weights created after context init.
        //
        // Without this, even with device_synchronize(), CUDA graph capture fails because
        // the capture stream's dependency graph doesn't include the model weight allocations.
        self.ctx.join_with_default_stream()?;

        // Begin stream capture in GLOBAL mode
        // GLOBAL mode allows all operations from any stream to be captured into the graph.
        // This is necessary when input tensors were allocated on a different stream.
        // With our singleton context, all memory is in the same virtual address space.
        self.capturing = true;
        self.ctx.stream()
            .begin_capture(CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_GLOBAL)
            .map_err(|e| TensorError::device(format!("Failed to begin graph capture: {:?}", e)))?;

        // Execute the forward pass (operations are recorded, not executed)
        let result = capture_fn();

        // End capture regardless of result
        let graph_result = self.ctx.stream()
            .end_capture(CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH);

        self.capturing = false;

        // Check if capture_fn succeeded
        result?;

        // Check if end_capture succeeded
        let graph = graph_result
            .map_err(|e| TensorError::device(format!("Failed to end graph capture: {:?}", e)))?
            .ok_or_else(|| TensorError::device("Graph capture returned empty (no operations captured)"))?;

        // Store the captured graph
        self.graph = Some(graph);
        self.input_shape = input_shape.to_vec();

        #[cfg(debug_assertions)]
        eprintln!(
            "[PinnGraph] CUDA graph captured for input shape {:?}",
            self.input_shape
        );

        Ok(())
    }

    /// Capture and immediately execute (for first iteration).
    ///
    /// This is a convenience method that captures the graph and then
    /// executes it, ensuring the first forward pass produces correct output.
    ///
    /// # Arguments
    /// * `input_shape` - Shape of the input tensor
    /// * `capture_fn` - Closure that executes the operations to capture
    pub fn capture_and_launch<F>(&mut self, input_shape: &[usize], capture_fn: F) -> Result<()>
    where
        F: FnOnce() -> Result<()>,
    {
        self.capture(input_shape, capture_fn)?;
        // Note: capture executes the operations during recording,
        // so we don't need an additional launch here.
        // The graph will be used for subsequent calls.
        Ok(())
    }

    /// Replay the captured graph.
    ///
    /// This launches all captured operations with a single driver call,
    /// eliminating kernel launch overhead.
    ///
    /// # Returns
    /// Error if the graph hasn't been captured yet.
    pub fn launch(&mut self) -> Result<()> {
        let graph = self.graph.as_ref()
            .ok_or_else(|| TensorError::runtime("Graph not captured - call capture() first"))?;

        graph.launch()
            .map_err(|e| TensorError::device(format!("Graph launch failed: {:?}", e)))?;

        self.launch_count += 1;

        Ok(())
    }

    /// Invalidate the captured graph.
    ///
    /// Call this when:
    /// - Batch size changes
    /// - Model topology changes
    /// - Workspace buffers are reallocated
    ///
    /// Note: You do NOT need to invalidate after optimizer steps if the
    /// optimizer uses in-place updates (which preserve memory addresses).
    pub fn invalidate(&mut self) {
        #[cfg(debug_assertions)]
        if self.graph.is_some() {
            eprintln!("[PinnGraph] Invalidating CUDA graph (was captured for {:?})", self.input_shape);
        }
        self.graph = None;
        self.input_shape.clear();
        // Note: launch_count is preserved for statistics
    }

    /// Check if capture is currently in progress.
    ///
    /// This is mainly for debugging and error messages.
    pub fn is_capturing(&self) -> bool {
        self.capturing
    }
}

#[cfg(feature = "cuda")]
impl std::fmt::Debug for PinnGraph {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PinnGraph")
            .field("captured", &self.is_captured())
            .field("input_shape", &self.input_shape)
            .field("launch_count", &self.launch_count)
            .field("capturing", &self.capturing)
            .finish()
    }
}

#[cfg(feature = "cuda")]
impl Drop for PinnGraph {
    fn drop(&mut self) {
        self.invalidate();
    }
}

// =============================================================================
// NON-CUDA STUB
// =============================================================================

#[cfg(not(feature = "cuda"))]
pub struct PinnGraph {
    _private: (),
}

#[cfg(not(feature = "cuda"))]
impl PinnGraph {
    pub fn new(_ctx: std::sync::Arc<super::cuda_stream_context::PinnStreamContext>) -> Self {
        Self { _private: () }
    }

    pub fn is_captured(&self) -> bool {
        false
    }

    pub fn launch_count(&self) -> u64 {
        0
    }

    pub fn capture<F>(&mut self, _input_shape: &[usize], _capture_fn: F) -> Result<()>
    where
        F: FnOnce() -> Result<()>,
    {
        Err(TensorError::device("CUDA feature not enabled"))
    }

    pub fn launch(&mut self) -> Result<()> {
        Err(TensorError::device("CUDA feature not enabled"))
    }

    pub fn invalidate(&mut self) {}
}

// =============================================================================
// TESTS
// =============================================================================

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

    #[test]
    #[cfg(feature = "cuda")]
    fn test_graph_creation() {
        use super::super::cuda_stream_context::PinnStreamContext;

        let ctx = Arc::new(PinnStreamContext::new(0).unwrap());
        let graph = PinnGraph::new(ctx);

        assert!(!graph.is_captured());
        assert_eq!(graph.launch_count(), 0);
        assert!(graph.input_shape().is_empty());
    }

    #[test]
    #[cfg(feature = "cuda")]
    fn test_graph_capture_simple() {
        use super::super::cuda_stream_context::PinnStreamContext;
        use rtx_tensor::{Tensor, Device};

        let ctx = Arc::new(PinnStreamContext::new(0).unwrap());
        let mut graph = PinnGraph::new(ctx.clone());
        let device = Device::cuda(0).unwrap();

        // Create test tensors
        let x = Tensor::randn(&[100, 1], &device).unwrap();
        let b = Tensor::randn(&[1, 64], &device).unwrap();
        let mut output = Tensor::zeros(&[100, 128], &device).unwrap();

        // CRITICAL: Device-wide synchronize before capture to ensure all tensor
        // allocations (which may use other streams) are complete. CUDA graph
        // capture cannot have cross-stream dependencies.
        ctx.device_synchronize().unwrap();

        // Capture a simple operation
        let result = graph.capture(&[100, 1], || {
            ctx.fourier_features_out(&x, &b, 6.28, &mut output)
        });

        assert!(result.is_ok(), "Capture failed: {:?}", result.err());
        assert!(graph.is_captured());
        assert_eq!(graph.input_shape(), &[100, 1]);
    }

    #[test]
    #[cfg(feature = "cuda")]
    fn test_graph_replay() {
        use super::super::cuda_stream_context::PinnStreamContext;
        use rtx_tensor::{Tensor, Device};

        let ctx = Arc::new(PinnStreamContext::new(0).unwrap());
        let mut graph = PinnGraph::new(ctx.clone());
        let device = Device::cuda(0).unwrap();

        let x = Tensor::randn(&[100, 1], &device).unwrap();
        let b = Tensor::randn(&[1, 64], &device).unwrap();
        let mut output = Tensor::zeros(&[100, 128], &device).unwrap();

        // CRITICAL: Device-wide synchronize before capture to ensure all tensor
        // allocations (which may use other streams) are complete. CUDA graph
        // capture cannot have cross-stream dependencies.
        ctx.device_synchronize().unwrap();

        // Capture
        graph.capture(&[100, 1], || {
            ctx.fourier_features_out(&x, &b, 6.28, &mut output)
        }).unwrap();

        // Replay multiple times
        for _ in 0..10 {
            graph.launch().unwrap();
        }

        ctx.synchronize().unwrap();
        assert_eq!(graph.launch_count(), 10);
    }

    #[test]
    #[cfg(feature = "cuda")]
    fn test_graph_invalidation() {
        use super::super::cuda_stream_context::PinnStreamContext;

        let ctx = Arc::new(PinnStreamContext::new(0).unwrap());
        let mut graph = PinnGraph::new(ctx.clone());

        // Capture with a no-op (empty graph will fail, but that's OK for this test)
        // We just want to test the invalidation logic
        let _ = graph.capture(&[100, 1], || Ok(()));

        // Invalidate
        graph.invalidate();

        assert!(!graph.is_captured());
        assert!(graph.input_shape().is_empty());
    }
}