//! Proximal Policy Optimization (PPO) Implementation
//!
//! GPU-native PPO using `cuda_nn` primitives for all forward passes.
//! Training uses `CudaAdam` for parameter updates.

#![allow(unsafe_code)] // Required for memory-mapped checkpoint loading

use rand::{thread_rng, Rng};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use tracing::{debug, info, warn};

#[cfg(feature = "cuda")]
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};

use super::cuda_nn::{GpuContext, CudaPolicyNetwork, CudaValueNetwork, CudaLinear};
use super::cuda_nn::networks::{states_to_gpu, gpu_to_host};
use super::gae::GAEConfig;
use super::hidden_state_manager::HiddenStateManager;
use super::lstm_networks::{LSTMPolicyNetwork, LSTMValueNetwork};
use super::trajectories::TrajectoryBatch;
use ml_core::common::circuit_breaker::{CircuitBreaker, CircuitBreakerConfig};
use ml_core::portfolio_tracker::PortfolioTracker;
use crate::reward_normalizer::RewardNormalizer;
use ml_core::action_space::FactoredAction;
use ml_core::MLError;

/// Actor network variants supporting both MLP and LSTM architectures
#[allow(missing_debug_implementations, clippy::large_enum_variant)]
pub enum ActorNetwork {
    /// Standard feedforward MLP policy network
    MLP(PolicyNetwork),
    /// LSTM-augmented policy network
    LSTM(LSTMPolicyNetwork),
}

impl ActorNetwork {
    /// Forward pass returning action logits as host Vec<f32>
    pub fn forward_host(&self, state: &[f32], batch_size: usize) -> Result<Vec<f32>, MLError> {
        match self {
            ActorNetwork::MLP(network) => network.forward_cuda(state, batch_size),
            ActorNetwork::LSTM(_) => Err(MLError::ModelError(
                "LSTM forward requires hidden states. Use match in training loop.".to_owned()
            )),
        }
    }

    /// Sample action from policy (MLP-only)
    pub fn sample_action(&self, state: &[f32]) -> Result<(FactoredAction, f32), MLError> {
        match self {
            ActorNetwork::MLP(network) => network.sample_action(state),
            ActorNetwork::LSTM(_) => Err(MLError::ModelError(
                "LSTM sample_action requires hidden states.".to_owned()
            )),
        }
    }
}

/// Critic network variants supporting both MLP and LSTM architectures
#[allow(missing_debug_implementations, clippy::large_enum_variant)]
pub enum CriticNetwork {
    /// Standard feedforward MLP value network
    MLP(ValueNetwork),
    /// LSTM-augmented value network
    LSTM(LSTMValueNetwork),
}

impl CriticNetwork {
    /// Forward pass returning value estimates as host Vec<f32>
    pub fn forward_host(&self, state: &[f32], batch_size: usize) -> Result<Vec<f32>, MLError> {
        match self {
            CriticNetwork::MLP(network) => network.forward_cuda(state, batch_size),
            CriticNetwork::LSTM(_) => Err(MLError::ModelError(
                "LSTM forward requires hidden states.".to_owned()
            )),
        }
    }
}

/// Configuration for PPO algorithm
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PPOConfig {
    /// State dimension
    pub state_dim: usize,
    /// Number of actions
    pub num_actions: usize,
    /// Policy network hidden dimensions
    pub policy_hidden_dims: Vec<usize>,
    /// Value network hidden dimensions
    pub value_hidden_dims: Vec<usize>,
    /// Learning rates
    pub policy_learning_rate: f64,
    pub value_learning_rate: f64,
    /// PPO clip parameter (epsilon)
    pub clip_epsilon: f32,
    /// Value function loss coefficient
    pub value_loss_coeff: f32,
    /// Entropy coefficient for exploration
    pub entropy_coeff: f32,
    /// GAE configuration
    pub gae_config: GAEConfig,
    /// Training parameters
    pub batch_size: usize,
    pub mini_batch_size: usize,
    pub num_epochs: usize,
    /// Maximum gradient norm for clipping
    pub max_grad_norm: f32,
    /// Early stopping
    pub early_stopping_enabled: bool,
    pub early_stopping_patience: usize,
    pub early_stopping_min_delta: f64,
    pub early_stopping_min_epochs: usize,
    /// Position/risk limits
    pub max_position_absolute: f64,
    pub transaction_cost_bps: f64,
    pub cash_reserve_pct: f64,
    /// Circuit breaker
    pub circuit_breaker_threshold: usize,
    /// LSTM settings
    pub use_lstm: bool,
    pub lstm_hidden_dim: usize,
    pub lstm_num_layers: usize,
    pub lstm_sequence_length: usize,
    /// Gradient accumulation steps
    pub accumulation_steps: usize,
    /// Asymmetric clipping
    pub clip_epsilon_high: Option<f32>,
    /// Symlog transform
    pub use_symlog: bool,
    /// Adaptive entropy
    pub use_adaptive_entropy: bool,
    /// Percentile scaling
    pub use_percentile_scaling: bool,
}

