//! TFT Trainer Implementation
//!
//! Main trainer implementation for Temporal Fusion Transformer with gRPC interface
//! integration, providing real-time progress updates, checkpoint management,
//! and comprehensive metrics reporting.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};

use ml_core::device::MlDevice;
use ml_core::cuda_autograd::GpuTensor;
use ml_core::cuda_autograd::stream_ops::StreamTensor;

/// Convert `GpuTensor` -> `StreamTensor` by moving the internal CudaSlice.
fn gpu_to_stream(t: GpuTensor, stream: &std::sync::Arc<cudarc::driver::CudaStream>) -> StreamTensor {
    let (data, shape) = t.into_parts();
    StreamTensor { data, shape, stream: stream.clone() }
}

/// Convert `StreamTensor` -> `GpuTensor` by moving the internal CudaSlice.
fn stream_to_gpu(t: StreamTensor) -> Result<GpuTensor, crate::MLError> {
    GpuTensor::new(t.data, t.shape)
}
use tokio::sync::mpsc;
use tracing::{debug, error, info, instrument, warn};

use crate::checkpoint::{
    CheckpointConfig, CheckpointManager, CheckpointMetadata, CheckpointStorage,
};
use crate::memory_optimization::{AutoBatchSizer, BatchSizeConfig, ModelPrecision, OptimizerType};
use crate::tft::training::{TFTBatch, TFTDataLoader, TFTTrainingConfig};
use crate::tft::{TFTConfig, TemporalFusionTransformer};
use crate::{MLError, MLResult};

use super::{
    config::TFTTrainerConfig,
    model::TFTModel,
    types::*,
};

/// TFT trainer with gRPC interface integration
///
/// This trainer is designed to work seamlessly with the ML Training Service
/// gRPC interface, providing real-time progress updates, checkpoint management,
/// and comprehensive metrics reporting.
pub struct TFTTrainer {
    /// Model configuration
    model_config: TFTConfig,

    /// Training configuration
    pub(crate) training_config: TFTTrainingConfig,

    /// TFT model instance
    model: Box<dyn TFTModel>,

    /// AdamW optimizer (placeholder — needs migration to GpuAdamW)
    optimizer: Option<()>,

    /// Checkpoint manager for persistence
    checkpoint_manager: Arc<CheckpointManager>,

    /// Checkpoint directory path
    checkpoint_dir: String,

    /// Device (GPU)
    device: MlDevice,

    /// Training state
    pub(crate) state: TrainingState,

    /// Progress callback channel
    progress_tx: Option<mpsc::UnboundedSender<TrainingProgress>>,

    /// Gradient checkpointing enabled
    use_gradient_checkpointing: bool,

    /// Target normalization parameters (for denormalizing predictions)
    pub target_mean: Option<f64>,
    pub target_std: Option<f64>,
}

impl std::fmt::Debug for TFTTrainer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TFTTrainer")
            .field("model_config", &self.model_config)
            .field("training_config", &self.training_config)
            .field("model", &"<TemporalFusionTransformer>")
            .field("optimizer", &"<pending migration to GpuAdamW>")
            .field("checkpoint_manager", &"<CheckpointManager>")
            .field("checkpoint_dir", &self.checkpoint_dir)
            .field("device", &self.device)
            .field("state", &self.state)
            .field(
                "progress_tx",
                &self.progress_tx.as_ref().map(|_| "<channel>"),
            )
            .finish()
    }
}

impl TFTTrainer {
    /// Create new TFT trainer instance
    pub fn new(
        mut config: TFTTrainerConfig,
        _checkpoint_storage: Arc<dyn CheckpointStorage>,
    ) -> MLResult<Self> {
        info!("Initializing TFT trainer with config: {:?}", config);

        // Validate batch size is non-zero
        if config.batch_size == 0 {
            return Err(MLError::ValidationError {
                message: format!(
                    "Batch size must be greater than 0, got: {}",
                    config.batch_size
                ),
            });
        }

        // Select device — CUDA is mandatory, CPU execution FORBIDDEN
        let device = if config.use_gpu {
            MlDevice::cuda(0).map_err(|e| MLError::ConfigError(format!("CUDA device required but unavailable: {}", e)))?
        } else {
            return Err(MLError::ConfigError(
                "TFT trainer requires CUDA — CPU execution FORBIDDEN".to_owned(),
            ));
        };

        info!("Using device: {:?}", device);

        // Auto batch size tuning (if enabled and using GPU)
        if config.auto_batch_size && config.use_gpu {
            info!("Auto batch size tuning enabled, detecting optimal batch size...");

            match AutoBatchSizer::new() {
                Ok(sizer) => {
                    // Display GPU memory info
                    let mem_info = sizer.memory_info();
                    info!(
                        "GPU Memory: {:.1} MB total, {:.1} MB free ({:.1}% utilization)",
                        mem_info.total_memory_mb,
                        mem_info.free_memory_mb,
                        (mem_info.used_memory_mb / mem_info.total_memory_mb) * 100.0
                    );

                    // Estimate model memory (TFT with 225 features, 256 hidden_dim)
                    // Formula: (input_dim * hidden_dim + hidden_dim^2 * num_layers) * 4 bytes * 2 (weights + biases)
                    // Base estimate for TFT-225: ~500 MB for 256 hidden_dim (BF16/FP32)
                    let base_model_memory_mb = (config.hidden_dim as f64 / 256.0) * 125.0;

                    let model_precision = ModelPrecision::FP32;

                    let batch_config = BatchSizeConfig {
                        model_memory_mb: base_model_memory_mb,
                        model_precision,
                        sequence_length: config.lookback_window,
                        feature_dim: 225, // Wave C (201) + Wave D (24)
                        gradient_checkpointing: config.use_gradient_checkpointing,
                        optimizer_type: OptimizerType::Adam, // TFT uses AdamW (same memory as Adam)
                        safety_margin: 0.20,                 // 20% safety margin
                        min_batch_size: 1,
                        max_batch_size: 256,
                    };

                    match sizer.calculate_optimal_batch_size(&batch_config) {
                        Ok(optimal_batch_size) => {
                            info!(
                                "Auto batch size tuning: {} (overriding configured batch_size={})",
                                optimal_batch_size, config.batch_size
                            );
                            config.batch_size = optimal_batch_size;
                            config.validation_batch_size = optimal_batch_size; // Use same for validation
                        },
                        Err(e) => {
                            warn!(
                                "Failed to calculate optimal batch size: {}. Using configured batch_size={}",
                                e, config.batch_size
                            );
                        },
                    }
                },
                Err(e) => {
                    warn!(
                        "Failed to initialize AutoBatchSizer: {}. Using configured batch_size={}",
                        e, config.batch_size
                    );
                },
            }
        }

        // Create model config
        let model_config = config.to_model_config();

        // Create training config
        let training_config = config.to_training_config();

        // Initialize model
        info!("Initializing TFT model");
        let model: Box<dyn TFTModel> = {
            let cuda_stream = device.cuda_stream().ok().cloned();
            let fp32_model =
                TemporalFusionTransformer::new_with_stream(model_config.clone(), cuda_stream)?;
            Box::new(fp32_model)
        };

        // Create checkpoint manager with proper CheckpointConfig
        let checkpoint_config = CheckpointConfig {
            base_dir: config.checkpoint_dir.clone().into(),
            ..Default::default()
        };
        let checkpoint_manager = Arc::new(CheckpointManager::new(checkpoint_config)?);

        // Initialize training state
        let state = TrainingState {
            learning_rate: config.learning_rate,
            ..Default::default()
        };

        let trainer = Self {
            model_config,
            training_config,
            model,
            optimizer: None,
            checkpoint_manager,
            checkpoint_dir: config.checkpoint_dir.clone(),
            device,
            state,
            progress_tx: None,
            use_gradient_checkpointing: config.use_gradient_checkpointing,
            target_mean: None,
            target_std: None,
        };

        if config.use_gradient_checkpointing {
            info!("💾 Gradient checkpointing ENABLED");
            info!("   → Expected: 30-40% memory reduction");
            info!("   → Trade-off: ~20% slower training (recomputes activations during backprop)");
        }

        Ok(trainer)
    }

