//! TFT Model Trait
//!
//! This module provides a trait-based abstraction for the TFT model used by TFTTrainer.

use ml_core::cuda_autograd::StreamTensor;

use crate::tft::{TFTConfig, TemporalFusionTransformer};
use crate::MLError;

/// Trait for TFT model abstraction
pub trait TFTModel: Send + Sync {
    /// Forward pass with optional gradient checkpointing
    ///
    /// # Arguments
    /// * `static_features` - Static features [batch, num_static_features]
    /// * `historical_ts` - Historical time series [batch, seq_len, num_unknown_features]
    /// * `future_ts` - Future time series [batch, horizon, num_known_features]
    /// * `use_checkpointing` - Enable gradient checkpointing (trades compute for memory)
    ///
    /// # Returns
    /// * Quantile predictions [batch, horizon, num_quantiles]
    fn forward(
        &mut self,
        static_features: &StreamTensor,
        historical_ts: &StreamTensor,
        future_ts: &StreamTensor,
        use_checkpointing: bool,
    ) -> Result<StreamTensor, MLError>;

    /// Get configuration
    fn get_config(&self) -> &TFTConfig;

    /// Clear attention cache to free memory
    /// Call this after training/inference batch to prevent memory accumulation
    fn clear_cache(&mut self);
}

/// Implement TFTModel for TemporalFusionTransformer
impl TFTModel for TemporalFusionTransformer {
    fn forward(
        &mut self,
        static_features: &StreamTensor,
        historical_ts: &StreamTensor,
        future_ts: &StreamTensor,
        use_checkpointing: bool,
    ) -> Result<StreamTensor, MLError> {
        self.forward_with_checkpointing(
            static_features,
            historical_ts,
            future_ts,
            use_checkpointing,
        )
    }

    fn get_config(&self) -> &TFTConfig {
        &self.config
    }

    fn clear_cache(&mut self) {
        // No-op: TemporalFusionTransformer doesn't expose a public clear_cache method
        // The attention cache is managed internally by TemporalSelfAttention
        // CUDA cache clearing is handled separately by sync_cuda_device()
    }
}
