foxhunt/ml/src/tft/trainable_adapter.rs

//! UnifiedTrainable Trait Implementation for TFT (Temporal Fusion Transformer)
//!
//! This adapter wraps the existing TemporalFusionTransformer implementation with the
//! UnifiedTrainable trait to enable standardized training orchestration. TFT is the most
//! complex architecture with attention mechanisms, variable selection, and quantile outputs.
//!
//! ## Key Features
//!
//! - Forward pass through multi-component architecture (VSN, GRN, attention, quantile)
//! - Quantile loss computation for uncertainty estimation
//! - Backward pass with gradient tracking across all components
//! - Checkpoint save/load using safetensors format
//! - Metrics collection including attention weights and feature importance
//! - Learning rate scheduling support
//!
//! ## Architecture Complexity
//!
//! TFT has the most complex architecture among all models:
//! - 3 Variable Selection Networks (static, historical, future)
//! - 3 Gated Residual Network stacks (encoding layers)
//! - LSTM encoder/decoder for temporal processing
//! - Multi-head self-attention mechanism
//! - Quantile output layer for uncertainty quantification
//!
//! This adapter provides standardized training orchestration while preserving
//! TFT's interpretability features (attention weights, feature importance).

use candle_core::{backprop::GradStore, Device, Tensor};
use candle_nn::{AdamW, Optimizer, ParamsAdamW};
use serde_json;
use std::collections::HashMap;

use super::{TFTConfig, TemporalFusionTransformer};
use crate::training::unified_trainer::{CheckpointMetadata, TrainingMetrics, UnifiedTrainable};
use crate::MLError;

/// Extended TFT with training infrastructure
///
/// This struct wraps TemporalFusionTransformer and adds necessary fields for training:
/// - Adam optimizer for parameter updates
/// - Step counter for learning rate scheduling
/// - Training loss history
/// - Gradient tracking
///
/// Note: TFT manages its own parameters through VarBuilder internally,
/// so we don't need a separate VarMap. Gradient computation is handled
/// by candle's automatic differentiation.
pub struct TrainableTFT {
    /// Core TFT model
    pub model: TemporalFusionTransformer,
    /// AdamW optimizer for parameter updates
    optimizer: AdamW,
    /// Last gradient store from backward pass
    last_grads: Option<GradStore>,
    /// Training step counter
    step_count: usize,
    /// Training loss history
    loss_history: Vec<f64>,
    /// Learning rate (mutable for scheduling)
    learning_rate: f64,
    /// Last computed gradient norm (for monitoring)
    last_grad_norm: f64,
}

impl std::fmt::Debug for TrainableTFT {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TrainableTFT")
            .field("model", &self.model)
            .field("step_count", &self.step_count)
            .field("loss_history_len", &self.loss_history.len())
            .field("learning_rate", &self.learning_rate)
            .field("last_grad_norm", &self.last_grad_norm)
            .finish_non_exhaustive()
    }
}

impl TrainableTFT {
    /// Create new trainable TFT model
    ///
    /// # Arguments
    /// * `config` - TFT configuration
    ///
    /// # Returns
    /// Trainable TFT wrapper ready for training
    pub fn new(config: TFTConfig) -> Result<Self, MLError> {
        // Create TFT model with internal VarBuilder
        let model = TemporalFusionTransformer::new(config.clone())?;
        let learning_rate = config.learning_rate;

        // Initialize AdamW optimizer with model parameters
        let params = model.varmap.all_vars();
        let optimizer = AdamW::new(
            params,
            ParamsAdamW {
                lr: learning_rate,
                beta1: 0.9,
                beta2: 0.999,
                eps: 1e-8,
                weight_decay: config.l2_regularization,
            },
        )
        .map_err(|e| MLError::ModelError(format!("Failed to initialize AdamW optimizer: {}", e)))?;

        Ok(Self {
            model,
            optimizer,
            last_grads: None,
            step_count: 0,
            loss_history: Vec::new(),
            learning_rate,
            last_grad_norm: 0.0,
        })
    }
}

impl UnifiedTrainable for TrainableTFT {
    /// Get model type identifier
    fn model_type(&self) -> &str {
        "TFT"
    }

    /// Get device model is on (CPU or CUDA)
    fn device(&self) -> &Device {
        &self.model.device
    }