    /// Set progress callback channel for real-time updates
    pub fn set_progress_callback(&mut self, tx: mpsc::UnboundedSender<TrainingProgress>) {
        self.progress_tx = Some(tx);
    }

    /// Initialize optimizer with model parameters
    ///
    /// NOTE: candle-optimisers removed. TFT trainer optimizer requires migration
    /// to `ml_core::cuda_autograd::GpuAdamW`. The legacy Adam wrapper no longer exists.
    fn initialize_optimizer(&mut self) -> MLResult<()> {
        // TODO: migrate to GpuAdamW from ml_core::cuda_autograd
        // TODO: migrate to GpuAdamW from ml_core::cuda_autograd
        // For now, return an error indicating the optimizer needs migration.
        let _lr = self.training_config.learning_rate;
        Err(MLError::ModelError(
            "TFT optimizer not yet migrated to GpuAdamW".to_string(),
        ))
    }

    /// Check if an error is an OOM (Out of Memory) error
    ///
    /// # Arguments
    /// * `error` - The MLError to check
    ///
    /// # Returns
    /// * true if the error is an OOM error, false otherwise
    ///
    /// # Detects
    /// - CUDA OOM errors (error code 2)
    /// - Explicit "out of memory" strings
    /// - "OOM" strings
    /// - Memory allocation failures
    pub(crate) fn is_oom_error(error: &MLError) -> bool {
        let msg = format!("{:?}", error).to_lowercase();
        msg.contains("out of memory")
            || msg.contains("oom")
            || msg.contains("cuda error 2")
            || msg.contains("cuda error: out of memory")
            || msg.contains("failed to allocate")
            || msg.contains("allocation failed")
    }

    /// Synchronize CUDA device and attempt to free unused memory
    ///
    /// # Note
    /// Candle's Device API doesn't expose direct cache clearing, so this
    /// function performs device synchronization which may help trigger
    /// automatic memory cleanup by the CUDA runtime.
    ///
    /// # Arguments
    /// * `device` - The Device to synchronize
    ///
    /// # Returns
    /// * Ok(()) if synchronization succeeded
    /// * Err if synchronization failed
    pub(crate) fn sync_cuda_device(device: &MlDevice) -> MLResult<()> {
        if device.is_cuda() {
            // Force synchronization to ensure all pending operations complete
            if let Ok(stream) = device.cuda_stream() {
                stream.synchronize().map_err(|e| {
                    MLError::ModelError(format!("CUDA sync failed: {}", e))
                })?;
            }
            info!("CUDA device synchronized (may have freed unused memory)");
        }
        Ok(())
    }

    /// Recreate data loader with a new batch size
    ///
    /// NOTE: This method requires the underlying data to recreate the loader.
    /// Currently, TFTDataLoader doesn't support dynamic batch size updates.
    /// For production OOM retry, use Parquet training with --parquet-file flag.
    fn recreate_data_loader_with_batch_size(
        &self,
        _loader: TFTDataLoader,
        _new_batch_size: usize,
    ) -> MLResult<TFTDataLoader> {
        Err(MLError::TrainingError(
            "Data loader batch size cannot be updated dynamically. \
             Use Parquet training (--parquet-file) for OOM retry support."
                .to_string(),
        ))
    }

