//! DQN Trainer — State/feature vector conversion

use anyhow::{Context, Result};
use ml_core::cuda_autograd::GpuTensor;
use common::CommonError;
use num_traits::ToPrimitive;

use super::DQNTrainer;
use crate::dqn::TradingState;
use crate::features::extraction::FeatureVector;

impl DQNTrainer {
    /// Convert feature vector to TradingState (42 market features → 45-dim state with portfolio)
    ///
    /// CRITICAL BUG FIX: Features 0-3 are LOG RETURNS (signed), not raw prices.
    /// Using .abs() destroys directional information (bullish vs bearish).
    /// We now use TradingState::from_normalized() to preserve sign information.
    ///
    /// Feature mapping:
    /// - Features 0-3: OHLC log returns → price_features (signed, normalized)
    /// - Features 4-224: All other features → technical_indicators (221 features including Wave D)
    ///
    /// # Arguments
    ///
    /// * `feature_vec` - 42-dimensional FeatureVector from extraction pipeline
    /// * `close_price` - Current close price for portfolio feature calculation (optional)
    ///
    /// # Bug #4 Fix
    ///
    /// Added close_price parameter to enable portfolio feature population from PortfolioTracker.
    pub(crate) fn feature_vector_to_state(
        &self,
        feature_vec: &FeatureVector,
        close_price: Option<rust_decimal::Decimal>,
    ) -> Result<TradingState> {
        self.feature_vector_to_state_with_ofi(feature_vec, close_price, None)
    }

    pub(crate) fn feature_vector_to_state_with_ofi(
        &self,
        feature_vec: &FeatureVector,
        close_price: Option<rust_decimal::Decimal>,
        ofi_index: Option<usize>,
    ) -> Result<TradingState> {
        // States are pre-normalized during data loading
        let normalized_features: Vec<f32> = feature_vec.iter().map(|&v| v as f32).collect();

        // Features 0-3 are LOG RETURNS - preserve sign information for price direction
        let price_features: Vec<f32> = vec![
            normalized_features[0], // open log return (can be negative)
            normalized_features[1], // high log return (can be negative)
            normalized_features[2], // low log return (can be negative)
            normalized_features[3], // close log return (can be negative)
        ];

        // 42-FEATURE ARCHITECTURE: Extract market features (indices 4-41)
        assert_eq!(
            normalized_features.len(),
            42,
            "Expected 42 market features (got {})",
            normalized_features.len()
        );
        let market_features: Vec<f32> = normalized_features[4..42]
            .iter()
            .map(|&x| x as f32)
            .collect();

        // Legacy technical_indicators (empty for 42-feature architecture)
        let technical_indicators = vec![];

        // BUG #36 FIX: Use NORMALIZED portfolio features to prevent Q-value explosion
        let portfolio_features = if let Some(price) = close_price {
            let price_f32 = price.to_f32().unwrap_or(0.0);
            self.portfolio_tracker
                .get_portfolio_features(price_f32)
                .to_vec() // cpu-side portfolio features
        } else {
            vec![0.0, 0.0, 0.0] // Fallback if no price provided
        };

        // OFI regime features: 8 features from MBP-10 order book data.
        // When OFI is enabled (mbp10_data_dir set), always return 8 features
        // (zeros if data didn't load) to match state_dim=53.
        let ofi_enabled = self.hyperparams.mbp10_data_dir.is_some();
        let regime_features: Vec<f32> = if let (Some(ofi), Some(idx)) = (&self.ofi_features, ofi_index) {
            ofi.get(idx)
                .map(|f| f.iter().map(|&v| v as f32).collect())
                .unwrap_or_else(|| vec![0.0; 8])
        } else if ofi_enabled {
            vec![0.0; 8]
        } else {
            vec![]
        };

        // Use from_normalized() to preserve sign information
        Ok(TradingState::from_normalized(
            price_features,
            technical_indicators,
            market_features,
            portfolio_features,
            regime_features,
        ))
    }

    /// Convert feature vector to state tensor for action selection
    ///
    /// Public wrapper around internal state conversion for hyperopt backtest integration.
    /// Converts a 42-dimensional feature vector to a 45-dimensional state tensor
    /// suitable for DQN agent's select_action method.
    ///
    /// # Arguments
    ///
    /// * `feature_vec` - 42-dimensional market feature vector
    /// * `close_price` - Current close price for portfolio feature calculation
    ///
    /// # Returns
    ///
    /// Result containing the 45-dimensional state tensor ready for model inference.
    /// Portfolio features (last 3 dimensions) are populated via PortfolioTracker.
    pub fn convert_to_state(
        &self,
        feature_vec: &FeatureVector,
        close_price: f64,
    ) -> Result<GpuTensor> {
        let close = rust_decimal::Decimal::try_from(close_price)
            .map_err(|e| CommonError::validation(&format!("Invalid close price: {}", e)))?;

        // Use internal conversion method (returns TradingState)
        let trading_state = self.feature_vector_to_state(feature_vec, Some(close))?;

        // Convert TradingState to flat vector, pad for tensor core alignment
        let state_vec = trading_state.to_vector();
        let raw_dim = state_vec.len();
        let aligned = (raw_dim + 7) & !7;
        let padded: Vec<f32> = if aligned > raw_dim {
            let mut v = state_vec.to_vec(); // cpu-side align pad
            v.resize(aligned, 0.0);
            v
        } else {
            state_vec.to_vec() // cpu-side vec clone
        };

        // Convert to GpuTensor using trainer's CUDA stream
        let stream = self.cuda_stream.as_ref()
            .ok_or_else(|| anyhow::anyhow!("CUDA stream required for state tensor"))?;
        GpuTensor::from_host(&padded, vec![padded.len()], stream)
            .context("Failed to create state tensor from TradingState")
    }

    /// Convert feature vector to flat state Vec<f32> (CPU only, no GPU tensor).
    ///
    /// Same as `convert_to_state` but returns the raw vector instead of a GPU tensor.
    /// Used by chunked batch inference to avoid per-bar GPU allocations.
    pub fn convert_to_state_vec(
        &self,
        feature_vec: &FeatureVector,
        close_price: f64,
    ) -> Result<Vec<f32>> {
        let close = rust_decimal::Decimal::try_from(close_price)
            .map_err(|e| CommonError::validation(&format!("Invalid close price: {}", e)))?;
        let trading_state = self.feature_vector_to_state(feature_vec, Some(close))?;
        Ok(trading_state.to_vector())
    }

    /// Convert feature vector to flat state Vec<f32> with OFI features at the given index.
    ///
    /// Same as `convert_to_state_vec` but injects OFI features from the preloaded
    /// array at `ofi_index`, preventing train/eval feature mismatch.
    pub fn convert_to_state_vec_with_ofi(
        &self,
        feature_vec: &FeatureVector,
        close_price: f64,
        ofi_index: usize,
    ) -> Result<Vec<f32>> {
        let close = rust_decimal::Decimal::try_from(close_price)
            .map_err(|e| CommonError::validation(&format!("Invalid close price: {}", e)))?;
        let trading_state = self.feature_vector_to_state_with_ofi(feature_vec, Some(close), Some(ofi_index))?;
        Ok(trading_state.to_vector())
    }

}