    /// Forward pass through model
    ///
    /// TFT requires 3 separate inputs (static, historical, future features).
    /// For unified interface, we assume input is concatenated and split internally.
    ///
    /// # Arguments
    /// * `input` - Concatenated input tensor [batch, total_features]
    ///
    /// # Returns
    /// Quantile predictions tensor [batch, prediction_horizon, num_quantiles]
    fn forward(&mut self, input: &Tensor) -> Result<Tensor, MLError> {
        // TFT's forward expects 3 separate tensors (static, historical, future)
        // For unified interface, we need to split the input tensor
        // This is a simplified version - real implementation would handle proper splitting

        let (batch_size, total_dim) = input.dims2().map_err(|e| MLError::TensorCreationError {
            operation: "forward: get input dims".to_string(),
            reason: e.to_string(),
        })?;

        // Calculate split points based on configuration
        let static_dim = self.model.config.num_static_features;
        let hist_dim = self.model.config.num_unknown_features * self.model.config.sequence_length;
        let future_dim =
            self.model.config.num_known_features * self.model.config.prediction_horizon;

        // Verify total dimension matches
        if total_dim != static_dim + hist_dim + future_dim {
            return Err(MLError::ValidationError {
                message: format!(
                    "Input dimension {} does not match expected {} (static={}, hist={}, future={})",
                    total_dim,
                    static_dim + hist_dim + future_dim,
                    static_dim,
                    hist_dim,
                    future_dim
                ),
            });
        }

        // Split input into 3 components
        let static_features =
            input
                .narrow(1, 0, static_dim)
                .map_err(|e| MLError::TensorCreationError {
                    operation: "forward: narrow static features".to_string(),
                    reason: e.to_string(),
                })?;

        let historical_features =
            input
                .narrow(1, static_dim, hist_dim)
                .map_err(|e| MLError::TensorCreationError {
                    operation: "forward: narrow historical features".to_string(),
                    reason: e.to_string(),
                })?;

        let future_features = input
            .narrow(1, static_dim + hist_dim, future_dim)
            .map_err(|e| MLError::TensorCreationError {
                operation: "forward: narrow future features".to_string(),
                reason: e.to_string(),
            })?;

        // Reshape historical and future to [batch, seq_len, features]
        let historical_reshaped = historical_features
            .reshape((
                batch_size,
                self.model.config.sequence_length,
                self.model.config.num_unknown_features,
            ))
            .map_err(|e| MLError::TensorCreationError {
                operation: "forward: reshape historical".to_string(),
                reason: e.to_string(),
            })?;

        let future_reshaped = future_features
            .reshape((
                batch_size,
                self.model.config.prediction_horizon,
                self.model.config.num_known_features,
            ))
            .map_err(|e| MLError::TensorCreationError {
                operation: "forward: reshape future".to_string(),
                reason: e.to_string(),
            })?;

        // Call TFT's forward method with 3 separate inputs
        self.model
            .forward(&static_features, &historical_reshaped, &future_reshaped)
    }

    /// Compute quantile loss for TFT
    ///
    /// Uses quantile regression loss for uncertainty estimation
    ///
    /// # Arguments
    /// * `predictions` - Quantile predictions [batch, horizon, num_quantiles]
    /// * `targets` - Ground truth tensor [batch, horizon]
    ///
    /// # Returns
    /// Scalar quantile loss tensor
    fn compute_loss(&self, predictions: &Tensor, targets: &Tensor) -> Result<Tensor, MLError> {
        // Delegate to TFT's quantile loss implementation
        self.model
            .quantile_outputs
            .quantile_loss(predictions, targets)
    }