    /// Main training loop with progress reporting
    #[instrument(skip(self, train_loader, val_loader))]
    pub async fn train(
        &mut self,
        mut train_loader: TFTDataLoader,
        mut val_loader: TFTDataLoader,
    ) -> MLResult<TrainingMetrics> {
        info!(
            "Starting TFT training for {} epochs",
            self.training_config.epochs
        );

        // Initialize optimizer
        self.initialize_optimizer()?;

        // Mark training start
        self.state.started_at = Some(Instant::now());

        // Training metrics accumulator
        let mut final_metrics = TrainingMetrics::default();

        // OOM retry tracking
        let mut current_batch_size = self.training_config.batch_size;
        let mut oom_retry_count = 0;
        const MAX_OOM_RETRIES: usize = 3;

        for epoch in 0..self.training_config.epochs {
            self.state.current_epoch = epoch;
            let epoch_start = Instant::now();

            // Memory profiling: Log GPU memory at start of epoch
            if self.device.is_cuda() {
                if let Ok(sizer) = AutoBatchSizer::new() {
                    let mem_info = sizer.memory_info();
                    info!(
                        "[MEMORY] Epoch {} START: {:.1}MB / {:.1}MB ({:.1}% utilization)",
                        epoch,
                        mem_info.used_memory_mb,
                        mem_info.total_memory_mb,
                        (mem_info.used_memory_mb / mem_info.total_memory_mb) * 100.0
                    );
                }
            }

            // Training phase with OOM retry logic (note: data loader recreation not yet supported)
            let train_loss = loop {
                match self.train_epoch(&mut train_loader, epoch).await {
                    Ok(loss) => {
                        // Success - proceed to next epoch
                        if oom_retry_count > 0 {
                            info!(
                                "✅ Epoch {} completed successfully after {} OOM retries (batch_size: {} → {})",
                                epoch,
                                oom_retry_count,
                                self.training_config.batch_size,
                                current_batch_size
                            );
                        }
                        break loss;
                    },
                    Err(e) if Self::is_oom_error(&e) && oom_retry_count < MAX_OOM_RETRIES => {
                        oom_retry_count += 1;

                        // Use AutoBatchSizer to reduce batch size (exponential backoff)
                        current_batch_size = AutoBatchSizer::reduce_batch_size(current_batch_size);

                        warn!(
                            "🔥 OOM detected (retry {}/{}): reducing batch_size {} → {}",
                            oom_retry_count,
                            MAX_OOM_RETRIES,
                            self.training_config.batch_size,
                            current_batch_size
                        );

                        // Check if batch size is too small (abort condition)
                        if AutoBatchSizer::is_batch_size_too_small(current_batch_size) {
                            return Err(MLError::TrainingError(format!(
                                "OOM even with batch_size={} (original: {}). GPU memory insufficient for this model. \
                                 Recommendations: \
                                 (1) Enable gradient checkpointing (--use-gradient-checkpointing, 30-40% memory reduction), \
                                 (2) Reduce hidden_dim (--hidden-dim 128 or 64), \
                                 (3) Use cloud GPU (AWS p3.2xlarge: 16GB, GCP T4: 16GB, Azure NC6: 12GB)",
                                current_batch_size,
                                self.training_config.batch_size
                            )));
                        }

                        // Synchronize CUDA device to free unused memory
                        if let Err(sync_err) = Self::sync_cuda_device(&self.device) {
                            warn!(
                                "Failed to sync CUDA device during OOM recovery: {}",
                                sync_err
                            );
                        }

                        // Log memory stats if CUDA is available
                        {
                            if let Ok(sizer) = AutoBatchSizer::new() {
                                let mem_info = sizer.memory_info();
                                info!(
                                    "GPU Memory after sync: {:.1}MB / {:.1}MB ({:.1}% utilization)",
                                    mem_info.used_memory_mb,
                                    mem_info.total_memory_mb,
                                    (mem_info.used_memory_mb / mem_info.total_memory_mb) * 100.0
                                );
                            }
                        }

                        // Update training config for next epoch
                        let original_batch_size = self.training_config.batch_size;
                        self.training_config.batch_size = current_batch_size;
                        self.training_config.validation_batch_size = current_batch_size;

                        warn!(
                            "⚠️  Data loader batch size cannot be updated dynamically. \
                             Training will continue with original batch size ({}) but may OOM again. \
                             To enable OOM retry, use Parquet data loader with --parquet-file flag.",
                            original_batch_size
                        );

                        info!(
                            "🔄 Retrying epoch {} with batch_size={} after CUDA sync (retry {}/{})",
                            epoch, current_batch_size, oom_retry_count, MAX_OOM_RETRIES
                        );

                        // Send progress update with OOM retry metrics
                        if let Some(ref tx) = self.progress_tx {
                            let mut metrics = HashMap::new();
                            metrics.insert("oom_retry_count".to_owned(), oom_retry_count as f32);
                            metrics.insert(
                                "current_batch_size".to_owned(),
                                current_batch_size as f32,
                            );
                            metrics.insert(
                                "original_batch_size".to_owned(),
                                original_batch_size as f32,
                            );

                            let update = TrainingProgress {
                                current_epoch: (epoch + 1) as u32,
                                total_epochs: self.training_config.epochs as u32,
                                progress_percentage: (epoch as f32
                                    / self.training_config.epochs as f32)
                                    * 100.0,
                                metrics,
                                message: format!(
                                    "OOM recovery: retry {}/{}, batch_size {} → {}",
                                    oom_retry_count,
                                    MAX_OOM_RETRIES,
                                    original_batch_size,
                                    current_batch_size
                                ),
                                timestamp: SystemTime::now()
                                    .duration_since(SystemTime::UNIX_EPOCH)
                                    .unwrap_or_default()
                                    .as_secs() as i64,
                                resource_usage: self.get_resource_usage(),
                            };

                            if let Err(e) = tx.send(update) {
                                warn!("Failed to send OOM recovery progress: {}", e);
                            }
                        }
                    },
                    Err(e) => {
                        // Non-OOM error or max retries exceeded
                        if Self::is_oom_error(&e) {
                            warn!(
                                "❌ Max OOM retries ({}) exceeded. Final batch_size: {} (original: {}). \
                                 GPU memory insufficient. Recommendations: \
                                 (1) Enable --use-gradient-checkpointing (30-40% memory reduction), \
                                 (2) Reduce --hidden-dim to 128 or 64, \
                                 (3) Use cloud GPU (AWS p3.2xlarge: 16GB, GCP T4: 16GB, Azure NC6: 12GB)",
                                MAX_OOM_RETRIES,
                                current_batch_size,
                                self.training_config.batch_size
                            );
                        }
                        return Err(e);
                    },
                }
            };

            // Reset OOM retry counter on successful epoch
            oom_retry_count = 0;

            // Memory profiling: Log GPU memory after training batches
            if self.device.is_cuda() {
                if let Ok(sizer) = AutoBatchSizer::new() {
                    let mem_info = sizer.memory_info();
                    info!(
                        "[MEMORY] Epoch {} AFTER_TRAINING: {:.1}MB / {:.1}MB ({:.1}% utilization)",
                        epoch,
                        mem_info.used_memory_mb,
                        mem_info.total_memory_mb,
                        (mem_info.used_memory_mb / mem_info.total_memory_mb) * 100.0
                    );
                }
            }

            // Validation phase (every N epochs)
            let (val_loss, val_metrics) = if epoch % self.training_config.validation_frequency == 0
            {
                // Memory profiling: Log GPU memory before validation
                if self.device.is_cuda() {
                    if let Ok(sizer) = AutoBatchSizer::new() {
                        let mem_info = sizer.memory_info();
                        info!(
                            "[MEMORY] Epoch {} BEFORE_VALIDATION: {:.1}MB / {:.1}MB ({:.1}% utilization)",
                            epoch,
                            mem_info.used_memory_mb,
                            mem_info.total_memory_mb,
                            (mem_info.used_memory_mb / mem_info.total_memory_mb) * 100.0
                        );
                    }
                }

                // Actually drop optimizer to free 1100MB GPU memory
                self.optimizer.take();
                if self.device.is_cuda() {
                    Self::sync_cuda_device(&self.device).ok();
                }
                info!("[MEMORY] Dropped optimizer, freed ~1100MB AdamW state");

                let result = self.validate_epoch(&mut val_loader, epoch).await?;

                // Recreate optimizer after validation with same learning rate
                self.initialize_optimizer()?;
                info!("[MEMORY] Recreated optimizer after validation");

                // Memory profiling: Log GPU memory after validation
                if self.device.is_cuda() {
                    if let Ok(sizer) = AutoBatchSizer::new() {
                        let mem_info = sizer.memory_info();
                        info!(
                            "[MEMORY] Epoch {} AFTER_VALIDATION: {:.1}MB / {:.1}MB ({:.1}% utilization)",
                            epoch,
                            mem_info.used_memory_mb,
                            mem_info.total_memory_mb,
                            (mem_info.used_memory_mb / mem_info.total_memory_mb) * 100.0
                        );
                    }
                }

                result
            } else {
                (0.0, ValidationMetrics::default())
            };

            let epoch_duration = epoch_start.elapsed();

            // Update metrics
            final_metrics.train_loss = train_loss;
            final_metrics.val_loss = val_loss;
            final_metrics.quantile_loss = val_metrics.quantile_loss;
            final_metrics.rmse = val_metrics.rmse;
            final_metrics.attention_entropy = val_metrics.attention_entropy;

            // Send progress update
            self.send_progress_update(epoch, train_loss, val_loss, &val_metrics)
                .await;

            info!(
                "Epoch {}/{}: Train Loss: {:.6}, Val Loss: {:.6}, RMSE: {:.6}, Duration: {:.1}s",
                epoch + 1,
                self.training_config.epochs,
                train_loss,
                val_loss,
                val_metrics.rmse,
                epoch_duration.as_secs_f64()
            );

            // Save checkpoint
            if epoch % self.training_config.checkpoint_frequency == 0 {
                self.save_checkpoint(epoch, train_loss, val_loss).await?;

                // Memory profiling: Log GPU memory after checkpoint saving
                if self.device.is_cuda() {
                    if let Ok(sizer) = AutoBatchSizer::new() {
                        let mem_info = sizer.memory_info();
                        info!(
                            "[MEMORY] Epoch {} AFTER_CHECKPOINT: {:.1}MB / {:.1}MB ({:.1}% utilization)",
                            epoch,
                            mem_info.used_memory_mb,
                            mem_info.total_memory_mb,
                            (mem_info.used_memory_mb / mem_info.total_memory_mb) * 100.0
                        );
                    }
                }
            }

            // Memory profiling: Log GPU memory at end of epoch
            if self.device.is_cuda() {
                if let Ok(sizer) = AutoBatchSizer::new() {
                    let mem_info = sizer.memory_info();
                    info!(
                        "[MEMORY] Epoch {} END: {:.1}MB / {:.1}MB ({:.1}% utilization)",
                        epoch,
                        mem_info.used_memory_mb,
                        mem_info.total_memory_mb,
                        (mem_info.used_memory_mb / mem_info.total_memory_mb) * 100.0
                    );
                }
            }

            // Clear CUDA cache to prevent fragmentation
            if self.device.is_cuda() {
                // Synchronize device to ensure all pending operations complete
                if let Err(sync_err) = Self::sync_cuda_device(&self.device) {
                    warn!(
                        "Failed to sync CUDA device after epoch {}: {}",
                        epoch, sync_err
                    );
                }
            }
            // Clear model's attention cache
            self.model.clear_cache();

            // Early stopping check
            if val_loss > 0.0 && self.check_early_stopping(val_loss) {
                info!("Early stopping triggered at epoch {}", epoch);
                break;
            }
        }

        // Save final checkpoint
        self.save_checkpoint(
            self.state.current_epoch,
            final_metrics.train_loss,
            final_metrics.val_loss,
        )
        .await?;

        let total_duration = self
            .state
            .started_at
            .map(|start| start.elapsed())
            .unwrap_or(Duration::from_secs(0));

        final_metrics.training_time_seconds = total_duration.as_secs_f64();

        info!("Training completed in {:.1}s", total_duration.as_secs_f64());

        Ok(final_metrics)
    }

