//! Automatic differentiation engine for RustyTorch++
//!
//! This crate provides the `Autodiff<B>` decorator pattern for automatic
//! differentiation, inspired by Burn's architecture.
//!
//! ## Key Features
//!
//! - **Zero inference overhead**: Use raw backend without `Autodiff` wrapper
//! - **No global state**: Per-tensor gradient nodes with `Arc` reference counting
//! - **Composable**: `Autodiff<Quantized<CudaBackend>>` is possible
//! - **Type-safe**: Training vs inference is a type-level distinction
//!
//! ## Usage
//!
//! ```rust,ignore
//! use rtx_backend_cuda::CudaBackend;
//! use rtx_autograd::{Autodiff, AutodiffBackend, no_grad};
//!
//! // Training with gradients
//! type TrainingBackend = Autodiff<CudaBackend>;
//! let x = TrainingBackend::from_data(&[1.0, 2.0], [2], &device);
//! let x = TrainingBackend::require_grad(x);
//! let y = TrainingBackend::mul(&x, &x);
//! let grads = TrainingBackend::backward(&y);
//!
//! // Inference without gradients (zero overhead!)
//! type InferenceBackend = CudaBackend;
//! let z = InferenceBackend::mul(&a, &b);  // No gradient tracking!
//!
//! // Disable gradients temporarily
//! let result = no_grad(|| {
//!     TrainingBackend::mul(&x, &y)  // No gradient node created
//! });
//! ```
//!
//! ## Architecture
//!
//! ```text
//! Autodiff<B: Backend>
//!     |
//!     +-- AutodiffTensor<B, D> (wrapper around B::TensorPrimitive<D>)
//!     |       |
//!     |       +-- inner: B::TensorPrimitive<D>
//!     |       +-- node: Option<Arc<AutodiffNode<B>>>
//!     |       +-- id: TensorId
//!     |       +-- requires_grad: bool
//!     |
//!     +-- AutodiffNode<B> (gradient computation node)
//!             |
//!             +-- parents: Vec<ParentRef<B>>
//!             +-- backward_fn: Box<dyn AutodiffBackwardFn<B>>
//!             +-- saved_tensors: type-erased saved inputs
//! ```

pub mod autodiff;
pub mod checkpoint;
pub mod compiled;
pub mod error;
pub mod forward_mode;
pub mod func;
pub mod generic_compat;
pub mod profiler;
pub mod vmap;

// Re-export primary types at crate root
pub use autodiff::{
    // Core types
    Autodiff,
    // Traits
    AutodiffBackwardFn,
    AutodiffDevice,
    AutodiffNode,
    AutodiffTensor,
    GradTensor,
    GradientStorage,
    TensorId,
    // Graph operations
    backward_impl,
    enable_grad,
    is_grad_enabled,
    // Context managers
    no_grad,
    set_grad_enabled,
    topological_sort,
};

// Re-export from rtx-backend
pub use rtx_backend::AutodiffBackend;

pub use error::{AutogradError, Result as AutogradResult};

// Re-export functional transforms (torch.func equivalent)
pub use func::{
    // Reverse-mode AD (implemented)
    grad,
    grad_with_argnums,
    hessian,
    jacfwd,
    jacfwd_with_argnums,
    jacrev,
    jacrev_with_argnums,
    // Forward-mode AD (now implemented)
    jvp,
    jvp_single,
    value_and_grad,
    vjp,
    // Vectorized map (now implemented)
    vmap,
    vmap_simple,
};

// Re-export forward-mode AD types
pub use forward_mode::{DualVariable, jacfwd_impl, jvp_impl, jvp_scalar};

// Re-export vmap types
pub use vmap::{BatchedVariable, vmap_impl};

// Re-export checkpoint types (gradient/activation checkpointing)
pub use checkpoint::{
    // Strategy traits and implementations
    AdaptiveCheckpointStrategy,
    // Configuration
    CheckpointConfig,
    CheckpointStrategy,
    EveryNthStrategy,
    ManualCheckpointStrategy,
    NoCheckpointStrategy,
    // Recomputation context
    RecomputeContext,
    RecomputeGuard,
    SqrtCheckpointStrategy,
    // Functions
    checkpoint,
    checkpoint_sequential,
    checkpoint_with_strategy,
    // Global state
    disable_checkpointing,
    enable_checkpointing,
    is_checkpointing_enabled,
};

