//! GPU-resident walk-forward data manager.
//!
//! Uploads the ENTIRE dataset to GPU VRAM once, then provides per-fold
//! views via index ranges. This eliminates per-fold re-upload latency
//! and enables zero-copy fold transitions on H100 (80 GB VRAM).
//!
//! # Architecture
//!
//! ```text
//! GPU VRAM (uploaded once):
//! ┌─────────────────────────────────────────────────────────┐
//! │  features [total_bars, 42]  │  targets [total_bars, 4]  │
//! │  ofi      [total_bars, 8]   │  (optional)               │
//! └─────────────────────────────────────────────────────────┘
//!
//! Fold 0: train [0..T0]          val [T0..V0]    test [V0..E0]
//! Fold 1: train [0..T1]          val [T1..V1]    test [V1..E1]
//! Fold 2: train [0..T2]          val [T2..V2]    test [V2..E2]
//!           (expanding window)
//!
//! Per-fold kernel launch:
//!   episode_starts[] = evenly spaced within [fold.train_start..fold.train_end]
//!   total_bars = fold.train_end  (kernel can't read past this)
//! ```
//!
//! No new CUDA kernels required — the existing `dqn_full_experience_kernel`
//! naturally restricts to any contiguous subrange via `episode_starts` and
//! `total_bars`.

use std::sync::Arc;
use cudarc::driver::{CudaSlice, CudaStream};
use tracing::info;

use crate::features::extraction::FeatureVector;
use crate::MLError;

/// A walk-forward fold's index range into the global GPU buffer.
#[derive(Debug, Clone)]
pub struct GpuFoldRange {
    /// Fold index (0-based).
    pub fold: usize,
    /// First bar index of the training window (inclusive).
    pub train_start: usize,
    /// Last bar index of the training window (exclusive).
    pub train_end: usize,
    /// First bar index of the validation window (inclusive).
    pub val_start: usize,
    /// Last bar index of the validation window (exclusive).
    pub val_end: usize,
    /// First bar index of the test window (inclusive).
    pub test_start: usize,
    /// Last bar index of the test window (exclusive).
    pub test_end: usize,
}

impl GpuFoldRange {
    /// Number of bars in the training window.
    pub fn train_len(&self) -> usize {
        self.train_end.saturating_sub(self.train_start)
    }

    /// Number of bars in the validation window.
    pub fn val_len(&self) -> usize {
        self.val_end.saturating_sub(self.val_start)
    }

    /// Number of bars in the test window.
    pub fn test_len(&self) -> usize {
        self.test_end.saturating_sub(self.test_start)
    }

    /// Compute episode start indices for the GPU experience collector kernel.
    ///
    /// Episodes are evenly spaced within `[train_start, train_end - timesteps)`.
    /// The kernel reads `timesteps` bars forward from each start, so starts must
    /// leave room for the full episode without exceeding `train_end`.
    pub fn episode_starts(&self, n_episodes: usize, timesteps: usize) -> Vec<i32> {
        let usable = self.train_len().saturating_sub(timesteps).max(1);
        let stride = (usable / n_episodes.max(1)).max(1);
        (0..n_episodes)
            .map(|i| {
                let offset = (i * stride) % usable;
                (self.train_start + offset) as i32
            })
            .collect()
    }
}

/// Walk-forward configuration for GPU-resident data.
#[derive(Debug, Clone)]
pub struct GpuWalkForwardConfig {
    /// Fraction of total data for the initial training window (default: 0.5).
    pub initial_train_fraction: f64,
    /// Fraction of total data per validation window (default: 0.125).
    pub val_fraction: f64,
    /// Fraction of total data per test window (default: 0.125).
    pub test_fraction: f64,
    /// Step size between folds as a fraction of total data (default: 0.125).
    pub step_fraction: f64,
}

impl Default for GpuWalkForwardConfig {
    fn default() -> Self {
        Self {
            initial_train_fraction: 0.5,
            val_fraction: 0.125,
            test_fraction: 0.125,
            step_fraction: 0.125,
        }
    }
}