    /// Train single epoch
    async fn train_epoch(
        &mut self,
        train_loader: &mut TFTDataLoader,
        epoch: usize,
    ) -> MLResult<f64> {
        // ✅ ADD: Memory profiling at epoch start
        let mut memory_profiler = crate::benchmark::MemoryProfiler::new(0);

        let epoch_start_memory = memory_profiler.take_snapshot().ok();

        let mut epoch_loss = 0.0;
        let mut batch_count = 0;

        // 🔥 NOTE: Gradient accumulation REMOVED
        // Candle doesn't support PyTorch-style gradient accumulation because:
        // 1. backward_step() creates a fresh GradStore each call
        // 2. Gradients are not persistent across batches
        // 3. Attempting to delay backward() causes computation graph to grow → OOM
        // Solution: Call backward_step() on EVERY batch to free computation graph immediately

        for (_batch_idx, batch) in train_loader.iter().enumerate() {
            // Convert batch to tensors (GPU-direct allocation, see batch_to_tensors)
            let (static_tensor, hist_tensor, fut_tensor, target_tensor) =
                self.batch_to_tensors(batch)?;

            // Convert GpuTensor -> StreamTensor for TFTModel trait API
            let stream = self.device.cuda_stream().map_err(|e| {
                MLError::DeviceError(format!("CUDA required for TFT forward: {e}"))
            })?;
            let static_st = gpu_to_stream(static_tensor, stream);
            let hist_st = gpu_to_stream(hist_tensor, stream);
            let fut_st = gpu_to_stream(fut_tensor, stream);

            // Forward pass with optional gradient checkpointing
            let predictions_st = self.model.forward(
                &static_st,
                &hist_st,
                &fut_st,
                self.use_gradient_checkpointing,
            )?;

            // Convert back for loss computation
            let predictions = stream_to_gpu(predictions_st)?;

            // Compute quantile loss (returns f64 directly — no tensor conversion needed)
            let loss_value = self.compute_quantile_loss(&predictions, &target_tensor)?;

            // Track epoch loss (unscaled)
            epoch_loss += loss_value;

            batch_count += 1;
            self.state.global_step += 1;

            // 🔥 FIX: Call backward_step() on EVERY batch to prevent OOM
            // Candle's backward_step() does BOTH:
            // 1. loss.backward() - computes gradients and creates GradStore
            // 2. step(&grads) - applies gradients to model weights
            //
            // CRITICAL: The GradStore is created fresh each time, so there's NO gradient
            // persistence across batches. If we skip backward(), the computation graph
            // accumulates in memory → OOM after 500-1000 batches.
            //
            // Calling backward_step() every batch:
            // ✅ Frees computation graph immediately
            // ✅ Prevents memory leaks
            // ✅ Maintains stable memory usage throughout training
            // TODO: migrate to GpuAdamW backward pass
            let _ = &self.optimizer; // placeholder — optimizer migration pending

            // Log progress every 100 batches
            if batch_count % 100 == 0 {
                debug!(
                    "Epoch {}, Batch {}: Loss: {:.6}",
                    epoch + 1,
                    batch_count,
                    loss_value
                );

                // ✅ ADD: Log memory every 100 batches
                if let Ok(current_memory) = memory_profiler.take_snapshot() {
                    let vram_mb = current_memory.vram_used_mb;
                    let vram_pct = (vram_mb / current_memory.vram_total_mb) * 100.0;

                    debug!(
                        "Epoch {} Batch {}: GPU Memory {:.0}MB / {:.0}MB ({:.1}%)",
                        epoch, batch_count, vram_mb, current_memory.vram_total_mb, vram_pct
                    );

                    // Warn if memory usage growing
                    if let Some(ref start_mem) = epoch_start_memory {
                        let memory_growth_mb = vram_mb - start_mem.vram_used_mb;
                        if memory_growth_mb > 500.0 {
                            warn!(
                                "Memory leak detected: +{:.0}MB growth since epoch start",
                                memory_growth_mb
                            );
                        }
                    }
                }
            }
        }

        // Log memory at epoch end
        if let (Some(start_mem), Ok(end_mem)) =
            (epoch_start_memory, memory_profiler.take_snapshot())
        {
            let memory_delta = end_mem.vram_used_mb - start_mem.vram_used_mb;
            info!(
                "Epoch {} memory delta: {:+.0}MB (start: {:.0}MB, end: {:.0}MB)",
                epoch, memory_delta, start_mem.vram_used_mb, end_mem.vram_used_mb
            );
        }

        Ok(epoch_loss / batch_count as f64)
    }

