//! TFT trainer type definitions
//!
//! This module contains shared type definitions used by the TFT trainer,
//! including metrics, state tracking, and resource monitoring structures.

use std::collections::HashMap;
use std::time::Instant;

use serde::{Deserialize, Serialize};

// ============================================================================
// Training State and Progress
// ============================================================================

/// Internal training state tracking
///
/// Tracks the current state of training including epoch progress, validation
/// metrics, and early stopping.
#[derive(Debug)]
pub(crate) struct TrainingState {
    /// Current epoch
    pub current_epoch: usize,

    /// Global step counter
    pub global_step: usize,

    /// Best validation loss
    pub best_val_loss: f64,

    /// Training start time
    pub started_at: Option<Instant>,

    /// Current learning rate
    pub learning_rate: f64,

    /// Early stopping patience counter
    pub patience_counter: usize,

    /// Last valid validation loss (for cached display)
    pub last_val_loss: Option<f64>,

    /// Last valid validation metrics (for cached display)
    pub last_val_metrics: ValidationMetrics,
}

impl Default for TrainingState {
    fn default() -> Self {
        Self {
            current_epoch: 0,
            global_step: 0,
            best_val_loss: f64::INFINITY,
            started_at: None,
            learning_rate: 0.0,
            patience_counter: 0,
            last_val_loss: None,
            last_val_metrics: ValidationMetrics::default(),
        }
    }
}

/// Public training progress for external monitoring
///
/// Provides real-time training progress information including metrics,
/// resource usage, and status updates.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingProgress {
    /// Current epoch (1-indexed for display)
    pub current_epoch: u32,

    /// Total epochs
    pub total_epochs: u32,

    /// Progress percentage (0.0 to 100.0)
    pub progress_percentage: f32,

    /// Current metrics
    pub metrics: HashMap<String, f32>,

    /// Status message
    pub message: String,

    /// Timestamp (Unix seconds)
    pub timestamp: i64,

    /// Resource usage
    pub resource_usage: ResourceUsage,
}

/// Resource usage statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceUsage {
    /// CPU usage percentage
    pub cpu_usage_percent: f32,

    /// Memory usage in GB
    pub memory_usage_gb: f32,

    /// GPU usage percentage
    pub gpu_usage_percent: f32,

    /// GPU memory usage in GB
    pub gpu_memory_usage_gb: f32,
}

impl Default for ResourceUsage {
    fn default() -> Self {
        Self {
            cpu_usage_percent: 0.0,
            memory_usage_gb: 0.0,
            gpu_usage_percent: 0.0,
            gpu_memory_usage_gb: 0.0,
        }
    }
}

// ============================================================================
// Validation Metrics
// ============================================================================

/// Internal validation metrics
///
/// Tracks validation performance metrics including quantile loss, RMSE,
/// and attention entropy for model interpretability.
#[derive(Debug, Clone)]
pub(crate) struct ValidationMetrics {
    pub quantile_loss: f64,
    pub rmse: f64,
    pub attention_entropy: f64,
}

impl Default for ValidationMetrics {
    fn default() -> Self {
        Self {
            quantile_loss: 0.0,
            rmse: 0.0,
            attention_entropy: 0.0,
        }
    }
}

// ============================================================================
// Training Metrics Result
// ============================================================================

/// Training metrics result
///
/// Final training results returned after training completion, including loss values,
/// RMSE, and attention entropy.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TrainingMetrics {
    /// Final training loss
    pub train_loss: f64,

    /// Final validation loss
    pub val_loss: f64,

    /// Quantile loss
    pub quantile_loss: f64,

    /// RMSE
    pub rmse: f64,

    /// Attention entropy (interpretability metric)
    pub attention_entropy: f64,

    /// Total training time in seconds
    pub training_time_seconds: f64,
}