impl GpuWalkForwardConfig {
    /// Generate fold ranges for the given total number of bars.
    ///
    /// Returns an empty `Vec` if the data is too small for even one fold.
    pub fn generate_folds(&self, total_bars: usize) -> Vec<GpuFoldRange> {
        let initial_train = (total_bars as f64 * self.initial_train_fraction) as usize;
        let val_size = (total_bars as f64 * self.val_fraction) as usize;
        let test_size = (total_bars as f64 * self.test_fraction) as usize;
        let step_size = (total_bars as f64 * self.step_fraction).max(1.0) as usize;
        let min_split = 50; // minimum bars per split

        if initial_train < min_split || val_size < min_split || test_size < min_split {
            return Vec::new();
        }

        let mut folds = Vec::new();
        let mut fold_idx = 0;

        loop {
            let train_end = initial_train + fold_idx * step_size;
            let val_start = train_end;
            let val_end = val_start + val_size;
            let test_start = val_end;
            let test_end = test_start + test_size;

            if test_end > total_bars {
                break;
            }

            folds.push(GpuFoldRange {
                fold: fold_idx,
                train_start: 0, // expanding window: always starts at 0
                train_end,
                val_start,
                val_end,
                test_start,
                test_end,
            });

            fold_idx += 1;
        }

        folds
    }
}

/// GPU-resident walk-forward data store.
///
/// Holds the entire dataset on GPU as cudarc `CudaSlice<f32>` buffers.
/// Folds are expressed as index ranges into these buffers — no re-upload needed.
// CudaSlice<f32> doesn't implement Debug, so manual impl.
impl std::fmt::Debug for GpuWalkForwardData {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GpuWalkForwardData")
            .field("total_bars", &self.total_bars)
            .field("feature_dim", &self.feature_dim)
            .field("target_dim", &self.target_dim)
            .field("folds", &self.folds.len())
            .field("vram_bytes", &self.vram_bytes)
            .field("has_ofi", &self.ofi.is_some())
            .finish()
    }
}

pub struct GpuWalkForwardData {
    /// Market features [total_bars * FEATURE_DIM] (FEATURE_DIM = 42).
    pub features: CudaSlice<f32>,
    /// Target prices [total_bars * TARGET_DIM] (TARGET_DIM = 4).
    pub targets: CudaSlice<f32>,
    /// OFI features [total_bars * 8] (optional, from MBP-10).
    pub ofi: Option<CudaSlice<f32>>,
    /// Total number of bars in the dataset.
    pub total_bars: usize,
    /// Feature dimension (42).
    pub feature_dim: usize,
    /// Target dimension (4).
    pub target_dim: usize,
    /// Generated fold ranges.
    pub folds: Vec<GpuFoldRange>,
    /// VRAM usage in bytes.
    pub vram_bytes: usize,
}

impl GpuWalkForwardData {
    /// Upload the entire dataset to GPU VRAM and generate walk-forward folds.
    ///
    /// On H100 (80 GB), a typical dataset of 1M bars uses:
    ///   features: 1M × 42 × 4 = 168 MB
    ///   targets:  1M × 4  × 4 =  16 MB
    ///   ofi:      1M × 8  × 4 =  32 MB
    ///   total:                   216 MB (~0.27% of VRAM)
    pub fn upload(
        training_data: &[(FeatureVector, Vec<f64>)],
        ofi_data: Option<&[[f64; 8]]>,
        wf_config: &GpuWalkForwardConfig,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let total_bars = training_data.len();
        if total_bars == 0 {
            return Err(MLError::ModelError("Empty training data".to_owned()));
        }

        let feature_dim = 42;
        let target_dim = 4;

        // Flatten features: [total_bars * 42]
        let mut flat_features = Vec::with_capacity(total_bars * feature_dim);
        for (features, _) in training_data {
            for &v in features.iter() {
                flat_features.push(v as f32);
            }
        }

        // Flatten targets: [total_bars * 4]
        let mut flat_targets = Vec::with_capacity(total_bars * target_dim);
        for (_, targets) in training_data {
            for i in 0..target_dim {
                flat_targets.push(targets.get(i).copied().unwrap_or(0.0) as f32);
            }
        }

        // Upload features
        let features_gpu = stream
            .clone_htod(&flat_features)
            .map_err(|e| MLError::ModelError(format!("WF features upload: {e}")))?;

        let mut vram = total_bars * feature_dim * 4;

        // Upload targets
        let targets_gpu = stream
            .clone_htod(&flat_targets)
            .map_err(|e| MLError::ModelError(format!("WF targets upload: {e}")))?;

        vram += total_bars * target_dim * 4;

        // Upload OFI (optional)
        let ofi_gpu = if let Some(ofi) = ofi_data {
            let n = ofi.len().min(total_bars);
            let mut flat_ofi = Vec::with_capacity(total_bars * 8);
            for i in 0..total_bars {
                if i < n {
                    for &v in ofi[i].iter() {
                        flat_ofi.push(v as f32);
                    }
                } else {
                    flat_ofi.extend_from_slice(&[0.0_f32; 8]);
                }
            }
            let buf = stream
                .clone_htod(&flat_ofi)
                .map_err(|e| MLError::ModelError(format!("WF OFI upload: {e}")))?;
            vram += total_bars * 8 * 4;
            Some(buf)
        } else {
            None
        };

        // Generate walk-forward folds
        let folds = wf_config.generate_folds(total_bars);

        info!(
            "GPU walk-forward data uploaded: {} bars, {} folds, {:.1} MB VRAM",
            total_bars,
            folds.len(),
            vram as f64 / 1_048_576.0,
        );

        for fold in &folds {
            info!(
                "  Fold {}: train [{}..{}] ({} bars), val [{}..{}] ({} bars), test [{}..{}] ({} bars)",
                fold.fold,
                fold.train_start, fold.train_end, fold.train_len(),
                fold.val_start, fold.val_end, fold.val_len(),
                fold.test_start, fold.test_end, fold.test_len(),
            );
        }

        Ok(Self {
            features: features_gpu,
            targets: targets_gpu,
            ofi: ofi_gpu,
            total_bars,
            feature_dim,
            target_dim,
            folds,
            vram_bytes: vram,
        })
    }