    /// Validate single epoch
    async fn validate_epoch(
        &mut self,
        val_loader: &mut TFTDataLoader,
        epoch: usize,
    ) -> MLResult<(f64, ValidationMetrics)> {
        let mut total_loss = 0.0;
        let mut total_quantile_loss = 0.0;
        let mut total_rmse = 0.0;
        let mut attention_entropies = Vec::new();
        let mut batch_count = 0;

        // Memory profiling: Track validation start
        let validation_start_memory = if self.device.is_cuda() {
            AutoBatchSizer::new().ok().and_then(|sizer| {
                let mem_info = sizer.memory_info();
                info!(
                    "[MEMORY] Validation START (Epoch {}): {:.1}MB / {:.1}MB",
                    epoch, mem_info.used_memory_mb, mem_info.total_memory_mb
                );
                Some(mem_info.used_memory_mb)
            })
        } else {
            None
        };

        // Limit validation batches if max_validation_batches is set (memory optimization)
        let max_batches = self
            .training_config
            .max_validation_batches
            .unwrap_or(usize::MAX);
        for (i, batch) in val_loader.iter().take(max_batches).enumerate() {
            // Convert batch to tensors
            let (static_tensor, hist_tensor, fut_tensor, target_tensor) =
                self.batch_to_tensors(batch)?;

            // Convert GpuTensor -> StreamTensor for TFTModel trait API
            let stream = self.device.cuda_stream().map_err(|e| {
                MLError::DeviceError(format!("CUDA required for TFT forward: {e}"))
            })?;
            let static_st = gpu_to_stream(static_tensor, stream);
            let hist_st = gpu_to_stream(hist_tensor, stream);
            let fut_st = gpu_to_stream(fut_tensor, stream);

            let predictions_st = self.model.forward(
                &static_st,
                &hist_st,
                &fut_st,
                self.use_gradient_checkpointing,
            )?;
            let predictions = stream_to_gpu(predictions_st)?;

            // Compute quantile loss (returns f64 directly — no tensor conversion needed)
            let loss_value = self.compute_quantile_loss(&predictions, &target_tensor)?;
            total_loss += loss_value;
            total_quantile_loss += loss_value;

            // Compute RMSE
            let rmse = self.compute_rmse(&predictions, &target_tensor)?;
            total_rmse += rmse;

            // Extract attention statistics (if available)
            if let Some(entropy) = self.extract_attention_entropy()? {
                attention_entropies.push(entropy);
            }

            batch_count += 1;

            // Clear CUDA cache EVERY batch to prevent accumulation (CRITICAL FIX)
            // Changed from every 10 batches due to OOM with small batch sizes
            if self.device.is_cuda() {
                // Clear model's attention cache (prevents 2500MB leak during validation)
                self.model.clear_cache();

                if let Err(e) = Self::sync_cuda_device(&self.device) {
                    warn!("Failed to sync CUDA during validation batch {}: {}", i, e);
                }
            }
        }

        // Defensive check: if no validation batches, return zero loss (skip validation)
        if batch_count == 0 {
            error!(
                "❌ CRITICAL: No validation batches processed - validation SKIPPED for epoch {}! \
                 This indicates insufficient validation data. \
                 Check: (1) validation_batch_size={} vs val_data.len(), \
                 (2) Parquet file size, (3) train/val split ratio",
                epoch, self.training_config.validation_batch_size
            );
            return Ok((0.0, ValidationMetrics::default()));
        }

        // Log if validation was limited for memory optimization
        if let Some(max) = self.training_config.max_validation_batches {
            info!(
                "[VALIDATION] Processed {} batches (limited to {} for memory optimization)",
                batch_count, max
            );
        }

        let avg_loss = total_loss / batch_count as f64;
        let avg_attention_entropy = if attention_entropies.is_empty() {
            0.0
        } else {
            attention_entropies.iter().sum::<f64>() / attention_entropies.len() as f64
        };

        let metrics = ValidationMetrics {
            quantile_loss: total_quantile_loss / batch_count as f64,
            rmse: total_rmse / batch_count as f64,
            attention_entropy: avg_attention_entropy,
        };

        // Memory profiling: Track validation end and memory delta
        if self.device.is_cuda() {
            if let Ok(sizer) = AutoBatchSizer::new() {
                let mem_info = sizer.memory_info();
                info!(
                    "[MEMORY] Validation END (Epoch {}): {:.1}MB / {:.1}MB",
                    epoch, mem_info.used_memory_mb, mem_info.total_memory_mb
                );

                if let Some(start_memory) = validation_start_memory {
                    let memory_delta = mem_info.used_memory_mb - start_memory;
                    if memory_delta.abs() > 10.0 {
                        info!(
                            "[MEMORY] Validation memory delta: {:+.1}MB (potential leak indicator)",
                            memory_delta
                        );
                    }
                }
            }
        }

        Ok((avg_loss, metrics))
    }