// Re-export profiler types
pub use profiler::{
    AutogradProfiler, BottleneckInfo, BottleneckType, EventType, GradFlowEdge, GradFlowNode,
    GradientFlow, MemorySnapshot, OperationStats, ProfiledEvent, ProfilerConfig, ProfilerReport,
    RecordGuard, get_profiler,
};

// Re-export compiled autograd types
pub use compiled::{
    BackwardEdge,
    // Capture and optimization
    BackwardGraphCapture,
    BackwardGraphOptimizer,
    BackwardNode,
    BackwardNodeId,
    BackwardOpType,
    CompiledAutogradContext,
    // Executor
    CompiledBackward,
    CompiledBackwardConfig,
    // Graph types
    CompiledBackwardGraph,
    ExecutionStats,
    GraphSignature,
    OptimizationConfig,
    OptimizationStats,
    SavedBackwardData,
    disable_compiled_autograd,
    // Global state
    enable_compiled_autograd,
    is_compiled_autograd_enabled,
};

// Re-export GenericTensor compatibility layer
pub use generic_compat::{
    AutodiffTensorGenericExt, GenericTensorAutodiffExt, from_autodiff, from_autodiff_ref,
    into_autodiff,
};

// ============================================================================
// COMPATIBILITY ALIASES
// These aliases provide backward compatibility with older code that uses
// the tape-based API style. The decorator pattern doesn't use a global tape,
// so some of these are no-ops.
// ============================================================================

/// Type alias for backward compatibility.
/// In the decorator pattern, tensors with gradients are `AutodiffTensor`.
pub type TensorAutograd<B, const D: usize> = AutodiffTensor<B, D>;

/// Type alias for node IDs. Same as `TensorId`.
pub type NodeId = TensorId;

/// Backward pass - alias for `backward_impl`.
///
/// This is a convenience wrapper that handles common backward pass patterns.
/// For the decorator pattern, use `Autodiff::<B>::backward(&tensor)` directly.
pub fn backward<T>(
    _loss: T,
    _retain_graph: Option<bool>,
) -> std::collections::HashMap<TensorId, Vec<f32>> {
    // The decorator pattern computes gradients differently.
    // This stub returns an empty map for compilation compatibility.
    // Real gradient computation should use Autodiff::<B>::backward().
    std::collections::HashMap::new()
}

/// Clear the computation tape.
///
/// In the decorator pattern, there is no global tape to clear.
/// This function is a no-op for backward compatibility.
#[inline]
pub fn clear_tape() {
    // No-op: The decorator pattern doesn't use a global tape.
    // Gradients are stored per-tensor and cleared automatically
    // when tensors go out of scope.
}

// ============================================================================
// ADDITIONAL COMPATIBILITY TYPES
// These types provide backward compatibility with older code that expects
// a tape-based autograd API.
// ============================================================================

use rtx_tensor::Tensor;
use std::sync::atomic::{AtomicUsize, Ordering};

/// Autograd context for tracking operations.
///
/// This is a stub implementation for backward compatibility.
/// The decorator pattern (`Autodiff<B>`) handles gradient tracking differently.
#[derive(Debug, Default)]
pub struct AutogradContext {
    _marker: std::marker::PhantomData<()>,
}

impl AutogradContext {
    /// Create a new autograd context
    pub fn new() -> Self {
        Self {
            _marker: std::marker::PhantomData,
        }
    }

    /// Get mutable reference to current tape (stub - returns None)
    pub fn current_tape_mut(&self) -> Option<&mut AutogradTape> {
        // Stub: decorator pattern doesn't use tapes
        None
    }

    /// Clear the autograd tape (no-op in decorator pattern)
    pub fn clear_tape(&self) {
        // No-op: decorator pattern doesn't use global tapes
    }
}

/// Stub type for backward compatibility with tape-based API
#[derive(Debug, Default)]
pub struct AutogradTape {
    _marker: std::marker::PhantomData<()>,
}

impl AutogradTape {
    /// Create a new tape
    pub fn new() -> Self {
        Self {
            _marker: std::marker::PhantomData,
        }
    }

    /// Backward pass (stub)
    pub fn backward(
        &self,
        _node_id: TensorId,
        _config: Option<BackwardConfig>,
    ) -> error::Result<std::collections::HashMap<TensorId, Tensor>> {
        // Stub: returns empty gradients
        Ok(std::collections::HashMap::new())
    }
}