    /// Get the number of folds.
    pub fn num_folds(&self) -> usize {
        self.folds.len()
    }

    /// Get a fold's range by index.
    pub fn fold(&self, idx: usize) -> Option<&GpuFoldRange> {
        self.folds.get(idx)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_fold_generation_basic() {
        let config = GpuWalkForwardConfig::default();
        // 1000 bars: initial_train=500, val=125, test=125, step=125
        let folds = config.generate_folds(1000);

        assert!(!folds.is_empty(), "should generate at least one fold");

        let f0 = &folds[0];
        assert_eq!(f0.train_start, 0);
        assert_eq!(f0.train_end, 500);
        assert_eq!(f0.val_start, 500);
        assert_eq!(f0.val_end, 625);
        assert_eq!(f0.test_start, 625);
        assert_eq!(f0.test_end, 750);

        // Check expanding window for fold 1
        if folds.len() > 1 {
            let f1 = &folds[1];
            assert_eq!(f1.train_start, 0);
            assert_eq!(f1.train_end, 625); // expanded by step_size=125
            assert!(f1.train_len() > f0.train_len());
        }
    }

    #[test]
    fn test_fold_generation_insufficient_data() {
        let config = GpuWalkForwardConfig::default();
        // 100 bars: initial_train=50, val=12, test=12 — val/test too small (< 50)
        let folds = config.generate_folds(100);
        assert!(folds.is_empty(), "100 bars too few for default config");
    }

    #[test]
    fn test_episode_starts_within_range() {
        let fold = GpuFoldRange {
            fold: 0,
            train_start: 0,
            train_end: 500,
            val_start: 500,
            val_end: 625,
            test_start: 625,
            test_end: 750,
        };
        let timesteps = 100;
        let n_episodes = 64;
        let starts = fold.episode_starts(n_episodes, timesteps);

        assert_eq!(starts.len(), n_episodes);
        for &s in &starts {
            assert!(s >= 0, "start must be non-negative");
            assert!(
                (s as usize) + timesteps <= fold.train_end,
                "episode must not exceed train window: start={s} + timesteps={timesteps} > train_end={}",
                fold.train_end,
            );
        }
    }

    #[test]
    fn test_fold_generation_many_bars() {
        let config = GpuWalkForwardConfig {
            initial_train_fraction: 0.5,
            val_fraction: 0.1,
            test_fraction: 0.1,
            step_fraction: 0.1,
        };
        // 10000 bars: should generate multiple folds
        let folds = config.generate_folds(10000);
        assert!(folds.len() >= 3, "should generate >=3 folds for 10K bars");

        // Each successive fold has strictly more training data
        for i in 1..folds.len() {
            assert!(
                folds[i].train_len() > folds[i - 1].train_len(),
                "fold {} train ({}) should exceed fold {} train ({})",
                i,
                folds[i].train_len(),
                i - 1,
                folds[i - 1].train_len(),
            );
        }
    }
}