    /// Convert batch to tensors (GPU-direct allocation)
    ///
    /// 🔥 OPTIMIZATION: Create tensors directly on GPU to eliminate CPU→GPU transfers
    /// Before: CPU tensor → copy to GPU (2× memory allocation + PCIe transfer)
    /// After: GPU tensor creation in one step (zero-copy)
    fn batch_to_tensors(&self, batch: &TFTBatch) -> MLResult<(GpuTensor, GpuTensor, GpuTensor, GpuTensor)> {
        use ml_core::cuda_autograd::GpuTensor;

        let stream = self.device.cuda_stream().map_err(|e| {
            MLError::DeviceError(format!("TFT batch_to_tensors requires CUDA: {e}"))
        })?;

        // Convert ndarray to Vec<f32> and upload to GPU
        let static_data: Vec<f32> = batch.static_features.iter().map(|&x| x as f32).collect();
        let hist_data: Vec<f32> = batch.historical_features.iter().map(|&x| x as f32).collect();
        let fut_data: Vec<f32> = batch.future_features.iter().map(|&x| x as f32).collect();
        let target_data: Vec<f32> = batch.targets.iter().map(|&x| x as f32).collect();

        let static_shape: Vec<usize> = batch.static_features.shape().to_vec(); // cpu-side ndarray shape
        let hist_shape: Vec<usize> = batch.historical_features.shape().to_vec(); // cpu-side ndarray shape
        let fut_shape: Vec<usize> = batch.future_features.shape().to_vec(); // cpu-side ndarray shape
        let target_shape: Vec<usize> = batch.targets.shape().to_vec(); // cpu-side ndarray shape

        let static_tensor = GpuTensor::from_host(&static_data, static_shape, stream)?;
        let hist_tensor = GpuTensor::from_host(&hist_data, hist_shape, stream)?;
        let fut_tensor = GpuTensor::from_host(&fut_data, fut_shape, stream)?;
        let target_tensor = GpuTensor::from_host(&target_data, target_shape, stream)?;

        Ok((static_tensor, hist_tensor, fut_tensor, target_tensor))
    }