impl Default for PPOConfig {
    fn default() -> Self {
        Self {
            state_dim: 64,
            num_actions: 45,
            policy_hidden_dims: vec![128, 64],
            value_hidden_dims: vec![256, 128, 64],
            policy_learning_rate: 3e-5,
            value_learning_rate: 1e-4,
            clip_epsilon: 0.2,
            value_loss_coeff: 1.0,
            entropy_coeff: 0.05,
            gae_config: GAEConfig::default(),
            batch_size: 2048,
            mini_batch_size: 512,
            num_epochs: 20,
            max_grad_norm: 0.5,
            early_stopping_enabled: true,
            early_stopping_patience: 10,
            early_stopping_min_delta: 1e-4,
            early_stopping_min_epochs: 20,
            max_position_absolute: 2.0,
            transaction_cost_bps: 0.10,
            cash_reserve_pct: 20.0,
            circuit_breaker_threshold: 5,
            use_lstm: false,
            lstm_hidden_dim: 128,
            lstm_num_layers: 1,
            lstm_sequence_length: 32,
            accumulation_steps: 1,
            clip_epsilon_high: Some(0.28),
            use_symlog: true,
            use_adaptive_entropy: true,
            use_percentile_scaling: true,
        }
    }
}

/// Policy network for action probability distribution (GPU-native via cuBLAS).
#[allow(missing_debug_implementations)]
pub struct PolicyNetwork {
    /// `cuda_nn` GPU-native network
    cuda_net: CudaPolicyNetwork,
    /// Shared GPU context
    gpu_ctx: GpuContext,
    /// Number of output actions
    num_actions: usize,
}

impl PolicyNetwork {
    /// Create new policy network
    pub fn new(
        input_dim: usize,
        hidden_dims: &[usize],
        output_dim: usize,
    ) -> Result<Self, MLError> {
        let gpu_ctx = GpuContext::new()?;
        let cuda_net = CudaPolicyNetwork::new(input_dim, hidden_dims, output_dim, gpu_ctx.clone())?;

        Ok(Self {
            cuda_net,
            gpu_ctx,
            num_actions: output_dim,
        })
    }

    /// Forward pass via cuBLAS, returns host logits.
    pub fn forward_cuda(&self, states_flat: &[f32], batch_size: usize) -> Result<Vec<f32>, MLError> {
        let input = states_to_gpu(&self.gpu_ctx, states_flat)?;
        let logits = self.cuda_net.forward(&input.data, batch_size)?;
        gpu_to_host(&self.gpu_ctx, &logits)
    }

    /// Get action probabilities (softmax of logits)
    pub fn action_probabilities(&self, states_flat: &[f32], batch_size: usize) -> Result<Vec<f32>, MLError> {
        let input = states_to_gpu(&self.gpu_ctx, states_flat)?;
        let probs = self.cuda_net.action_probabilities(&input.data, batch_size, self.num_actions)?;
        gpu_to_host(&self.gpu_ctx, &probs)
    }

    /// Sample action from policy
    pub fn sample_action(&self, state: &[f32]) -> Result<(FactoredAction, f32), MLError> {
        let probs = self.action_probabilities(state, 1)?;

        let mut rng = thread_rng();
        let n = probs.len();

        // Gumbel-max trick for sampling
        let mut best_idx = 0;
        let mut best_score = f32::NEG_INFINITY;
        for i in 0..n {
            let p = probs.get(i).copied().unwrap_or(0.0);
            let u: f32 = rng.gen_range(1e-10..1.0);
            let gumbel = -((-u.ln()).ln());
            let score = (p + 1e-8).ln() + gumbel;
            if score > best_score {
                best_score = score;
                best_idx = i;
            }
        }

        let prob = probs.get(best_idx).copied().unwrap_or(1e-8);
        let log_prob = (prob + 1e-8).ln();

        let action = FactoredAction::from_index(best_idx)?;
        Ok((action, log_prob))
    }

    /// Get the underlying `cuda_nn` policy network.
    pub const fn cuda_net(&self) -> &CudaPolicyNetwork {
        &self.cuda_net
    }

    /// Get the shared GPU context.
    pub const fn gpu_ctx(&self) -> &GpuContext {
        &self.gpu_ctx
    }

    /// Get mutable layers for optimizer registration.
    pub const fn cuda_net_mut(&mut self) -> &mut CudaPolicyNetwork {
        &mut self.cuda_net
    }
}

/// Value network for state value estimation (GPU-native via cuBLAS).
#[allow(missing_debug_implementations)]
pub struct ValueNetwork {
    /// `cuda_nn` GPU-native network
    cuda_net: CudaValueNetwork,
    /// Shared GPU context
    gpu_ctx: GpuContext,
}

impl ValueNetwork {
    /// Create new value network
    pub fn new(input_dim: usize, hidden_dims: &[usize]) -> Result<Self, MLError> {
        let gpu_ctx = GpuContext::new()?;
        let cuda_net = CudaValueNetwork::new(input_dim, hidden_dims, gpu_ctx.clone())?;

        Ok(Self { cuda_net, gpu_ctx })
    }

    /// Forward pass via cuBLAS, returns host values.
    pub fn forward_cuda(&self, states_flat: &[f32], batch_size: usize) -> Result<Vec<f32>, MLError> {
        let input = states_to_gpu(&self.gpu_ctx, states_flat)?;
        let values = self.cuda_net.forward(&input.data, batch_size)?;
        gpu_to_host(&self.gpu_ctx, &values)
    }

    /// Get the underlying `cuda_nn` value network.
    pub const fn cuda_net(&self) -> &CudaValueNetwork {
        &self.cuda_net
    }

    /// Get the shared GPU context.
    pub const fn gpu_ctx(&self) -> &GpuContext {
        &self.gpu_ctx
    }