    /// Backward pass to compute gradients
    ///
    /// # Arguments
    /// * `loss` - Scalar loss tensor from compute_loss
    ///
    /// # Returns
    /// Gradient norm for monitoring gradient explosion/vanishing
    fn backward(&mut self, loss: &Tensor) -> Result<f64, MLError> {
        // Trigger backward pass and get gradients
        let grads = loss.backward().map_err(|e| MLError::TensorCreationError {
            operation: "backward: loss.backward()".to_string(),
            reason: e.to_string(),
        })?;

        // Calculate L2 norm of gradients FIRST (before moving grads): ||∇L||₂ = √(Σ grad_i²)
        let mut total_norm_squared = 0.0_f64;

        // Iterate through all model parameters in VarMap
        let varmap_data = self
            .model
            .varmap
            .data()
            .lock()
            .map_err(|e| MLError::TrainingError(format!("Failed to lock VarMap: {}", e)))?;

        for (_name, var) in varmap_data.iter() {
            // Get gradient for this parameter
            if let Some(grad) = grads.get(var.as_tensor()) {
                // Compute squared L2 norm of this parameter's gradient
                let grad_norm_sq = grad
                    .sqr()
                    .and_then(|t| t.sum_all())
                    .and_then(|t| t.to_dtype(candle_core::DType::F64))
                    .and_then(|t| t.to_scalar::<f64>())
                    .map_err(|e| MLError::TensorCreationError {
                        operation: "backward: compute gradient norm".to_string(),
                        reason: e.to_string(),
                    })?;

                total_norm_squared += grad_norm_sq;
            }
        }

        // Compute final L2 norm
        let grad_norm = total_norm_squared.sqrt();

        // Detect gradient explosion/vanishing
        if grad_norm.is_nan() || grad_norm.is_infinite() {
            return Err(MLError::TrainingError(
                "Gradient norm is NaN or Inf - gradient explosion detected".to_string(),
            ));
        }

        self.last_grad_norm = grad_norm;

        // Store gradients for optimizer_step() (move happens here)
        self.last_grads = Some(grads);

        Ok(grad_norm)
    }

    /// Update model parameters using optimizer
    ///
    /// Applies Adam optimizer updates to all trainable parameters in the TFT model.
    /// Uses the AdamW variant with weight decay for regularization.
    ///
    /// Adam update rule: θ = θ - α * m̂ / (√v̂ + ε)
    /// Where:
    /// - m̂ = exponential moving average of gradients (momentum)
    /// - v̂ = exponential moving average of squared gradients (RMSprop)
    /// - α = learning rate
    /// - ε = small constant for numerical stability (1e-8)
    ///
    /// # Returns
    /// Ok(()) on success, MLError on failure
    fn optimizer_step(&mut self) -> Result<(), MLError> {
        // Get gradients from last backward() call
        let grads = self.last_grads.as_ref().ok_or_else(|| {
            MLError::TrainingError(
                "No gradients available. Call backward() before optimizer_step()".to_string(),
            )
        })?;

        // Use Candle's built-in step() method which performs parameter updates
        // This method internally:
        // 1. Uses gradients from the GradStore
        // 2. Updates Adam state (m, v, step count)
        // 3. Computes parameter updates using Adam formula
        // 4. Applies updates to all parameters in the VarMap
        self.optimizer
            .step(grads)
            .map_err(|e| MLError::TrainingError(format!("Optimizer step failed: {}", e)))?;

        self.step_count += 1;

        // Clear gradients after update
        self.last_grads = None;

        Ok(())
    }

    /// Zero gradients before next backward pass
    ///
    /// In Candle, gradients are managed through the automatic differentiation system.
    /// Each call to `backward()` creates a new gradient computation graph, so gradients
    /// don't automatically accumulate between batches like in PyTorch.
    ///
    /// However, we implement explicit gradient zeroing for two reasons:
    /// 1. Defense in depth - ensures no gradient accumulation if training loop is modified
    /// 2. Unified interface compliance - matches expected behavior across all trainable models
    ///
    /// This implementation verifies that the VarMap is accessible and could be extended
    /// in the future if Candle adds explicit gradient accumulation features.
    fn zero_grad(&mut self) -> Result<(), MLError> {
        // Verify VarMap is accessible (defensive check)
        let _varmap_check = self.model.varmap.data().lock().map_err(|e| {
            MLError::TrainingError(format!("Failed to lock VarMap for gradient zeroing: {}", e))
        })?;

        // In Candle, gradients are not stored in VarMap but managed by GradStore
        // returned from backward(). Each backward() call creates a fresh gradient
        // computation, so explicit zeroing is not needed for correctness.
        //
        // However, we maintain this method for:
        // - Interface compliance with UnifiedTrainable trait
        // - Future-proofing if Candle adds gradient accumulation
        // - Documentation of gradient management strategy

        // Reset gradient norm tracking
        self.last_grad_norm = 0.0;

        Ok(())
    }

    /// Get current learning rate
    fn get_learning_rate(&self) -> f64 {
        self.learning_rate
    }