/// Configuration for backward pass (stub)
#[derive(Debug, Clone, Default)]
pub struct BackwardConfig {
    /// Retain computation graph after backward
    pub retain_graph: bool,
    /// Create graph for higher-order gradients
    pub create_graph: bool,
}

/// Variable wrapping a tensor with gradient tracking support.
///
/// This is a compatibility type for code that uses the older Variable API.
/// For new code, prefer using `Autodiff<B>` directly.
#[derive(Debug, Clone)]
pub struct Variable {
    tensor: Tensor,
    requires_grad: bool,
    grad: Option<Tensor>,
    id: usize,
}

static NEXT_VAR_ID: AtomicUsize = AtomicUsize::new(1);

impl Variable {
    /// Create a new Variable from a tensor
    pub fn new(tensor: Tensor, requires_grad: bool) -> Self {
        Self {
            tensor,
            requires_grad,
            grad: None,
            id: NEXT_VAR_ID.fetch_add(1, Ordering::SeqCst),
        }
    }

    /// Create a Variable from a tensor (defaults to requiring gradients)
    pub fn from_tensor(tensor: Tensor) -> Self {
        Self::new(tensor, true)
    }

    /// Get a reference to the underlying tensor
    pub fn tensor(&self) -> &Tensor {
        &self.tensor
    }

    /// Get the underlying tensor (consuming self)
    pub fn into_tensor(self) -> Tensor {
        self.tensor
    }

    /// Check if this variable requires gradients
    pub fn requires_grad(&self) -> bool {
        self.requires_grad
    }

    /// Set whether this variable requires gradients
    pub fn set_requires_grad(&mut self, requires_grad: bool) {
        self.requires_grad = requires_grad;
    }

    /// Get the gradient tensor if it exists
    pub fn grad(&self) -> Option<&Tensor> {
        self.grad.as_ref()
    }

    /// Set the gradient tensor
    pub fn set_grad(&mut self, grad: Tensor) {
        self.grad = Some(grad);
    }

    /// Clear the gradient
    pub fn zero_grad(&mut self) {
        self.grad = None;
    }

    /// Get the variable ID
    pub fn id(&self) -> usize {
        self.id
    }

    /// Get the shape of the underlying tensor
    pub fn shape(&self) -> Vec<usize> {
        self.tensor.shape().to_vec()
    }