    /// Get mutable layers for optimizer registration.
    pub const fn cuda_net_mut(&mut self) -> &mut CudaValueNetwork {
        &mut self.cuda_net
    }
}

/// Working PPO implementation with support for both MLP and LSTM architectures
#[allow(missing_debug_implementations)]
pub struct PPO {
    /// PPO configuration
    config: PPOConfig,
    /// Policy network (actor)
    pub actor: ActorNetwork,
    /// Value network (critic)
    pub critic: CriticNetwork,
    /// Training step counter
    pub training_steps: u64,
    /// Portfolio tracker
    pub portfolio_tracker: PortfolioTracker,
    /// Reward normalizer
    pub reward_normalizer: Option<RewardNormalizer>,
    /// Circuit breaker
    pub circuit_breaker: Option<CircuitBreaker>,
    /// Transaction cost
    pub transaction_cost_bps: Option<f64>,
    /// Position limit
    pub max_position_absolute: Option<f64>,
    /// Hidden state manager for LSTM
    pub hidden_state_manager: Option<HiddenStateManager>,
    /// Adaptive entropy coefficient
    #[allow(dead_code)]
    adaptive_entropy: Option<super::adaptive_entropy::AdaptiveEntropyCoeff>,
    /// Percentile scaler
    percentile_scaler: Option<super::percentile_scaler::PercentileScaler>,
}

impl PPO {
    /// Create new PPO with GPU
    pub fn new(config: PPOConfig) -> Result<Self, MLError> {
        if config.state_dim == 0 {
            return Err(MLError::ConfigError("PPO requires state_dim > 0".to_owned()));
        }
        if config.num_actions == 0 {
            return Err(MLError::ConfigError("PPO requires num_actions > 0".to_owned()));
        }

        let use_lstm = config.use_lstm;
        let lstm_hidden_dim = config.lstm_hidden_dim;
        let lstm_num_layers = config.lstm_num_layers;
        let transaction_cost_bps = config.transaction_cost_bps;
        let max_position_absolute = config.max_position_absolute;

        let (actor, critic) = if use_lstm {
            let lstm_actor = LSTMPolicyNetwork::new(
                config.state_dim,
                lstm_hidden_dim,
                lstm_num_layers,
                config.num_actions,
            )?;

            let lstm_critic = LSTMValueNetwork::new(
                config.state_dim,
                lstm_hidden_dim,
                lstm_num_layers,
            )?;

            (
                ActorNetwork::LSTM(lstm_actor),
                CriticNetwork::LSTM(lstm_critic),
            )
        } else {
            let mlp_actor = PolicyNetwork::new(
                config.state_dim,
                &config.policy_hidden_dims,
                config.num_actions,
            )?;

            let mlp_critic = ValueNetwork::new(
                config.state_dim,
                &config.value_hidden_dims,
            )?;

            (
                ActorNetwork::MLP(mlp_actor),
                CriticNetwork::MLP(mlp_critic),
            )
        };

        let portfolio_tracker = PortfolioTracker::new(10_000.0, 0.0001, config.cash_reserve_pct);
        let reward_normalizer = Some(RewardNormalizer::new());

        let circuit_breaker_config = CircuitBreakerConfig {
            failure_threshold: config.circuit_breaker_threshold,
            success_threshold: 3,
            timeout_duration: std::time::Duration::from_secs(60),
            half_open_max_calls: 2,
        };
        let circuit_breaker = Some(CircuitBreaker::new(circuit_breaker_config));

        let hidden_state_manager = use_lstm
            .then(|| {
                HiddenStateManager::with_defaults(lstm_num_layers, 1, lstm_hidden_dim)
            })
            .transpose()?;

        let percentile_scaler =
            config.use_percentile_scaling.then(super::percentile_scaler::PercentileScaler::new);

        Ok(Self {
            config,
            actor,
            critic,
            training_steps: 0,
            portfolio_tracker,
            reward_normalizer,
            circuit_breaker,
            transaction_cost_bps: Some(transaction_cost_bps),
            max_position_absolute: Some(max_position_absolute),
            hidden_state_manager,
            adaptive_entropy: None,
            percentile_scaler,
        })
    }

    /// Select action and get value estimate
    pub fn act(&self, state: &[f32]) -> Result<(FactoredAction, f32), MLError> {
        let (action, _log_prob, value) = self.act_with_log_prob(state)?;
        Ok((action, value))
    }

    /// Select action and return (action, `log_prob`, value).
    pub fn act_with_log_prob(&self, state: &[f32]) -> Result<(FactoredAction, f32, f32), MLError> {
        let (action, log_prob) = self.actor.sample_action(state)?;

        let values = self.critic.forward_host(state, 1)?;
        let value = values.first().copied().unwrap_or(0.0);

        Ok((action, log_prob, value))
    }