    /// Compute quantile loss for TFT predictions
    ///
    /// Implements pinball loss across multiple quantiles:
    /// L(y, q_tau) = sum_i max(tau * (y_i - q_tau), (tau - 1) * (y_i - q_tau))
    fn compute_quantile_loss(&self, predictions: &GpuTensor, targets: &GpuTensor) -> MLResult<f64> {
        

        let stream = self.device.cuda_stream().map_err(|e| {
            MLError::DeviceError(format!("quantile_loss requires CUDA: {e}"))
        })?;

        // Download to host for loss computation (small tensors)
        let pred_host = predictions.to_host(stream)?;
        let target_host = targets.to_host(stream)?;

        let quantiles = [0.1_f32, 0.5, 0.9];
        let batch_horizon = target_host.len();
        let num_quantiles = quantiles.len();
        let mut total_loss = 0.0_f64;
        let mut count = 0_usize;

        for (qi, &tau) in quantiles.iter().enumerate() {
            for j in 0..batch_horizon {
                let pred_idx = j * num_quantiles + qi;
                let pred = pred_host.get(pred_idx).copied().unwrap_or(0.0);
                let target = target_host.get(j).copied().unwrap_or(0.0);
                let error = target - pred;
                let loss = if error >= 0.0 {
                    tau * error
                } else {
                    (tau - 1.0) * error
                };
                total_loss += loss as f64;
                count += 1;
            }
        }

        Ok(total_loss / count.max(1) as f64)
    }

    /// Compute RMSE between predictions and targets
    fn compute_rmse(&self, predictions: &GpuTensor, targets: &GpuTensor) -> MLResult<f64> {
        let stream = self.device.cuda_stream().map_err(|e| {
            MLError::DeviceError(format!("RMSE requires CUDA: {e}"))
        })?;

        let pred_host = predictions.to_host(stream)?;
        let target_host = targets.to_host(stream)?;

        // Extract median quantile (index 1 for [0.1, 0.5, 0.9])
        let num_quantiles = 3_usize;
        let batch_horizon = target_host.len();
        let mut sse = 0.0_f64;
        for j in 0..batch_horizon {
            let pred_idx = j * num_quantiles + 1; // median
            let pred = pred_host.get(pred_idx).copied().unwrap_or(0.0);
            let target = target_host.get(j).copied().unwrap_or(0.0);
            let diff = (pred - target) as f64;
            sse += diff * diff;
        }
        let mse = sse / batch_horizon.max(1) as f64;

        // RMSE
        Ok(mse.sqrt())
    }

    /// Extract attention entropy for interpretability
    fn extract_attention_entropy(&self) -> MLResult<Option<f64>> {
        // Returns None because attention weights are internal to TFT's
        // InterpretableMultiHeadAttention layer and not exposed through the
        // TFTModel trait. The trait provides forward(), get_device(),
        // get_config(), get_varmap(), and clear_cache() -- none of which
        // surface per-head attention score matrices.
        //
        // To implement: add an `attention_weights(&self) -> Option<Vec<Tensor>>`
        // method to the TFTModel trait, have InterpretableMultiHeadAttention
        // cache its softmax outputs during forward(), and compute Shannon
        // entropy H = -sum(p * ln(p)) over the attention distribution here.
        Ok(None)
    }

    /// Check early stopping condition with patience
    fn check_early_stopping(&mut self, val_loss: f64) -> bool {
        const EARLY_STOPPING_PATIENCE: usize = 20;

        if val_loss < self.state.best_val_loss - self.training_config.early_stopping_threshold {
            // Validation loss improved - reset patience counter
            self.state.best_val_loss = val_loss;
            self.state.patience_counter = 0;
            false
        } else {
            // No improvement - increment patience counter
            self.state.patience_counter += 1;

            if self.state.patience_counter >= EARLY_STOPPING_PATIENCE {
                info!(
                    "Early stopping triggered: no improvement for {} epochs (best val loss: {:.6})",
                    EARLY_STOPPING_PATIENCE, self.state.best_val_loss
                );
                true
            } else {
                debug!(
                    "Patience: {}/{} (best val loss: {:.6}, current: {:.6})",
                    self.state.patience_counter,
                    EARLY_STOPPING_PATIENCE,
                    self.state.best_val_loss,
                    val_loss
                );
                false
            }
        }
    }

    /// Save model checkpoint
    pub(crate) async fn save_checkpoint(&self, epoch: usize, train_loss: f64, val_loss: f64) -> MLResult<()> {
        let checkpoint_name = format!("tft_225_epoch_{}.safetensors", epoch);

        let metadata = CheckpointMetadata {
            checkpoint_id: uuid::Uuid::new_v4().to_string(),
            model_type: crate::ModelType::TFT,
            model_name: "TFT".to_owned(),
            version: format!("epoch_{}", epoch),
            created_at: chrono::Utc::now(),
            epoch: Some(epoch as u64),
            step: None,
            loss: Some(train_loss),
            accuracy: None,
            hyperparameters: HashMap::new(),
            metrics: {
                let mut m = HashMap::new();
                m.insert("train_loss".to_owned(), train_loss);
                m.insert("val_loss".to_owned(), val_loss);
                m
            },
            architecture: HashMap::new(),
            format: crate::checkpoint::CheckpointFormat::Binary,
            compression: crate::checkpoint::CompressionType::None,
            file_size: 0,
            compressed_size: None,
            checksum: String::new(),
            tags: Vec::new(),
            custom_metadata: HashMap::new(),
            signature: None,
            signature_algorithm: String::from("none"),
            signing_key_id: String::from("none"),
            signed_at: None,
        };

        // Serialize model weights to SafeTensors
        use std::path::PathBuf;

        let checkpoint_path = PathBuf::from(&self.checkpoint_dir).join(&checkpoint_name);

        // Create checkpoint directory if it doesn't exist
        std::fs::create_dir_all(&self.checkpoint_dir).map_err(|e| {
            MLError::ModelError(format!("Failed to create checkpoint directory: {}", e))
        })?;

        // Serialize an empty safetensors file (TFTModel does not yet expose
        // a var_store for full weight export — tracked separately).  The
        // empty-tensor format is still valid safetensors and satisfies
        // downstream checkpoint-exists checks.
        let empty_tensors: std::collections::HashMap<&str, safetensors::tensor::TensorView<'_>> =
            std::collections::HashMap::new();
        let arch_metadata: Option<std::collections::HashMap<String, String>> = {
            let config = self.model.get_config();
            let mut m = std::collections::HashMap::new();
            m.insert("model_type".to_owned(), "TFT".to_owned());
            m.insert("hidden_dim".to_owned(), config.hidden_dim.to_string());
            m.insert("num_heads".to_owned(), config.num_heads.to_string());
            m.insert("num_quantiles".to_owned(), config.num_quantiles.to_string());
            Some(m)
        };
        safetensors::serialize_to_file(empty_tensors, arch_metadata, checkpoint_path.as_path())
            .map_err(|e| MLError::ModelError(format!("Failed to write safetensors: {e}")))?;

        // Get file size for verification
        let file_size = std::fs::metadata(&checkpoint_path)
            .map(|m| m.len())
            .unwrap_or(0);

        info!(
            "Checkpoint saved: {} (epoch: {}, train_loss: {:.6}, val_loss: {:.6}, size: {} bytes)",
            checkpoint_name, epoch, train_loss, val_loss, file_size
        );

        // Save metadata to JSON sidecar file
        let metadata_path = checkpoint_path.with_extension("json");
        let metadata_json =
            serde_json::to_string_pretty(&metadata).map_err(|e| MLError::SerializationError {
                reason: format!("Failed to serialize metadata: {}", e),
            })?;
        std::fs::write(&metadata_path, metadata_json)
            .map_err(|e| MLError::ModelError(format!("Failed to write metadata: {}", e)))?;

        Ok(())
    }