    /// Set learning rate (for scheduling)
    ///
    /// Updates both the cached learning rate and the optimizer's internal learning rate.
    /// This enables learning rate scheduling strategies like step decay, cosine annealing, etc.
    fn set_learning_rate(&mut self, lr: f64) -> Result<(), MLError> {
        if lr <= 0.0 || lr > 1.0 {
            return Err(MLError::ValidationError {
                message: format!("Invalid learning rate: {}. Must be in range (0.0, 1.0]", lr),
            });
        }

        // Update cached learning rate
        self.learning_rate = lr;

        // Update optimizer's learning rate (modifies in-place, no Result returned)
        self.optimizer.set_learning_rate(lr);

        Ok(())
    }

    /// Get current training step count
    fn get_step(&self) -> usize {
        self.step_count
    }

    /// Collect current training metrics
    ///
    /// Includes TFT-specific metrics like attention weights and feature importance
    fn collect_metrics(&self) -> TrainingMetrics {
        let model_metrics = self.model.get_metrics();

        let mut custom_metrics = HashMap::new();
        for (key, value) in model_metrics.iter() {
            custom_metrics.insert(key.clone(), *value);
        }

        // Add training-specific metrics
        custom_metrics.insert("step_count".to_string(), self.step_count as f64);
        custom_metrics.insert("last_grad_norm".to_string(), self.last_grad_norm);

        // Calculate approximate number of parameters from VarMap
        let num_params = self
            .model
            .varmap
            .data()
            .lock()
            .map(|data| {
                data.iter()
                    .map(|(_, var)| var.as_tensor().elem_count())
                    .sum::<usize>()
            })
            .unwrap_or(0);
        custom_metrics.insert("num_parameters".to_string(), num_params as f64);

        TrainingMetrics {
            loss: self.loss_history.last().copied().unwrap_or(0.0),
            val_loss: None, // Will be set by orchestrator during validation
            accuracy: None, // TFT uses quantile loss, not classification accuracy
            learning_rate: self.learning_rate,
            grad_norm: None, // Will be updated by backward() call
            custom_metrics,
        }
    }

    /// Save model checkpoint in standardized format
    ///
    /// Saves model weights via safetensors and metadata via JSON.
    ///
    /// # Arguments
    /// * `checkpoint_path` - Path to save checkpoint (without extension)
    ///
    /// # Returns
    /// Path to saved checkpoint
    fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError> {
        // Create and save checkpoint metadata
        let metadata = CheckpointMetadata {
            model_type: "TFT".to_string(),
            version: self.model.metadata.version.clone(),
            epoch: self.loss_history.len(), // Use loss history length as proxy for epochs
            step: self.step_count,
            timestamp: std::time::SystemTime::now(),
            config: serde_json::to_value(&self.model.config)
                .map_err(|e| MLError::ModelError(format!("Failed to serialize config: {}", e)))?,
            metrics: self.collect_metrics(),
        };

        // Save metadata to JSON
        crate::training::unified_trainer::checkpoint::save_metadata(&metadata, checkpoint_path)?;

        // Save model weights to safetensors format
        let safetensors_path = format!("{}.safetensors", checkpoint_path);

        // Extract tensors from VarMap
        let vars_data = self
            .model
            .varmap
            .data()
            .lock()
            .map_err(|e| MLError::TrainingError(format!("Failed to lock VarMap for checkpoint save: {}", e)))?;

        let mut tensors: HashMap<String, Tensor> = HashMap::new();
        for (name, var) in vars_data.iter() {
            tensors.insert(name.clone(), var.as_tensor().clone());
        }

        // Save using safetensors
        candle_core::safetensors::save(&tensors, &safetensors_path)
            .map_err(|e| MLError::ModelError(format!("Failed to save safetensors: {}", e)))?;

        tracing::info!(
            "Saved TFT checkpoint to {} (step {}, {} tensors)",
            checkpoint_path,
            self.step_count,
            tensors.len()
        );

        Ok(safetensors_path)
    }

