//! GPU-native policy and value networks using cuBLAS-backed linear layers.
//!
//! Drop-in replacements for the Candle-based `PolicyNetwork` and `ValueNetwork`
//! in `ppo.rs`. All computation stays on GPU via cuBLAS sgemm + CUDA kernels.

use cudarc::driver::CudaSlice;

use ml_core::MLError;

use super::linear::CudaLinear;
use super::activations::cuda_relu;
use super::softmax::{cuda_softmax, cuda_log_softmax};
use super::tensor_util::{CudaVec, cuda_from_slice};
use super::GpuContext;


// ---------------------------------------------------------------------------
// CudaPolicyNetwork
// ---------------------------------------------------------------------------

/// GPU-native policy network: `state -> action logits`
///
/// Architecture: `Linear -> ReLU -> Linear -> ReLU -> ... -> Linear (logits)`
/// All layers use cuBLAS sgemm internally. No Candle tensors.
pub struct CudaPolicyNetwork {
    layers: Vec<CudaLinear>,
    ctx: GpuContext,
}

impl std::fmt::Debug for CudaPolicyNetwork {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CudaPolicyNetwork")
            .field("num_layers", &self.layers.len())
            .finish()
    }
}

impl CudaPolicyNetwork {
    /// Create a new policy network.
    ///
    /// # Arguments
    /// * `input_dim` - State dimension
    /// * `hidden_dims` - Hidden layer dimensions
    /// * `output_dim` - Number of actions
    /// * `ctx` - Shared GPU context
    pub fn new(
        input_dim: usize,
        hidden_dims: &[usize],
        output_dim: usize,
        ctx: GpuContext,
    ) -> Result<Self, MLError> {
        let mut layers = Vec::new();
        let mut current_dim = input_dim;

        for &hidden_dim in hidden_dims {
            layers.push(CudaLinear::new(ctx.clone(), current_dim, hidden_dim)?);
            current_dim = hidden_dim;
        }

        // Output layer (logits, no activation)
        layers.push(CudaLinear::new(ctx.clone(), current_dim, output_dim)?);

        Ok(Self { layers, ctx })
    }

    /// Forward pass: state -> logits
    ///
    /// Input: `[batch_size * input_dim]` flattened row-major
    /// Output: `CudaVec` of `[batch_size * num_actions]`
    pub fn forward(&self, input: &CudaSlice<f32>, batch_size: usize) -> Result<CudaVec, MLError> {
        let mut x = self.layers.first()
            .ok_or_else(|| MLError::ConfigError("Policy network has no layers".to_owned()))?
            .forward(input, batch_size)?;

        for layer in self.layers.iter().skip(1) {
            // ReLU activation (not on last layer)
            let last = self.layers.last()
                .ok_or_else(|| MLError::ConfigError("Policy network has no layers".to_owned()))?;
            let is_last = std::ptr::eq(layer, last);
            if !is_last {
                // Apply ReLU to previous layer output, then run next linear
                let relu_out = cuda_relu(&self.ctx.stream, &x.data, x.len)?;
                x = layer.forward(&relu_out.data, batch_size)?;
            } else {
                // Apply ReLU to the penultimate output, then pass to final layer
                let relu_out = cuda_relu(&self.ctx.stream, &x.data, x.len)?;
                x = layer.forward(&relu_out.data, batch_size)?;
            }
        }

        Ok(x)
    }

    /// Compute action probabilities: softmax(logits)
    pub fn action_probabilities(&self, input: &CudaSlice<f32>, batch_size: usize, num_actions: usize) -> Result<CudaVec, MLError> {
        let logits = self.forward(input, batch_size)?;
        cuda_softmax(&self.ctx.stream, &logits.data, batch_size, num_actions)
    }