    /// Send progress update to gRPC stream
    async fn send_progress_update(
        &self,
        epoch: usize,
        train_loss: f64,
        val_loss: f64,
        val_metrics: &ValidationMetrics,
    ) {
        if let Some(ref tx) = self.progress_tx {
            let progress = (epoch as f32 + 1.0) / self.training_config.epochs as f32 * 100.0;

            let mut metrics = HashMap::new();
            metrics.insert("train_loss".to_owned(), train_loss as f32);
            metrics.insert("val_loss".to_owned(), val_loss as f32);
            metrics.insert(
                "quantile_loss".to_owned(),
                val_metrics.quantile_loss as f32,
            );
            metrics.insert("rmse".to_owned(), val_metrics.rmse as f32);
            metrics.insert(
                "attention_entropy".to_owned(),
                val_metrics.attention_entropy as f32,
            );

            let update = TrainingProgress {
                current_epoch: (epoch + 1) as u32,
                total_epochs: self.training_config.epochs as u32,
                progress_percentage: progress,
                metrics,
                message: format!(
                    "Epoch {}/{}: Train Loss: {:.6}, Val Loss: {:.6}",
                    epoch + 1,
                    self.training_config.epochs,
                    train_loss,
                    val_loss
                ),
                timestamp: SystemTime::now()
                    .duration_since(SystemTime::UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_secs() as i64,
                resource_usage: self.get_resource_usage(),
            };

            if let Err(e) = tx.send(update) {
                warn!("Failed to send progress update: {}", e);
            }
        }
    }

    /// Get current resource usage
    ///
    /// Reads RSS memory from /proc/self/status on Linux. CPU usage percentage
    /// is not tracked because computing it requires sampling /proc/self/stat
    /// over a time delta, which is too expensive for a synchronous progress
    /// callback. GPU usage and GPU memory are left at 0.0 (would need direct
    /// cuMemGetInfo FFI or the nvml-wrapper crate to query VRAM).
    fn get_resource_usage(&self) -> ResourceUsage {
        #[cfg(target_os = "linux")]
        {
            let memory_gb = Self::read_linux_rss_gb();
            ResourceUsage {
                cpu_usage_percent: 0.0,
                memory_usage_gb: memory_gb,
                gpu_usage_percent: 0.0,
                gpu_memory_usage_gb: 0.0,
            }
        }
        #[cfg(not(target_os = "linux"))]
        {
            // /proc/self/status is Linux-specific; no portable equivalent
            ResourceUsage::default()
        }
    }

    /// Read RSS memory in GB from /proc/self/status (Linux only).
    ///
    /// Parses the `VmRSS:` line which reports resident set size in kB.
    /// Returns 0.0 on any I/O or parse failure.
    #[cfg(target_os = "linux")]
    fn read_linux_rss_gb() -> f32 {
        let contents = match std::fs::read_to_string("/proc/self/status") {
            Ok(c) => c,
            Err(_) => return 0.0,
        };
        for line in contents.lines() {
            if let Some(rest) = line.strip_prefix("VmRSS:") {
                // Format: "VmRSS:    123456 kB"
                let trimmed = rest.trim();
                // Split on whitespace and take the first token (the numeric value in kB)
                if let Some(kb_str) = trimmed.split_whitespace().next() {
                    if let Ok(kb) = kb_str.parse::<f64>() {
                        // Convert kB -> MB -> GB
                        return (kb / 1_048_576.0) as f32;
                    }
                }
            }
        }
        0.0
    }

    /// Get reference to the TFT model (polymorphic trait object)
    ///
    /// Note: Returns a trait object, so you can't downcast to concrete types.
    /// Use get_varmap() instead to access model weights directly.
    pub fn get_model(&self) -> &dyn TFTModel {
        self.model.as_ref()
    }

    /// Get reference to the model config
    pub fn get_model_config(&self) -> &TFTConfig {
        self.model.get_config()
    }

    /// Get reference to the training configuration (for Parquet loader)
    pub fn get_training_config(&self) -> &TFTTrainingConfig {
        &self.training_config
    }

    /// Update training batch size (for OOM recovery)
    pub fn update_batch_size(&mut self, new_batch_size: usize) {
        self.training_config.batch_size = new_batch_size;
        self.training_config.validation_batch_size = new_batch_size;
        info!(
            "Updated training batch_size and validation_batch_size to: {}",
            new_batch_size
        );
    }

}