    /// Load model checkpoint from standardized format
    ///
    /// Loads model weights from safetensors and metadata from JSON.
    ///
    /// # Arguments
    /// * `checkpoint_path` - Path to checkpoint (without extension)
    ///
    /// # Returns
    /// Loaded checkpoint metadata
    fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result<CheckpointMetadata, MLError> {
        // Load metadata from JSON
        let metadata =
            crate::training::unified_trainer::checkpoint::load_metadata(checkpoint_path)?;

        // Validate model type
        if metadata.model_type != "TFT" {
            return Err(MLError::ModelError(format!(
                "Invalid model type in checkpoint: expected TFT, got {}",
                metadata.model_type
            )));
        }

        // Load model weights from safetensors
        let safetensors_path = format!("{}.safetensors", checkpoint_path);
        let tensors = candle_core::safetensors::load(&safetensors_path, &self.model.device)
            .map_err(|e| MLError::ModelError(format!("Failed to load safetensors: {}", e)))?;

        // Load tensors into VarMap
        let vars_data = self
            .model
            .varmap
            .data()
            .lock()
            .map_err(|e| MLError::TrainingError(format!("Failed to lock VarMap for checkpoint load: {}", e)))?;

        for (name, tensor) in &tensors {
            if let Some(var) = vars_data.get(name) {
                var.set(tensor).map_err(|e| {
                    MLError::ModelError(format!("Failed to set var {}: {}", name, e))
                })?;
            } else {
                tracing::warn!("Checkpoint contains unknown variable: {}", name);
            }
        }

        // Update model state from metadata
        self.step_count = metadata.step;
        self.model.is_trained = true;
        self.learning_rate = metadata.metrics.learning_rate;

        tracing::info!(
            "Loaded TFT checkpoint from {} (step {}, {} tensors)",
            checkpoint_path,
            metadata.step,
            tensors.len()
        );

        Ok(metadata)
    }

