//! DQN Trainer -- Action selection, routing, and fill simulation

use anyhow::Result;
use ml_core::cuda_autograd::GpuTensor;
use tracing::{debug, info};

use super::DQNTrainer;
use crate::dqn::action_space::{ExposureLevel, FactoredAction};
use crate::dqn::TradingState;
use crate::dqn::order_router::OrderRouter;
use ml_core::fill_simulator::FillResult;

/// Extract the raw `&CudaSlice<f32>` from a `GpuTensor` reference.
/// GpuTensor stores data as `CudaSlice<f32>` directly — no conversion needed.
macro_rules! extract_cuda_f32 {
    ($tensor:expr) => {{
        $tensor.data()
    }};
}

impl DQNTrainer {
    pub(crate) async fn select_action(&self, state: &TradingState) -> Result<FactoredAction> {
        let _agent = self.agent.read().await;
        let state_vec = state.to_vector();
        let raw_dim = state_vec.len();
        let aligned_dim = (raw_dim + 7) & !7;
        let padded: Vec<f32> = if aligned_dim > raw_dim {
            let mut v = state_vec.to_vec(); v.resize(aligned_dim, 0.0); v // cpu-side align pad
        } else { state_vec.to_vec() }; // cpu-side vec clone
        let stream = self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream required"))?;
        let state_tensor = GpuTensor::from_host(&padded, vec![1, aligned_dim], stream)
            .map_err(|e| anyhow::anyhow!("Failed to create state tensor: {}", e))?;
        let action_idx = self.epsilon_greedy_action(&state_tensor).await?;
        let exposure = ExposureLevel::from_index(action_idx)
            .map_err(|e| anyhow::anyhow!("Invalid action index {}: {}", action_idx, e))?;
        Ok(self.route_action(exposure, self.hyperparams.avg_spread as f32))
    }

    pub(crate) fn route_action(&self, exposure: ExposureLevel, current_spread: f32) -> FactoredAction {
        OrderRouter::route(exposure, current_spread, self.hyperparams.avg_spread as f32, self.vol_ema as f32, self.median_vol as f32)
    }

    pub(crate) fn simulate_fill(&self, action: FactoredAction, step: usize) -> (FactoredAction, FillResult) {
        let normalized_vol = if self.median_vol > 0.0 { (self.vol_ema / self.median_vol) as f32 } else { 1.0 };
        let spread_bps = self.hyperparams.avg_spread * 10000.0;
        let fill_result = self.fill_simulator.simulate_fill(action.order, action.urgency, normalized_vol, spread_bps, step, action.exposure as usize);
        if fill_result.filled { (action, fill_result) } else { (OrderRouter::route_default(ExposureLevel::Flat), fill_result) }
    }

