//! 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
//! - Checkpoint save/load using JSON metadata
//! - 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 std::collections::HashMap;

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

/// Extended TFT with training infrastructure
///
/// This struct wraps TemporalFusionTransformer and adds necessary fields for training:
/// - Step counter for learning rate scheduling
/// - Training loss history
/// - Gradient tracking
///
/// Note: TFT manages its own parameters through its internal variable store.
/// Gradient computation is handled via the GPU autograd system.
pub struct TrainableTFT {
    /// Core TFT model
    pub model: TemporalFusionTransformer,
    /// 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> {
        let model = TemporalFusionTransformer::new(config.clone())?;
        let learning_rate = config.learning_rate;

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

impl UnifiedTrainable for TrainableTFT {
    fn model_type(&self) -> &str {
        "TFT"
    }

    fn device_name(&self) -> String {
        "cuda:0".to_owned()
    }

    fn forward_loss(&mut self, input: &[f32], target: &[f32]) -> Result<f64, MLError> {
        use ml_supervised::gpu_tensor::{GpuTensor, gpu_sub, gpu_sqr, gpu_mean_all};
        let stream = self.model.stream().clone();
        let t = GpuTensor::from_vec(target.to_vec(), &[1, target.len()], &stream)?; // cpu-side: initial upload
        let static_dim = self.model.config.num_static_features;
        let hist_dim = input.len().saturating_sub(static_dim);
        let static_t = GpuTensor::from_vec(input.get(..static_dim).unwrap_or(&[]).to_vec(), &[1, static_dim], &stream)?; // cpu-side: initial upload
        let hist_t = GpuTensor::from_vec(input.get(static_dim..).unwrap_or(&[]).to_vec(), &[1, 1, hist_dim], &stream)?; // cpu-side: initial upload
        let fut_t = GpuTensor::zeros(&[1, 1, self.model.config.num_known_features], &stream)?;
        let output = self.model.forward(&static_t, &hist_t, &fut_t)?;
        // Reshape output to match target shape (flatten horizon dimension).
        // TFT forward returns [batch, horizon, quantiles], target is [batch, quantiles].
        let output = if output.shape != t.shape && output.numel() == t.numel() {
            output.reshape(&t.shape)?
        } else {
            output
        };
        let diff = gpu_sub(&output, &t)?;
        let sq = gpu_sqr(&diff)?;
        let loss = gpu_mean_all(&sq)?; // gpu-exit: 1 scalar loss
        Ok(loss as f64)
    }

    fn backward(&mut self, loss_value: f64) -> Result<f64, MLError> {
        // TFT manages parameters internally via TemporalFusionTransformer::train().
        // The UnifiedTrainable path only supports forward_loss for evaluation.
        // Use TFTTrainer or TemporalFusionTransformer::train() for actual training.
        self.loss_history.push(loss_value);
        Err(MLError::TrainingError(
            "TFT backward not supported via UnifiedTrainable — use TFTTrainer::train() instead".to_owned()
        ))
    }

    fn optimizer_step(&mut self) -> Result<(), MLError> {
        Err(MLError::TrainingError(
            "TFT optimizer not supported via UnifiedTrainable — use TFTTrainer::train() instead".to_owned()
        ))
    }

    fn zero_grad(&mut self) -> Result<(), MLError> {
        self.last_grad_norm = 0.0;
        Ok(())
    }

    fn get_learning_rate(&self) -> f64 {
        self.learning_rate
    }

    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),
            });
        }
        self.learning_rate = lr;
        Ok(())
    }

    fn get_step(&self) -> usize {
        self.step_count
    }

    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);
        }

        custom_metrics.insert("step_count".to_owned(), self.step_count as f64);
        custom_metrics.insert("last_grad_norm".to_owned(), self.last_grad_norm);

        TrainingMetrics {
            loss: self.loss_history.last().copied().unwrap_or(0.0),
            val_loss: None,
            accuracy: None,
            learning_rate: self.learning_rate,
            grad_norm: None,
            custom_metrics,
        }
    }

    fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError> {
        let metadata = CheckpointMetadata {
            model_type: "TFT".to_owned(),
            version: self.model.metadata.version.clone(),
            epoch: self.loss_history.len(),
            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(),
        };

        checkpoint::save_metadata(&metadata, checkpoint_path)?;

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

        Ok(format!("{}.json", checkpoint_path))
    }

    fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result<CheckpointMetadata, MLError> {
        let metadata = checkpoint::load_metadata(checkpoint_path)?;

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

        self.step_count = metadata.step;
        self.model.is_trained = true;
        self.learning_rate = metadata.metrics.learning_rate;

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

        Ok(metadata)
    }
}

#[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,
            learning_rate: 1e-3,
            ..Default::default()
        };
        let model = TrainableTFT::new(config)?;

        assert_eq!(model.model_type(), "TFT");
        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,
            ..Default::default()
        };
        let mut model = TrainableTFT::new(config)?;

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

        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,
            ..Default::default()
        };
        let model = TrainableTFT::new(config)?;

        let metrics = model.collect_metrics();

        assert!(metrics.loss >= 0.0);
        assert_eq!(metrics.learning_rate, model.get_learning_rate());
        assert!(!metrics.custom_metrics.is_empty());
        assert!(metrics.custom_metrics.contains_key("step_count"));

        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,
            ..Default::default()
        };
        let model = TrainableTFT::new(config.clone())?;

        let temp_dir = tempfile::tempdir()?;
        let checkpoint_path = temp_dir.path().join("tft_test_checkpoint");
        let checkpoint_path_str = checkpoint_path.to_str().ok_or_else(|| {
            MLError::ModelError("Invalid path".to_owned())
        })?;

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

        assert!(std::path::Path::new(&format!("{}.json", checkpoint_path_str)).exists());

        let mut loaded_model = TrainableTFT::new(config)?;
        let metadata = loaded_model.load_checkpoint(checkpoint_path_str)?;

        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,
            ..Default::default()
        };
        let mut model = TrainableTFT::new(config)?;

        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,
            ..Default::default()
        };
        let mut model = TrainableTFT::new(config)?;

        model.last_grad_norm = 1.5;
        assert_eq!(model.last_grad_norm, 1.5);

        model.zero_grad()?;
        assert_eq!(model.last_grad_norm, 0.0);

        model.zero_grad()?;
        assert_eq!(model.last_grad_norm, 0.0);

        Ok(())
    }
}