    /// Validate model on validation set
    ///
    /// # Arguments
    /// * `val_data` - Validation dataset (input, target) pairs
    ///
    /// # Returns
    /// Validation loss (quantile loss)
    fn validate(&mut self, val_data: &[(Tensor, Tensor)]) -> Result<f64, MLError> {
        let mut total_loss = 0.0;
        let mut count = 0;

        for (input, target) in val_data {
            // Forward pass
            let predictions = self.forward(input)?;

            // Compute loss
            let loss = self.compute_loss(&predictions, target)?;
            let loss_value = loss
                .to_scalar::<f64>()
                .map_err(|e| MLError::TensorCreationError {
                    operation: "validate: loss.to_scalar()".to_string(),
                    reason: e.to_string(),
                })?;

            total_loss += loss_value;
            count += 1;
        }

        if count == 0 {
            return Err(MLError::ValidationError {
                message: "Validation set is empty".to_string(),
            });
        }

        Ok(total_loss / count as f64)
    }
}

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

    #[test]
    fn test_tft_trainable_creation() -> anyhow::Result<()> {
        let config = TFTConfig {
            input_dim: 64,
            hidden_dim: 32,
            num_heads: 4,
            num_layers: 2,
            prediction_horizon: 5,
            sequence_length: 20,
            num_quantiles: 5,
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 49, // 64 - 5 - 10 = 49
            learning_rate: 1e-3,
            ..Default::default()
        };
        let model = TrainableTFT::new(config)?;

        // Test trait methods
        assert_eq!(model.model_type(), "TFT");
        // Device can be CPU or CUDA depending on availability
        let device_str = format!("{:?}", model.device());
        assert!(device_str.contains("Cpu") || device_str.contains("Cuda"));
        assert_eq!(model.get_step(), 0);
        assert_eq!(model.get_learning_rate(), 1e-3);

        Ok(())
    }

    #[test]
    fn test_tft_learning_rate_validation() -> anyhow::Result<()> {
        let config = TFTConfig {
            input_dim: 225,
            hidden_dim: 32,
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 210, // 225 - 5 - 10 = 210
            ..Default::default()
        };
        let mut model = TrainableTFT::new(config)?;

        // Valid learning rate
        assert!(model.set_learning_rate(5e-4).is_ok());
        assert_eq!(model.get_learning_rate(), 5e-4);

        // Invalid learning rates
        assert!(model.set_learning_rate(0.0).is_err());
        assert!(model.set_learning_rate(-0.1).is_err());
        assert!(model.set_learning_rate(1.5).is_err());

        Ok(())
    }

    #[test]
    fn test_tft_metrics_collection() -> anyhow::Result<()> {
        let config = TFTConfig {
            input_dim: 225,
            hidden_dim: 32,
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 210, // 225 - 5 - 10 = 210
            ..Default::default()
        };
        let model = TrainableTFT::new(config)?;

        let metrics = model.collect_metrics();

        // Check standardized metrics
        assert!(metrics.loss >= 0.0);
        assert_eq!(metrics.learning_rate, model.get_learning_rate());
        assert!(!metrics.custom_metrics.is_empty());

        // Check TFT-specific metrics
        assert!(metrics.custom_metrics.contains_key("step_count"));
        assert!(metrics.custom_metrics.contains_key("num_parameters"));

        Ok(())
    }

    #[test]
    fn test_tft_checkpoint_save_load() -> anyhow::Result<()> {
        let config = TFTConfig {
            input_dim: 64,
            hidden_dim: 32,
            num_heads: 4,
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 49, // 64 - 5 - 10 = 49
            ..Default::default()
        };
        let model = TrainableTFT::new(config.clone())?;

        // Create temporary checkpoint directory
        let temp_dir = tempfile::tempdir()?;
        let checkpoint_path = temp_dir.path().join("tft_test_checkpoint");
        let checkpoint_path_str = checkpoint_path.to_str().unwrap();

        // Save checkpoint
        let saved_path = model.save_checkpoint(checkpoint_path_str)?;
        assert!(saved_path.contains("tft_test_checkpoint"));

        // Verify checkpoint files exist
        assert!(std::path::Path::new(&format!("{}.safetensors", checkpoint_path_str)).exists());
        assert!(std::path::Path::new(&format!("{}.json", checkpoint_path_str)).exists());

        // Load checkpoint into new model
        let mut loaded_model = TrainableTFT::new(config)?;
        let metadata = loaded_model.load_checkpoint(checkpoint_path_str)?;

        // Verify metadata
        assert_eq!(metadata.model_type, "TFT");
        assert!(loaded_model.model.is_trained);
        assert_eq!(loaded_model.get_step(), model.get_step());

        Ok(())
    }

    #[test]
    fn test_tft_zero_grad() -> anyhow::Result<()> {
        let config = TFTConfig {
            input_dim: 225,
            hidden_dim: 32,
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 210, // 225 - 5 - 10 = 210
            ..Default::default()
        };
        let mut model = TrainableTFT::new(config)?;

        // Zero gradients should succeed even with no prior gradients
        model.zero_grad()?;

        Ok(())
    }

    #[test]
    fn test_tft_zero_grad_resets_norm() -> anyhow::Result<()> {
        let config = TFTConfig {
            input_dim: 225,
            hidden_dim: 32,
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 210, // 225 - 5 - 10 = 210
            ..Default::default()
        };
        let mut model = TrainableTFT::new(config)?;

        // Set a non-zero gradient norm to simulate post-backward state
        model.last_grad_norm = 1.5;
        assert_eq!(model.last_grad_norm, 1.5);

        // Zero gradients should reset gradient norm tracking
        model.zero_grad()?;
        assert_eq!(model.last_grad_norm, 0.0);

        // Multiple calls should be idempotent
        model.zero_grad()?;
        assert_eq!(model.last_grad_norm, 0.0);

        Ok(())
    }

    #[test]
    fn test_tft_zero_grad_with_training_simulation() -> anyhow::Result<()> {
        let config = TFTConfig {
            input_dim: 64,
            hidden_dim: 32,
            num_heads: 4,
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 49, // 64 - 5 - 10 = 49
            sequence_length: 10,
            prediction_horizon: 5,
            ..Default::default()
        };
        let mut model = TrainableTFT::new(config)?;

        // Create dummy input tensor
        let batch_size = 4;
        // static + historical (unknown only) * seq_len + future (known) * pred_horizon
        let total_dim = 5 + 49 * 10 + 10 * 5; // 5 + 490 + 50 = 545
        let input = Tensor::randn(0f32, 1.0, (batch_size, total_dim), model.device())?;
        let target = Tensor::randn(0f32, 1.0, (batch_size, 5), model.device())?;

        // Simulate training step
        let predictions = model.forward(&input)?;
        let loss = model.compute_loss(&predictions, &target)?;
        let grad_norm = model.backward(&loss)?;

        // Verify gradient norm was computed
        assert!(grad_norm > 0.0);
        assert_eq!(model.last_grad_norm, grad_norm);

        // Zero gradients before next iteration
        model.zero_grad()?;
        assert_eq!(model.last_grad_norm, 0.0);

        Ok(())
    }
}