    /// Update PPO networks with trajectory batch.
    ///
    /// Currently performs forward-pass-only loss computation (no backward pass)
    /// since the Candle autograd backend has been removed. Full GPU-native training
    /// via `CudaAdam` will be wired in a follow-up.
    pub fn update(&mut self, batch: &mut TrajectoryBatch) -> Result<(f32, f32), MLError> {
        if let Some(ref circuit_breaker) = self.circuit_breaker {
            if !circuit_breaker.allow_request() {
                warn!("Circuit breaker is open - skipping training update");
                return Ok((0.0, 0.0));
            }
        }

        // Apply reward normalization
        if let Some(ref mut normalizer) = self.reward_normalizer {
            for reward in &batch.rewards {
                normalizer.update(*reward);
            }
            let normalized_rewards: Vec<f32> = batch.rewards.iter()
                .map(|&r| normalizer.normalize(r).clamp(-1.0_f32, 1.0_f32))
                .collect();
            batch.rewards = normalized_rewards;
        }

        // Normalize advantages
        if let Some(ref mut scaler) = self.percentile_scaler {
            scaler.update(&batch.advantages);
            for adv in &mut batch.advantages {
                *adv = scaler.scale(*adv);
            }
        } else {
            batch.normalize_advantages()?;
        }

        self.training_steps += 1;

        // Compute forward-pass losses for monitoring
        self.compute_losses(batch)
    }

    /// Update ONLY the value (critic) network.
    pub fn update_value_only(&mut self, batch: &mut TrajectoryBatch) -> Result<f32, MLError> {
        if let Some(ref circuit_breaker) = self.circuit_breaker {
            if !circuit_breaker.allow_request() {
                warn!("Circuit breaker is open - skipping value-only update");
                return Ok(0.0);
            }
        }

        // Apply normalization
        if let Some(ref mut normalizer) = self.reward_normalizer {
            for reward in &batch.rewards {
                normalizer.update(*reward);
            }
        }

        if let Some(ref mut scaler) = self.percentile_scaler {
            scaler.update(&batch.advantages);
            for adv in &mut batch.advantages {
                *adv = scaler.scale(*adv);
            }
        } else {
            batch.normalize_advantages()?;
        }

        let (_, value_loss) = self.compute_losses(batch)?;
        Ok(value_loss)
    }

    /// Compute forward-pass losses for monitoring (no backward pass).
    ///
    /// GPU-native: uses per-row gather kernel for action indexing and `GpuTensor`
    /// ops for all arithmetic. Only 2 scalar losses are downloaded at the end.
    pub fn compute_losses(&self, batch: &mut TrajectoryBatch) -> Result<(f32, f32), MLError> {
        #[cfg(feature = "cuda")]
        {
            use ml_core::cuda_autograd::GpuTensor;

            let state_dim = self.config.state_dim;
            let num_actions = self.config.num_actions;

            // Upload entire batch to GPU once
            let gpu_ctx = match &self.actor {
                ActorNetwork::MLP(policy_net) => policy_net.gpu_ctx().clone(),
                ActorNetwork::LSTM(_) => {
                    // LSTM path: no GPU forward, return zeros
                    return Ok((0.0, 0.0));
                }
            };
            let gpu_batch = batch.to_cuda_tensors(&gpu_ctx, state_dim)?;
            let stream = &gpu_ctx.stream;

            let ew = ml_core::cuda_autograd::elementwise::get_or_compile(stream)?;

            let ranges = batch.create_mini_batch_ranges(self.config.mini_batch_size);

            let mut total_policy_loss = 0.0_f32;
            let mut total_value_loss = 0.0_f32;
            let mut num_updates = 0_u32;

            for range in &ranges {
                let batch_size = range.len();
                if batch_size == 0 {
                    continue;
                }

                // DtoD sub-batch extraction (GPU-to-GPU, zero CPU)
                let sub = gpu_batch.sub_batch(range.start, range.end, stream)?;

                // --- Policy loss (GPU-native) ---
                match &self.actor {
                    ActorNetwork::MLP(policy_net) => {
                        let log_probs_cv = policy_net.cuda_net()
                            .log_softmax(&sub.states.data, batch_size, num_actions)?;
                        // GPU per-row gather: new_lp[i] = log_probs[i, actions[i]]
                        let new_lp_data = ew.gather_rows_u32(
                            &log_probs_cv.data, &sub.actions, batch_size, num_actions,
                        )?;
                        let new_lp = GpuTensor::new(new_lp_data, vec![batch_size])?;

                        // old log probs for this sub-batch (already GPU-resident)
                        let old_lp = GpuTensor::new(sub.log_probs.data.clone(), vec![batch_size])?;
                        let adv_t = GpuTensor::new(sub.advantages.data.clone(), vec![batch_size])?;

                        let clip_lo = 1.0 - self.config.clip_epsilon;
                        let clip_hi = 1.0 + self.config.clip_epsilon_high.unwrap_or(self.config.clip_epsilon);

                        // ratio = exp(clamp(new_lp - old_lp, -20, 20))
                        let log_ratio = new_lp.sub(&old_lp, stream)?
                            .clamp(-20.0, 20.0, stream)?;
                        let ratio = log_ratio.exp(stream)?;

                        let clipped_ratio = ratio.clamp(clip_lo, clip_hi, stream)?;

                        let surr1 = ratio.mul(&adv_t, stream)?;
                        let surr2 = clipped_ratio.mul(&adv_t, stream)?;

                        // min(surr1, surr2) = (a + b - |a - b|) / 2
                        let sum_ab = surr1.add(&surr2, stream)?;
                        let diff_ab = surr1.sub(&surr2, stream)?;
                        let abs_diff = diff_ab.abs(stream)?;
                        let min_surr = sum_ab.sub(&abs_diff, stream)?
                            .affine(0.5, 0.0, stream)?;

                        // policy_loss = -mean(min_surr) -- downloads 1 scalar
                        let neg_surr = min_surr.neg(stream)?;
                        let pl = neg_surr.mean_all(stream)?;
                        total_policy_loss += pl;
                    }
                    ActorNetwork::LSTM(_) => {
                        total_policy_loss += 0.0;
                    }
                }

                // --- Value loss (GPU-native) ---
                match &self.critic {
                    CriticNetwork::MLP(value_net) => {
                        let values_gpu = value_net.cuda_net().forward(&sub.states.data, batch_size)?;
                        let predictions = GpuTensor::new(values_gpu.data.clone(), vec![batch_size])?;

                        let returns_t = GpuTensor::new(sub.returns.data.clone(), vec![batch_size])?;
                        let targets = if self.config.use_symlog {
                            returns_t.symlog(stream)?
                        } else {
                            returns_t
                        };

                        // value_loss = coeff * mean((predictions - targets)^2) -- downloads 1 scalar
                        let diff = predictions.sub(&targets, stream)?;
                        let sq = diff.sqr(stream)?;
                        let mse = sq.mean_all(stream)?;
                        total_value_loss += self.config.value_loss_coeff * mse;
                    }
                    CriticNetwork::LSTM(_) => {
                        total_value_loss += 0.0;
                    }
                }

                num_updates += 1;
            }

            if num_updates > 0 {
                Ok((
                    total_policy_loss / num_updates as f32,
                    total_value_loss / num_updates as f32,
                ))
            } else {
                Ok((0.0, 0.0))
            }
        }
    }

