//! TFT DBN Training Extension
//!
//! Adds OHLCV bar training support to TFTTrainer with sliding window
//! sample creation (lookback=60, horizon=10) and OOM-resilient batch retry.

use ndarray::{Array1, Array2};
use tracing::info;

use crate::features::extraction::{FeatureExtractor, OHLCVBar};
use crate::tft::training::TFTDataLoader;
use crate::trainers::tft::{TFTTrainer, TrainingMetrics};
use crate::{MLError, MLResult};

impl TFTTrainer {
    /// Train TFT on pre-loaded OHLCV bars (from DBN or any other source).
    pub async fn train_from_bars(&mut self, bars: &[OHLCVBar]) -> MLResult<TrainingMetrics> {
        info!("Starting TFT training from {} OHLCV bars", bars.len());

        let training_data = self.prepare_training_data(bars)?;

        self.train_from_samples(training_data).await
    }

    /// Shared training loop: split data, create loaders, train with OOM retry.
    async fn train_from_samples(
        &mut self,
        training_data: Vec<(Array1<f64>, Array2<f64>, Array2<f64>, Array1<f64>)>,
    ) -> MLResult<TrainingMetrics> {
        info!(
            "Target normalization applied: mean={:.2}, std={:.2}",
            self.target_mean.unwrap_or(0.0),
            self.target_std.unwrap_or(1.0)
        );

        info!(
            "Loaded {} training samples (lookback=60, horizon=10)",
            training_data.len()
        );

        // Split into train/val (80/20)
        let split_idx = (training_data.len() as f64 * 0.8) as usize;
        let train_data = training_data[..split_idx].to_vec(); // cpu-side split
        let val_data = training_data[split_idx..].to_vec(); // cpu-side split

        info!(
            "Split data: {} train samples, {} val samples",
            train_data.len(),
            val_data.len()
        );

        // Validate sufficient validation data
        let validation_batch_size = self.get_training_config().validation_batch_size;
        let expected_val_batches = val_data.len() / validation_batch_size;

        if expected_val_batches == 0 {
            return Err(MLError::InsufficientData(
                format!(
                    "Insufficient validation data: {} samples with batch_size={} produces 0 batches. \
                     Need at least {} validation samples. \
                     Current dataset: {} total samples, {} validation (20%). \
                     Recommendations: (1) Use larger dataset, (2) reduce validation_batch_size, or (3) increase train/val split ratio",
                    val_data.len(),
                    validation_batch_size,
                    validation_batch_size,
                    training_data.len(),
                    val_data.len()
                )
            ));
        }

        info!(
            "Validation batches expected: {} (batch_size={})",
            expected_val_batches, validation_batch_size
        );

        // OOM retry loop: Automatically reduce batch size if OOM occurs during training
        let mut current_batch_size = self.get_training_config().batch_size;
        let mut oom_retry_count = 0;
        const MAX_OOM_RETRIES: usize = 3;
        let min_batch_size = 2; // Minimum batch size for OOM recovery

        loop {
            info!(
                "Creating data loaders (batch_size={}, attempt={})",
                current_batch_size,
                oom_retry_count + 1
            );

            // Create data loaders with current batch size
            let train_loader = TFTDataLoader::new(train_data.clone(), current_batch_size, true);
            let val_loader = TFTDataLoader::new(
                val_data.clone(),
                self.get_training_config().validation_batch_size,
                false,
            );

            // Attempt training
            match self.train(train_loader, val_loader).await {
                Ok(metrics) => {
                    if oom_retry_count > 0 {
                        info!(
                            "Training completed successfully after {} OOM retries (final batch_size={})",
                            oom_retry_count, current_batch_size
                        );
                    } else {
                        info!(
                            "Training completed successfully (batch_size={})",
                            current_batch_size
                        );
                    }
                    return Ok(metrics);
                },
                Err(e) => {
                    // Check if error is OOM-related
                    if Self::is_oom_error(&e)
                        && oom_retry_count < MAX_OOM_RETRIES
                        && current_batch_size > min_batch_size
                    {
                        oom_retry_count += 1;
                        let old_batch_size = current_batch_size;
                        current_batch_size = current_batch_size / 2;

                        // Enforce minimum batch size
                        if current_batch_size < min_batch_size {
                            current_batch_size = min_batch_size;
                        }

                        tracing::warn!(
                            "OOM detected (attempt {}/{}), reducing batch_size: {} -> {}",
                            oom_retry_count,
                            MAX_OOM_RETRIES,
                            old_batch_size,
                            current_batch_size
                        );

                        // Update training config with reduced batch size
                        self.update_batch_size(current_batch_size);

                        // Retry with smaller batch size
                        continue;
                    } else {
                        // Non-OOM error OR retries exhausted OR batch size at minimum
                        if Self::is_oom_error(&e) {
                            return Err(MLError::TrainingError(format!(
                                "Training OOM after {} retries (final batch_size={}). \
                                 Consider: (1) using a GPU with more VRAM, (2) reducing model size, or (3) using CPU",
                                oom_retry_count, current_batch_size
                            )));
                        }
                        return Err(e);
                    }
                },
            }
        }
    }