    /// Add two variables
    pub fn add(&self, other: &Variable) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .add(&other.tensor)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(
            result,
            self.requires_grad || other.requires_grad,
        ))
    }

    /// Multiply two variables
    pub fn mul(&self, other: &Variable) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .mul(&other.tensor)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(
            result,
            self.requires_grad || other.requires_grad,
        ))
    }

    /// Multiply by a scalar
    pub fn multiply_scalar(&self, scalar: f32) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .mul_scalar(scalar)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(result, self.requires_grad))
    }

    /// Matrix multiplication
    pub fn matmul(&self, other: &Variable) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .matmul(&other.tensor)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(
            result,
            self.requires_grad || other.requires_grad,
        ))
    }

    /// Sum all elements (or along a dimension)
    pub fn sum(&self, dim: Option<usize>) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .sum(dim)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(result, self.requires_grad))
    }

    /// Sum all elements (convenience method)
    pub fn sum_all(&self) -> Result<Variable, error::AutogradError> {
        self.sum(None)
    }

    /// Mean along specified dimensions
    pub fn mean_dims(&self, dims: &[i32], keepdim: bool) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .mean(dims, keepdim)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(result, self.requires_grad))
    }

    /// Mean of all elements (convenience method)
    pub fn mean(&self) -> Result<Variable, error::AutogradError> {
        self.mean_dims(&[], false)
    }

    /// Alias for mean()
    pub fn mean_all(&self) -> Result<Variable, error::AutogradError> {
        self.mean()
    }

    /// Backward pass (stub - returns empty gradients)
    pub fn backward(&self) -> std::collections::HashMap<usize, Tensor> {
        // Stub implementation - real implementation would compute gradients
        std::collections::HashMap::new()
    }

    /// Backward gradient computation (stub)
    pub fn backward_grad(&self, _grad: &Tensor) -> error::Result<()> {
        // Stub implementation
        Ok(())
    }

    /// Element-wise multiply (alias for mul)
    pub fn multiply(&self, other: &Variable) -> Result<Variable, error::AutogradError> {
        self.mul(other)
    }

    /// Get the value (underlying tensor) - alias for tensor()
    pub fn value(&self) -> &Tensor {
        self.tensor()
    }

    /// Get scalar value from a single-element tensor
    pub fn item(&self) -> f32 {
        // Get first element as scalar
        self.tensor
            .to_vec()
            .ok()
            .and_then(|v| v.first().copied())
            .unwrap_or(0.0)
    }

    /// Element-wise division
    pub fn div(&self, other: &Variable) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .div(&other.tensor)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(
            result,
            self.requires_grad || other.requires_grad,
        ))
    }

    /// Element-wise subtraction
    pub fn sub(&self, other: &Variable) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .sub(&other.tensor)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(
            result,
            self.requires_grad || other.requires_grad,
        ))
    }

    /// Negate the variable
    pub fn neg(&self) -> Result<Variable, error::AutogradError> {
        self.multiply_scalar(-1.0)
    }

    /// Exponentiation
    pub fn exp(&self) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .exp()
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(result, self.requires_grad))
    }

    /// Natural logarithm
    pub fn log(&self) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .log()
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(result, self.requires_grad))
    }

    /// Power function (with scalar exponent)
    pub fn pow(&self, exp: f32) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .pow_scalar(exp)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(result, self.requires_grad))
    }

    /// Power function (with tensor exponent)
    pub fn pow_tensor(&self, exp: &Variable) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .pow(&exp.tensor)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(
            result,
            self.requires_grad || exp.requires_grad,
        ))
    }

    /// Square root
    pub fn sqrt(&self) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .sqrt()
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(result, self.requires_grad))
    }

    /// Absolute value
    pub fn abs(&self) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .abs()
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(result, self.requires_grad))
    }

    /// Reshape the variable
    pub fn reshape(&self, shape: &[usize]) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .reshape(shape)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(result, self.requires_grad))
    }

    /// Create a zeros tensor
    pub fn zeros(
        shape: &[usize],
        device: &rtx_tensor::Device,
    ) -> Result<Variable, error::AutogradError> {
        let tensor = Tensor::zeros(shape, device)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(tensor, false))
    }

    /// Create a ones tensor
    pub fn ones(
        shape: &[usize],
        device: &rtx_tensor::Device,
    ) -> Result<Variable, error::AutogradError> {
        let tensor = Tensor::ones(shape, device)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(tensor, false))
    }

    /// Create from data
    pub fn from_data(
        data: Vec<f32>,
        shape: &[usize],
        device: &rtx_tensor::Device,
    ) -> Result<Variable, error::AutogradError> {
        let tensor = Tensor::from_data(data, shape.to_vec(), device)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(tensor, true))
    }

    /// Create from slice (convenience method)
    pub fn from_slice(
        data: &[f32],
        shape: &[usize],
        device: &rtx_tensor::Device,
    ) -> Result<Variable, error::AutogradError> {
        Self::from_data(data.to_vec(), shape, device)
    }

    /// Detach from computation graph (returns a new variable without gradient tracking)
    pub fn detach(&self) -> Variable {
        Variable::new(self.tensor.clone(), false)
    }

    /// Mean squared value (mean of x^2)
    pub fn mean_square(&self) -> Result<Variable, error::AutogradError> {
        let squared = self.mul(self)?;
        squared.mean_all()
    }

    /// Create a constant variable (no gradient tracking)
    pub fn constant(tensor: Tensor) -> Variable {
        Variable::new(tensor, false)
    }

    /// Get the device of the underlying tensor
    pub fn device(&self) -> &rtx_tensor::Device {
        self.tensor.device()
    }

    /// Squeeze dimensions
    pub fn squeeze(&self, dim: Option<i32>) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .squeeze(dim)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(result, self.requires_grad))
    }

    /// Squeeze all singleton dimensions
    pub fn squeeze_all(&self) -> Result<Variable, error::AutogradError> {
        self.squeeze(None)
    }

    /// Unsqueeze - add a dimension
    pub fn unsqueeze(&self, dim: i32) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .unsqueeze(dim)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(result, self.requires_grad))
    }

    /// Transpose the tensor
    pub fn transpose(&self, dim0: i32, dim1: i32) -> Result<Variable, error::AutogradError> {
        let result = self
            .tensor
            .transpose(dim0, dim1)
            .map_err(|e| error::AutogradError::ComputationError(e.to_string()))?;
        Ok(Variable::new(result, self.requires_grad))
    }
}