//! TFT Trainer Configuration
//!
//! Configuration types for the Temporal Fusion Transformer trainer, including
//! training hyperparameters and memory optimization settings.

use serde::{Deserialize, Serialize};

use crate::tft::training::TFTTrainingConfig;
use crate::tft::TFTConfig;

/// TFT trainer configuration from gRPC proto
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TFTTrainerConfig {
    /// Number of epochs
    pub epochs: usize,

    /// Learning rate
    pub learning_rate: f64,

    /// Batch size (overridden if auto_batch_size is true)
    pub batch_size: usize,

    /// Auto-detect optimal batch size based on available GPU memory
    pub auto_batch_size: bool,

    /// Hidden dimension (128, 256, 512)
    pub hidden_dim: usize,

    /// Number of attention heads (4, 8, 16)
    pub num_attention_heads: usize,

    /// Dropout rate (0.0-0.3)
    pub dropout_rate: f64,

    /// Number of LSTM layers
    pub lstm_layers: usize,

    /// Quantiles for quantile regression [0.1, 0.5, 0.9]
    pub quantiles: Vec<f64>,

    /// Lookback window length
    pub lookback_window: usize,

    /// Forecast horizon
    pub forecast_horizon: usize,

    /// Use GPU
    pub use_gpu: bool,

    /// Enable gradient checkpointing (trades compute for memory, 30-40% reduction)
    pub use_gradient_checkpointing: bool,

    /// Validation batch size
    pub validation_batch_size: usize,

    /// Maximum validation batches to run (None = unlimited)
    /// Limits validation to N batches to reduce memory usage on constrained GPUs.
    /// Example: 50 batches = ~500MB vs 1760MB for full validation (176 batches)
    pub max_validation_batches: Option<usize>,

    /// Validation frequency (run validation every N epochs, default: 1)
    /// Used in hyperparameter optimization to control validation overhead.
    /// Set to 1 for every epoch (normal), higher values for faster training.
    pub validation_frequency: usize,

    /// Checkpoint directory
    pub checkpoint_dir: String,
}

impl Default for TFTTrainerConfig {
    fn default() -> Self {
        let batch_size = 32; // Reduced for 4GB VRAM (overridden if auto_batch_size=true)
        Self {
            epochs: 100,
            learning_rate: 1e-3,
            batch_size,
            auto_batch_size: false, // Default: manual batch size
            hidden_dim: 256,
            num_attention_heads: 8,
            dropout_rate: 0.1,
            lstm_layers: 2,
            quantiles: vec![0.1, 0.5, 0.9],
            lookback_window: 60,
            forecast_horizon: 10,
            use_gpu: true,
            use_gradient_checkpointing: false, // Default: off (prioritize speed over memory)
            validation_batch_size: batch_size, // Match training batch_size to avoid memory spikes
            max_validation_batches: None, // Default: unlimited (use all validation data)
            validation_frequency: 1, // Default: validate every epoch
            checkpoint_dir: "/tmp/tft_checkpoints".to_owned(),
        }
    }
}

impl TFTTrainerConfig {
    /// Create TFT model config from trainer config
    pub fn to_model_config(&self) -> TFTConfig {
        TFTConfig {
            input_dim: 225, // 5 + 10 + 210 = 225 (static + known + unknown)
            hidden_dim: self.hidden_dim,
            num_heads: self.num_attention_heads,
            num_layers: self.lstm_layers,
            prediction_horizon: self.forecast_horizon,
            sequence_length: self.lookback_window,
            num_quantiles: 3, // [0.1, 0.5, 0.9]
            num_static_features: 5,
            num_known_features: 10,
            num_unknown_features: 210, // Wave C (201) + Wave D (24) = 225 total features
            learning_rate: self.learning_rate,
            batch_size: self.batch_size,
            dropout_rate: self.dropout_rate,
            l2_regularization: 1e-4,
            use_flash_attention: true,
            memory_efficient: true,
            max_inference_latency_us: 50,
            target_throughput_pps: 100_000,
        }
    }

    /// Create TFT training config from trainer config
    pub fn to_training_config(&self) -> TFTTrainingConfig {
        TFTTrainingConfig {
            epochs: self.epochs,
            batch_size: self.batch_size,
            learning_rate: self.learning_rate,
            dropout_rate: self.dropout_rate,
            gradient_checkpointing: self.use_gradient_checkpointing,
            validation_batch_size: self.validation_batch_size,
            max_validation_batches: self.max_validation_batches,
            validation_frequency: self.validation_frequency,
            ..Default::default()
        }
    }
}