    /// Compute log-softmax of logits (for log-probability computation)
    pub fn log_softmax(&self, input: &CudaSlice<f32>, batch_size: usize, num_actions: usize) -> Result<CudaVec, MLError> {
        let logits = self.forward(input, batch_size)?;
        cuda_log_softmax(&self.ctx.stream, &logits.data, batch_size, num_actions)
    }

    /// Get all layers for optimizer registration.
    pub fn layers_mut(&mut self) -> &mut [CudaLinear] {
        &mut self.layers
    }

    /// Get all layers (immutable).
    pub fn layers(&self) -> &[CudaLinear] {
        &self.layers
    }

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

// ---------------------------------------------------------------------------
// CudaValueNetwork
// ---------------------------------------------------------------------------

/// GPU-native value network: `state -> scalar value`
///
/// Architecture: `Linear -> ReLU -> Linear -> ReLU -> ... -> Linear (1)`
pub struct CudaValueNetwork {
    layers: Vec<CudaLinear>,
    ctx: GpuContext,
}

impl std::fmt::Debug for CudaValueNetwork {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CudaValueNetwork")
            .field("num_layers", &self.layers.len())
            .finish()
    }
}

impl CudaValueNetwork {
    /// Create a new value network.
    ///
    /// # Arguments
    /// * `input_dim` - State dimension
    /// * `hidden_dims` - Hidden layer dimensions
    /// * `ctx` - Shared GPU context
    pub fn new(
        input_dim: usize,
        hidden_dims: &[usize],
        ctx: GpuContext,
    ) -> Result<Self, MLError> {
        let mut layers = Vec::new();
        let mut current_dim = input_dim;

        for &hidden_dim in hidden_dims {
            layers.push(CudaLinear::new(ctx.clone(), current_dim, hidden_dim)?);
            current_dim = hidden_dim;
        }

        // Output layer: single value
        layers.push(CudaLinear::new(ctx.clone(), current_dim, 1)?);

        Ok(Self { layers, ctx })
    }

    /// Forward pass: state -> value `[batch_size]`
    ///
    /// Input: `[batch_size * input_dim]` flattened row-major
    /// Output: `CudaVec` of `[batch_size]` (one value per sample)
    pub fn forward(&self, input: &CudaSlice<f32>, batch_size: usize) -> Result<CudaVec, MLError> {
        let mut x = self.layers.first()
            .ok_or_else(|| MLError::ConfigError("Value network has no layers".to_owned()))?
            .forward(input, batch_size)?;

        for layer in self.layers.iter().skip(1) {
            let relu_out = cuda_relu(&self.ctx.stream, &x.data, x.len)?;
            x = layer.forward(&relu_out.data, batch_size)?;
        }

        // Output is [batch_size, 1], the CudaVec has batch_size elements
        // since CudaLinear outputs [batch * 1]
        Ok(x)
    }

    /// Get all layers for optimizer registration.
    pub fn layers_mut(&mut self) -> &mut [CudaLinear] {
        &mut self.layers
    }

    /// Get all layers (immutable).
    pub fn layers(&self) -> &[CudaLinear] {
        &self.layers
    }

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

// ---------------------------------------------------------------------------
// Conversion utilities: Candle Tensor <-> CudaSlice
// ---------------------------------------------------------------------------

/// Upload a flat f32 slice to GPU as `CudaVec`.
///
/// This is the bridge between host-side trajectory data (Vec<f32>) and
/// GPU-native network inputs.
pub fn states_to_gpu(
    ctx: &GpuContext,
    states_flat: &[f32],
) -> Result<CudaVec, MLError> {
    cuda_from_slice(&ctx.stream, states_flat)
}

/// Download a `CudaVec` to host `Vec<f32>`.
///
/// Used at API boundaries where callers need host-side results
/// (e.g., action selection during trajectory collection).
pub fn gpu_to_host(
    ctx: &GpuContext,
    gpu_vec: &CudaVec,
) -> Result<Vec<f32>, MLError> {
    gpu_vec.to_vec(&ctx.stream)
}