    pub(crate) async fn select_actions_batch(&mut self, states: &[TradingState]) -> Result<Vec<FactoredAction>> {
        if states.is_empty() { return Ok(Vec::new()); }
        let agent = self.agent.read().await;
        let batch_size = states.len();
        let first_vec = states.first().map(|s| s.to_vector()).ok_or_else(|| anyhow::anyhow!("Empty states slice"))?;
        let raw_state_dim = first_vec.len();
        let aligned_dim = (raw_state_dim + 7) & !7;
        let pad = aligned_dim - raw_state_dim;
        let mut flat_states = Vec::with_capacity(batch_size * aligned_dim);
        flat_states.extend_from_slice(&first_vec);
        flat_states.extend(std::iter::repeat_n(0.0_f32, pad));
        for (i, state) in states.iter().enumerate().skip(1) {
            let vec = state.to_vector();
            if vec.len() != raw_state_dim { return Err(anyhow::anyhow!("State {} dimension mismatch: expected {}, got {}", i, raw_state_dim, vec.len())); }
            flat_states.extend_from_slice(&vec);
            flat_states.extend(std::iter::repeat_n(0.0_f32, pad));
        }
        let stream = self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream required"))?;
        let batch_tensor = GpuTensor::from_host(&flat_states, vec![batch_size, aligned_dim], stream)
            .map_err(|e| anyhow::anyhow!("Failed to create batched state tensor: {}", e))?;
        let base_epsilon = agent.get_effective_epsilon() as f64;
        let adjusted_epsilon = self.calculate_volatility_adjusted_epsilon(base_epsilon);
        let epsilon = adjusted_epsilon as f32;
        debug!("Epsilon: base={:.4}, volatility-adjusted={:.4}", base_epsilon, adjusted_epsilon);
        let batch_q_values = agent.forward(&batch_tensor).map_err(|e| anyhow::anyhow!("Batched forward pass failed: {}", e))?;
        let branching_q_tensors: Option<(GpuTensor, GpuTensor, GpuTensor)> = if true {
            agent.batch_branching_q_values(&batch_tensor).map_err(|e| anyhow::anyhow!("Branching Q-values failed: {}", e))?
        } else { None };
        drop(agent);
        {
            if self.gpu_action_selector.is_none() && self.cuda_stream.is_some() {
                let stream = std::sync::Arc::clone(self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream required"))?);
                let mut selector = crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::new(stream, self.hyperparams.batch_size.max(batch_size).max(8192), 0xDEAD_BEEF_CAFE_u64)
                    .map_err(|e| anyhow::anyhow!("GPU fused action selector init failed: {e}"))?;
                selector.set_q_gap_threshold(self.hyperparams.q_gap_threshold as f32);
                info!("GPU action selector initialized for select_actions_batch (q_gap={:.3})", self.hyperparams.q_gap_threshold);
                self.gpu_action_selector = Some(selector);
            }
            let selector = self.gpu_action_selector.as_mut().ok_or_else(|| anyhow::anyhow!("GPU action selector requires CUDA device"))?;

            // Upload UCB count bonuses for branching exploration.
            // CountBonus tracks per-branch action counts and computes UCB bonuses.
            // Bonuses encourage trying under-explored actions (exploration pressure).
            if true {
                let agent_r = self.agent.read().await;
                if let Some(bonuses) = agent_r.get_count_bonuses_branched() {
                    let coeff = self.hyperparams.count_bonus_coefficient.unwrap_or(0.0) as f32;
                    if coeff > 0.0 {
                        let (be, bo, bu) = bonuses;
                        let be_f32: [f32; 5] = std::array::from_fn(|i| be.get(i).copied().unwrap_or(0.0) as f32 * coeff);
                        let bo_f32: [f32; 3] = std::array::from_fn(|i| bo.get(i).copied().unwrap_or(0.0) as f32 * coeff);
                        let bu_f32: [f32; 3] = std::array::from_fn(|i| bu.get(i).copied().unwrap_or(0.0) as f32 * coeff);
                        let _ = selector.set_count_bonuses(&be_f32, &bo_f32, &bu_f32);
                    }
                }
                drop(agent_r);
            }

            let factored_slice = if let Some((ref q_exp, ref q_ord, ref q_urg)) = branching_q_tensors {
                let qe_v = extract_cuda_f32!(q_exp);
                let qo_v = extract_cuda_f32!(q_ord);
                let qu_v = extract_cuda_f32!(q_urg);
                selector.select_actions_branching(&qe_v, &qo_v, &qu_v, epsilon, batch_size).map_err(|e| anyhow::anyhow!("GPU branching action selection failed: {e}"))?
            } else {
                let q_v = extract_cuda_f32!(&batch_q_values);
                let exposure_slice = selector.select_actions(&q_v, epsilon, batch_size, 5).map_err(|e| anyhow::anyhow!("GPU fused action selection failed: {e}"))?;
                selector.route_exposure_to_factored(&exposure_slice, batch_size, self.hyperparams.avg_spread as f32, self.hyperparams.avg_spread as f32, self.vol_ema as f32, self.median_vol as f32)
                    .map_err(|e| anyhow::anyhow!("GPU route exposure->factored: {e}"))?
            };
            let host_actions = crate::cuda_pipeline::gpu_action_selector::GpuActionSelector::readback_actions(selector.stream(), &factored_slice, batch_size)
                .map_err(|e| anyhow::anyhow!("Factored index readback: {e}"))?;
            let mut actions = Vec::with_capacity(batch_size);
            for (i, &idx) in host_actions.iter().enumerate() {
                let action = FactoredAction::from_index(idx as usize).map_err(|e| anyhow::anyhow!("Invalid factored index {idx} at [{i}]: {e}"))?;
                actions.push(action);
            }
            Ok(actions)
        }
    }

    // select_actions_batch_gpu removed: zero callers, was a Candle forward path.
    // GPU experience collection uses the experience_action_select kernel directly.

    // ── Ensemble consensus action selection ──────────────────────────────
    // Task 9 analysis: The ensemble heads (heads 1..K-1) live on FusedTrainingCtx
    // which is only accessible during the training loop (not during inference).
    // The ensemble already provides value via diversity gradient during training:
    // each head learns different strategies, and the KL-divergence loss encourages
    // disagreement. At inference time, head 0 (the primary network inside the
    // CUDA Graph) is the sole Q-value source for action selection.
    //
    // To use ensemble consensus at inference, the extra head weights would need
    // to be synced into separate forward-pass-capable networks — a significant
    // memory + compute cost that is not justified until production A/B testing
    // shows head-0-only inference is a bottleneck. The ensemble diversity gradient
    // already improves head 0's policy during training.
    //
    // Future: if ensemble consensus is needed at inference, extract head weights
    // at epoch boundary into lightweight forward-only modules and compute
    // mean Q + variance scaling here in select_actions_batch().

    pub(crate) async fn epsilon_greedy_action(&self, state: &GpuTensor) -> Result<usize> {
        use rand::Rng;
        let epsilon = self.get_epsilon().await? as f32;
        let mut rng = rand::thread_rng();
        if rng.gen::<f32>() < epsilon { Ok(rng.gen_range(0..5)) }
        else {
            let agent = self.agent.read().await;
            let q_values = agent.forward(state)?;
            let stream = self.cuda_stream.as_ref().ok_or_else(|| anyhow::anyhow!("CUDA stream required"))?;
            let indices = q_values.argmax(1, stream).map_err(|e| anyhow::anyhow!("argmax: {e}"))?;
            let best_action = indices.first().copied().unwrap_or(2) as usize;
            Ok(best_action)
        }
    }
}