    /// Convert OHLCV bars into TFT training samples (data-source-agnostic).
    fn prepare_training_data(
        &mut self,
        bars: &[OHLCVBar],
    ) -> MLResult<Vec<(Array1<f64>, Array2<f64>, Array2<f64>, Array1<f64>)>> {
        let mut all_ohlcv_bars = bars.to_vec(); // cpu-side clone for sort

        // Sort bars by timestamp (critical for rolling window feature extraction)
        all_ohlcv_bars.sort_by_key(|bar| bar.timestamp);

        // Extract features using full 51-feature extractor (core features)
        let feature_vectors = self.extract_full_features(&all_ohlcv_bars)?;

        info!(
            "Extracted {} feature vectors from {} bars",
            feature_vectors.len(),
            all_ohlcv_bars.len()
        );

        // Compute normalization parameters from close prices
        let all_closes: Vec<f64> = all_ohlcv_bars.iter().map(|b| b.close).collect();

        if all_closes.is_empty() {
            return Err(MLError::InsufficientData(
                "No close prices available for normalization".to_owned(),
            ));
        }

        let price_mean = all_closes.iter().sum::<f64>() / all_closes.len() as f64;
        let price_variance = all_closes
            .iter()
            .map(|c| (c - price_mean).powi(2))
            .sum::<f64>()
            / all_closes.len() as f64;
        let price_std = price_variance.sqrt();

        if price_std < 1e-8 {
            return Err(MLError::InvalidInput(format!(
                "Price std_dev too small ({:.2e}), data may be constant",
                price_std
            )));
        }

        info!(
            "Target normalization: mean={:.2}, std={:.2}",
            price_mean, price_std
        );

        self.target_mean = Some(price_mean);
        self.target_std = Some(price_std);

        const LOOKBACK: usize = 60;
        const HORIZON: usize = 10;

        let mut tft_samples = Vec::new();

        for i in 0..(feature_vectors.len().saturating_sub(LOOKBACK + HORIZON)) {
            let static_feats = Array1::from_vec(feature_vectors[i + LOOKBACK][0..5].to_vec()); // cpu-side slice

            // Historical features: fv[15..42] = 27 unknown features per timestep
            // Feature vector is [f64; 42]: [0..5] static, [5..15] known, [15..42] unknown
            const HIST_START: usize = 15;
            const HIST_END: usize = 42;
            const NUM_HIST_FEATURES: usize = HIST_END - HIST_START; // 27

            let mut hist_data = Vec::with_capacity(LOOKBACK * NUM_HIST_FEATURES);
            for j in i..(i + LOOKBACK) {
                hist_data.extend_from_slice(&feature_vectors[j][HIST_START..HIST_END]);
            }
            let historical_feats =
                Array2::from_shape_vec((LOOKBACK, NUM_HIST_FEATURES), hist_data).map_err(|e| {
                    MLError::ModelError(format!(
                        "Failed to create historical features array: {}",
                        e
                    ))
                })?;

            let mut fut_data = Vec::new();
            for j in (i + LOOKBACK)..(i + LOOKBACK + HORIZON) {
                fut_data.extend_from_slice(&feature_vectors[j][5..15]);
            }
            let future_feats = Array2::from_shape_vec((HORIZON, 10), fut_data).map_err(|e| {
                MLError::ModelError(format!("Failed to create future features array: {}", e))
            })?;

            let mut targets = Vec::new();
            for j in (i + LOOKBACK)..(i + LOOKBACK + HORIZON) {
                let raw_price = all_ohlcv_bars[j + 50].close;
                let normalized = (raw_price - price_mean) / (price_std + 1e-8);
                targets.push(normalized);
            }
            let target_feats = Array1::from_vec(targets);

            tft_samples.push((static_feats, historical_feats, future_feats, target_feats));
        }

        info!(
            "Created {} TFT training samples (lookback={}, horizon={})",
            tft_samples.len(),
            LOOKBACK,
            HORIZON
        );

        Ok(tft_samples)
    }

    /// Extract full 51 features (core features, Proxy OFI removed in WAVE 10) from OHLCV bars
    ///
    /// Uses the same feature extraction pipeline as DQN/PPO for consistency
    fn extract_full_features(&self, bars: &[OHLCVBar]) -> MLResult<Vec<[f64; 42]>> {
        if bars.is_empty() {
            return Err(MLError::InsufficientData(
                "Cannot extract features from empty bar sequence".to_owned(),
            ));
        }

        const WARMUP_PERIOD: usize = 50;
        if bars.len() < WARMUP_PERIOD {
            return Err(MLError::InsufficientData(format!(
                "Insufficient data: {} bars provided, {} required for warmup",
                bars.len(),
                WARMUP_PERIOD
            )));
        }

        let mut extractor = FeatureExtractor::new();
        let mut feature_vectors = Vec::with_capacity(bars.len() - WARMUP_PERIOD);

        // Feed bars sequentially to build rolling windows
        for (i, bar) in bars.iter().enumerate() {
            extractor.update(bar).map_err(|e| {
                MLError::ModelError(format!("Feature extraction failed at bar {}: {}", i, e))
            })?;

            // Start extracting features after warmup
            if i >= WARMUP_PERIOD {
                let features_51 = extractor.extract_current_features().map_err(|e| {
                    MLError::ModelError(format!("Failed to extract features at bar {}: {}", i, e))
                })?;
                feature_vectors.push(features_51);
            }
        }

        Ok(feature_vectors)
    }
}