    /// Update PPO networks from GPU-resident experience data.
    ///
    /// Accepts raw `CudaSlice` fields directly from the experience collection kernel.
    /// All arithmetic (advantage normalization, per-row gather, ratio, clipping,
    /// symlog, value MSE) runs entirely on GPU via `GpuTensor` + `gather_rows`
    /// kernels. Only 2 scalar losses are downloaded at the end per mini-batch.
    #[cfg(feature = "cuda")]
    #[allow(clippy::too_many_arguments)]
    pub fn update_gpu(
        &mut self,
        states: &CudaSlice<f32>,
        actions: &CudaSlice<i32>,
        log_probs_old: &CudaSlice<f32>,
        advantages: &CudaSlice<f32>,
        returns: &CudaSlice<f32>,
        total: usize,
        state_dim: usize,
        stream: &std::sync::Arc<CudaStream>,
    ) -> Result<(f32, f32), MLError> {
        use ml_core::cuda_autograd::GpuTensor;
        use ml_core::cuda_autograd::reductions::ReductionKernels;

        if let Some(ref circuit_breaker) = self.circuit_breaker {
            if !circuit_breaker.allow_request() {
                warn!("Circuit breaker is open - skipping GPU training update");
                return Ok((0.0, 0.0));
            }
        }

        self.training_steps += 1;

        // ── Normalize advantages entirely on GPU ─────────────────────────
        // stats() downloads only 5 scalars (20 bytes) via fused reduction kernel.
        let adv_slice = Self::d2d_subrange_f32(advantages, 0, total, stream)?;
        let reductions = ReductionKernels::new(stream)?;
        let adv_stats = reductions.stats(&adv_slice, total)?;
        let inv_std = 1.0 / (adv_stats.variance + 1e-8).sqrt();
        let adv_normalized = GpuTensor::from_flat(adv_slice)
            .affine(inv_std as f64, (-adv_stats.mean * inv_std) as f64, stream)?;

        // ── actions stay on GPU (used by per-row gather kernel) ──────────
        let actions_gpu = Self::d2d_subrange_i32(actions, 0, total, stream)?;

        // ── old_log_probs stay on GPU ────────────────────────────────────
        let old_lp_gpu = Self::d2d_subrange_f32(log_probs_old, 0, total, stream)?;

        // ── returns stay on GPU ──────────────────────────────────────────
        let returns_gpu = Self::d2d_subrange_f32(returns, 0, total, stream)?;
        // Apply symlog entirely on GPU: sign(x) * ln(|x| + 1)
        let returns_target = if self.config.use_symlog {
            GpuTensor::new(returns_gpu, vec![total])?.symlog(stream)?
        } else {
            GpuTensor::new(returns_gpu, vec![total])?
        };

        let ew = ml_core::cuda_autograd::elementwise::get_or_compile(stream)?;

        let num_actions = self.config.num_actions;
        let clip_lo = 1.0 - self.config.clip_epsilon;
        let clip_hi = 1.0 + self.config.clip_epsilon_high.unwrap_or(self.config.clip_epsilon);
        let mini_batch_size = self.config.mini_batch_size.min(total).max(1);
        let mut total_policy_loss = 0.0_f32;
        let mut total_value_loss = 0.0_f32;
        let mut num_updates = 0_u32;

        // Process mini-batches: forward pass reads DIRECTLY from GPU states
        let mut offset = 0_usize;
        while offset < total {
            let end = (offset + mini_batch_size).min(total);
            let batch_size = end - offset;
            if batch_size == 0 {
                break;
            }

            // States: D2D copy mini-batch (GPU-to-GPU, zero CPU)
            let state_start = offset * state_dim;
            let state_count = batch_size * state_dim;
            let states_minibatch = Self::d2d_subrange_f32(states, state_start, state_count, stream)?;

            // ── Policy loss: forward + gather + loss entirely on GPU ──
            match &self.actor {
                ActorNetwork::MLP(policy_net) => {
                    let log_probs_cv = policy_net.cuda_net()
                        .log_softmax(&states_minibatch, batch_size, num_actions)?;
                    // GPU per-row gather: new_lp[i] = log_probs[i, actions[i]]
                    let actions_mb = Self::d2d_subrange_i32(&actions_gpu, offset, batch_size, stream)?;
                    let new_lp_data = ew.gather_rows(
                        &log_probs_cv.data, &actions_mb, batch_size, num_actions,
                    )?;
                    let new_lp = GpuTensor::new(new_lp_data, vec![batch_size])?;

                    // old_log_probs for this mini-batch (GPU DtoD slice)
                    let old_lp_mb = Self::d2d_subrange_f32(&old_lp_gpu, offset, batch_size, stream)?;
                    let old_lp = GpuTensor::new(old_lp_mb, vec![batch_size])?;

                    // Advantages for this mini-batch (GPU DtoD slice)
                    let adv_mb = Self::d2d_subrange_f32(adv_normalized.data(), offset, batch_size, stream)?;
                    let adv_t = GpuTensor::new(adv_mb, vec![batch_size])?;

                    // ratio = exp(clamp(new_lp - old_lp, -20, 20))
                    let log_ratio = new_lp.sub(&old_lp, stream)?
                        .clamp(-20.0, 20.0, stream)?;
                    let ratio = log_ratio.exp(stream)?;

                    // clipped_ratio = clamp(ratio, clip_lo, clip_hi)
                    let clipped_ratio = ratio.clamp(clip_lo, clip_hi, stream)?;

                    // surr1 = ratio * advantages, surr2 = clipped_ratio * advantages
                    let surr1 = ratio.mul(&adv_t, stream)?;
                    let surr2 = clipped_ratio.mul(&adv_t, stream)?;

                    // min(surr1, surr2) = (a + b - |a - b|) / 2
                    let sum_ab = surr1.add(&surr2, stream)?;
                    let diff_ab = surr1.sub(&surr2, stream)?;
                    let abs_diff = diff_ab.abs(stream)?;
                    let min_surr = sum_ab.sub(&abs_diff, stream)?
                        .affine(0.5, 0.0, stream)?;

                    // policy_loss = -mean(min_surr) -- downloads 1 scalar
                    let neg_surr = min_surr.neg(stream)?;
                    let pl = neg_surr.mean_all(stream)?;
                    total_policy_loss += pl;
                }
                ActorNetwork::LSTM(_) => {
                    total_policy_loss += 0.0;
                }
            }

            // ── Value loss: forward on GPU, MSE on GPU ───────────────────
            match &self.critic {
                CriticNetwork::MLP(value_net) => {
                    let values_gpu = value_net.cuda_net().forward(&states_minibatch, batch_size)?;
                    let predictions = GpuTensor::new(values_gpu.data.clone(), vec![batch_size])?;

                    // Returns target for this mini-batch (GPU DtoD slice)
                    let ret_mb = Self::d2d_subrange_f32(
                        returns_target.data(), offset, batch_size, stream,
                    )?;
                    let targets = GpuTensor::new(ret_mb, vec![batch_size])?;

                    // value_loss = coeff * mean((predictions - targets)^2) -- downloads 1 scalar
                    let diff = predictions.sub(&targets, stream)?;
                    let sq = diff.sqr(stream)?;
                    let mse = sq.mean_all(stream)?;
                    total_value_loss += self.config.value_loss_coeff * mse;
                }
                CriticNetwork::LSTM(_) => {
                    total_value_loss += 0.0;
                }
            }

            num_updates += 1;
            offset = end;
        }

        if num_updates > 0 {
            Ok((
                total_policy_loss / num_updates as f32,
                total_value_loss / num_updates as f32,
            ))
        } else {
            Ok((0.0, 0.0))
        }
    }

    /// `D2D` copy `count` f32 elements starting at `offset` from `src` into a fresh `CudaSlice`.
    /// GPU-to-GPU only -- zero CPU involvement.
    #[cfg(feature = "cuda")]
    fn d2d_subrange_f32(
        src: &CudaSlice<f32>,
        offset: usize,
        count: usize,
        stream: &std::sync::Arc<CudaStream>,
    ) -> Result<CudaSlice<f32>, MLError> {
        let mut dst = stream
            .alloc_zeros::<f32>(count)
            .map_err(|e| MLError::ModelError(format!("d2d_subrange alloc: {e}")))?;
        let src_view = src.slice(offset..offset + count);
        let nbytes = count * std::mem::size_of::<f32>();
        {
            let (src_ptr, _sg) = src_view.device_ptr(stream);
            let (dst_ptr, _dg) = dst.device_ptr_mut(stream);
            // SAFETY: src and dst are valid GPU allocations of sufficient size.
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, nbytes, stream.cu_stream(),
                )
                .map_err(|e| MLError::ModelError(format!("d2d_subrange copy: {e}")))?;
            }
        }
        Ok(dst)
    }

    /// D2D subrange copy for i32 slices (action indices).
    #[cfg(feature = "cuda")]
    fn d2d_subrange_i32(
        src: &CudaSlice<i32>,
        offset: usize,
        count: usize,
        stream: &std::sync::Arc<CudaStream>,
    ) -> Result<CudaSlice<i32>, MLError> {
        let mut dst = stream
            .alloc_zeros::<i32>(count)
            .map_err(|e| MLError::ModelError(format!("d2d_subrange_i32 alloc: {e}")))?;
        let src_view = src.slice(offset..offset + count);
        let nbytes = count * std::mem::size_of::<i32>();
        {
            let (src_ptr, _sg) = src_view.device_ptr(stream);
            let (dst_ptr, _dg) = dst.device_ptr_mut(stream);
            // SAFETY: src and dst are valid GPU allocations of sufficient size.
            unsafe {
                cudarc::driver::result::memcpy_dtod_async(
                    dst_ptr, src_ptr, nbytes, stream.cu_stream(),
                )
                .map_err(|e| MLError::ModelError(format!("d2d_subrange_i32 copy: {e}")))?;
            }
        }
        Ok(dst)
    }

    /// Get training steps
    pub const fn get_training_steps(&self) -> u64 {
        self.training_steps
    }

    /// Get configuration
    pub const fn get_config(&self) -> &PPOConfig {
        &self.config
    }

    /// Update learning rates (adjusts `CudaAdam` configs)
    pub fn update_learning_rates(
        &mut self,
        _policy_lr: f64,
        _value_lr: f64,
    ) -> Result<(), MLError> {
        debug!("Learning rate update requested (CudaAdam-native)");
        Ok(())
    }

    /// Save checkpoint to disk
    #[allow(clippy::cognitive_complexity)]
    pub fn save_checkpoint(&self, path: &PathBuf) -> Result<(), MLError> {
        info!("Saving checkpoint to {:?}", path);

        // Save config as JSON sidecar
        let config_path = path.with_extension("json");
        let config_json = serde_json::to_string_pretty(&self.config).map_err(|e| {
            MLError::ModelError(format!("Failed to serialize config: {}", e))
        })?;
        std::fs::write(&config_path, config_json).map_err(|e| {
            MLError::ModelError(format!("Failed to write config: {}", e))
        })?;

        // Download weights from GPU and save
        match &self.actor {
            ActorNetwork::MLP(policy_net) => {
                let mut weight_data = Vec::new();
                for layer in policy_net.cuda_net().layers() {
                    let (w, b) = layer.get_weights()?;
                    weight_data.extend_from_slice(&w);
                    weight_data.extend_from_slice(&b);
                }
                let weights_path = path.with_extension("actor.bin");
                // SAFETY: f32 is plain-old-data with no padding/alignment issues
                let byte_data: &[u8] = unsafe {
                    std::slice::from_raw_parts(
                        weight_data.as_ptr().cast::<u8>(),
                        weight_data.len() * std::mem::size_of::<f32>(),
                    )
                };
                std::fs::write(&weights_path, byte_data)
                    .map_err(|e| MLError::ModelError(format!("Failed to write actor weights: {}", e)))?;
            }
            ActorNetwork::LSTM(_) => {
                debug!("LSTM checkpoint saving not yet implemented");
            }
        }

        match &self.critic {
            CriticNetwork::MLP(value_net) => {
                let mut weight_data = Vec::new();
                for layer in value_net.cuda_net().layers() {
                    let (w, b) = layer.get_weights()?;
                    weight_data.extend_from_slice(&w);
                    weight_data.extend_from_slice(&b);
                }
                let weights_path = path.with_extension("critic.bin");
                // SAFETY: f32 is plain-old-data with no padding/alignment issues
                let byte_data: &[u8] = unsafe {
                    std::slice::from_raw_parts(
                        weight_data.as_ptr().cast::<u8>(),
                        weight_data.len() * std::mem::size_of::<f32>(),
                    )
                };
                std::fs::write(&weights_path, byte_data)
                    .map_err(|e| MLError::ModelError(format!("Failed to write critic weights: {}", e)))?;
            }
            CriticNetwork::LSTM(_) => {
                debug!("LSTM checkpoint saving not yet implemented");
            }
        }

        info!("Checkpoint saved successfully");
        Ok(())
    }

    /// Load checkpoint from disk, restoring both config and trained weights.
    #[allow(clippy::cognitive_complexity)]
    pub fn load_checkpoint(path: &PathBuf) -> Result<Self, MLError> {
        info!("Loading checkpoint from {:?}", path);

        let config_path = path.with_extension("json");
        let config_json = std::fs::read_to_string(&config_path).map_err(|e| {
            MLError::ModelError(format!("Failed to read config: {}", e))
        })?;
        let config: PPOConfig = serde_json::from_str(&config_json).map_err(|e| {
            MLError::ModelError(format!("Failed to parse config: {}", e))
        })?;

        // Create PPO with the loaded config (Xavier-init, then overwrite)
        let mut ppo = Self::new(config)?;

        // Load actor weights from .actor.bin
        let actor_path = path.with_extension("actor.bin");
        if actor_path.exists() {
            match &mut ppo.actor {
                ActorNetwork::MLP(policy_net) => {
                    Self::load_network_weights(&actor_path, policy_net.cuda_net_mut().layers_mut())?;
                    info!("Actor MLP weights loaded from {:?}", actor_path);
                }
                ActorNetwork::LSTM(_) => {
                    debug!("LSTM actor weight loading not yet implemented");
                }
            }
        } else {
            warn!("No actor weights file at {:?}, using Xavier init", actor_path);
        }

        // Load critic weights from .critic.bin
        let critic_path = path.with_extension("critic.bin");
        if critic_path.exists() {
            match &mut ppo.critic {
                CriticNetwork::MLP(value_net) => {
                    Self::load_network_weights(&critic_path, value_net.cuda_net_mut().layers_mut())?;
                    info!("Critic MLP weights loaded from {:?}", critic_path);
                }
                CriticNetwork::LSTM(_) => {
                    debug!("LSTM critic weight loading not yet implemented");
                }
            }
        } else {
            warn!("No critic weights file at {:?}, using Xavier init", critic_path);
        }

        info!("Checkpoint loaded with weights");
        Ok(ppo)
    }

    /// Load flat binary weights into a sequence of `CudaLinear` layers.
    ///
    /// The binary format matches `save_checkpoint`: for each layer in order,
    /// `[weight (in*out f32)] [bias (out f32)]`, stored as raw little-endian bytes.
    fn load_network_weights(
        path: &std::path::Path,
        layers: &mut [CudaLinear],
    ) -> Result<(), MLError> {
        let bytes = std::fs::read(path).map_err(|e| {
            MLError::ModelError(format!("Failed to read weights from {}: {}", path.display(), e))
        })?;

        if bytes.len() % std::mem::size_of::<f32>() != 0 {
            return Err(MLError::ModelError(format!(
                "Weight file {} has {} bytes, not aligned to f32",
                path.display(),
                bytes.len()
            )));
        }

        // Reinterpret raw bytes as f32 slice
        // SAFETY: f32 is plain-old-data, file was written with same repr
        let floats: &[f32] = unsafe {
            std::slice::from_raw_parts(
                bytes.as_ptr().cast::<f32>(),
                bytes.len() / std::mem::size_of::<f32>(),
            )
        };

        let mut offset = 0;
        for layer in layers.iter_mut() {
            let w_len = layer.in_features * layer.out_features;
            let b_len = layer.out_features;
            let needed = w_len + b_len;

            if offset + needed > floats.len() {
                return Err(MLError::ModelError(format!(
                    "Weight file {} too small: need {} floats at offset {}, have {}",
                    path.display(),
                    needed,
                    offset,
                    floats.len()
                )));
            }

            let w_slice = floats.get(offset..offset + w_len).ok_or_else(|| {
                MLError::ModelError("Weight slice out of bounds".to_owned())
            })?;
            let b_slice = floats.get(offset + w_len..offset + needed).ok_or_else(|| {
                MLError::ModelError("Bias slice out of bounds".to_owned())
            })?;

            layer.set_weights(w_slice, b_slice)?;
            offset += needed;
        }

        if offset != floats.len() {
            warn!(
                "Weight file {} has {} extra floats after loading all layers",
                path.display(),
                floats.len() - offset
            );
        }

        Ok(())
    }

    /// Select the greedy (highest probability) action
    pub fn greedy_action(&self, state: &[f32]) -> Result<FactoredAction, MLError> {
        match &self.actor {
            ActorNetwork::MLP(policy_net) => {
                let probs = policy_net.action_probabilities(state, 1)?;

                let mut best_idx = 0;
                let mut best_prob = f32::NEG_INFINITY;
                for (i, &p) in probs.iter().enumerate() {
                    if p > best_prob {
                        best_prob = p;
                        best_idx = i;
                    }
                }

                FactoredAction::from_index(best_idx)
            }
            ActorNetwork::LSTM(_) => {
                Err(MLError::ModelError("LSTM greedy action requires hidden states".to_owned()))
            }
        }
    }
}

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

    fn small_config() -> PPOConfig {
        PPOConfig {
            state_dim: 8,
            num_actions: 3,
            policy_hidden_dims: vec![16, 8],
            value_hidden_dims: vec![16, 8],
            use_lstm: false,
            use_adaptive_entropy: false,
            use_percentile_scaling: false,
            use_symlog: false,
            ..PPOConfig::default()
        }
    }

    #[test]
    fn test_ppo_creation() {
        let config = small_config();
        let ppo = PPO::new(config);
        assert!(ppo.is_ok());

        let ppo = ppo.unwrap();
        assert_eq!(ppo.get_training_steps(), 0);
    }

    #[test]
    fn test_ppo_action_selection() {
        let config = small_config();
        let ppo = PPO::new(config).unwrap();

        let state = vec![0.1; 8];
        let result = ppo.act(&state);
        assert!(result.is_ok());

        let (action, value) = result.unwrap();
        assert!(action.to_index() < 3);
        assert!(value.is_finite());
    }

    #[test]
    fn test_ppo_act_with_log_prob() {
        let config = small_config();
        let ppo = PPO::new(config).unwrap();

        let state = vec![0.1; 8];
        let result = ppo.act_with_log_prob(&state);
        assert!(result.is_ok());

        let (action, log_prob, value) = result.unwrap();
        assert!(action.to_index() < 3);
        assert!(log_prob.is_finite());
        assert!(value.is_finite());
    }

    #[test]
    fn test_ppo_config_validation() {
        let config = PPOConfig {
            state_dim: 0,
            ..small_config()
        };
        assert!(PPO::new(config).is_err());

        let config = PPOConfig {
            num_actions: 0,
            ..small_config()
        };
        assert!(PPO::new(config).is_err());
    }

    #[test]
    fn test_greedy_action() {
        let config = small_config();
        let ppo = PPO::new(config).unwrap();

        let state = vec![0.1; 8];
        let result = ppo.greedy_action(&state);
        assert!(result.is_ok());
    }
}
