quantrs2_core/qml/

reinforcement_learning.rs

//! Quantum Reinforcement Learning Algorithms
//!
//! This module implements quantum reinforcement learning algorithms that leverage
//! quantum advantage for policy optimization, value function approximation, and
//! exploration strategies in reinforcement learning tasks.

use crate::{
    error::QuantRS2Result, gate::multi::*, gate::single::*, gate::GateOp, qubit::QubitId,
    variational::VariationalOptimizer,
};
use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::random::{rngs::StdRng, Rng, RngExt, SeedableRng};
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;

/// Configuration for quantum reinforcement learning
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QuantumRLConfig {
    /// Number of qubits for state representation
    pub state_qubits: usize,
    /// Number of qubits for action representation
    pub action_qubits: usize,
    /// Number of qubits for value function
    pub value_qubits: usize,
    /// Learning rate for policy optimization
    pub learning_rate: f64,
    /// Discount factor (gamma)
    pub discount_factor: f64,
    /// Exploration rate (epsilon for epsilon-greedy)
    pub exploration_rate: f64,
    /// Exploration decay rate
    pub exploration_decay: f64,
    /// Minimum exploration rate
    pub min_exploration_rate: f64,
    /// Replay buffer size
    pub replay_buffer_size: usize,
    /// Batch size for training
    pub batch_size: usize,
    /// Number of circuit layers
    pub circuit_depth: usize,
    /// Whether to use quantum advantage techniques
    pub use_quantum_advantage: bool,
    /// Random seed for reproducibility
    pub random_seed: Option<u64>,
}

impl Default for QuantumRLConfig {
    fn default() -> Self {
        Self {
            state_qubits: 4,
            action_qubits: 2,
            value_qubits: 3,
            learning_rate: 0.01,
            discount_factor: 0.99,
            exploration_rate: 1.0,
            exploration_decay: 0.995,
            min_exploration_rate: 0.01,
            replay_buffer_size: 10000,
            batch_size: 32,
            circuit_depth: 6,
            use_quantum_advantage: true,
            random_seed: None,
        }
    }
}

/// Experience tuple for replay buffer
#[derive(Debug, Clone)]
pub struct Experience {
    /// Current state
    pub state: Array1<f64>,
    /// Action taken
    pub action: usize,
    /// Reward received
    pub reward: f64,
    /// Next state
    pub next_state: Array1<f64>,
    /// Whether episode ended
    pub done: bool,
}

/// Replay buffer for experience storage
pub struct ReplayBuffer {
    /// Storage for experiences
    buffer: VecDeque<Experience>,
    /// Maximum buffer size
    max_size: usize,
    /// Random number generator
    rng: StdRng,
}

impl ReplayBuffer {
    /// Create a new replay buffer
    pub fn new(max_size: usize, seed: Option<u64>) -> Self {
        let rng = match seed {
            Some(s) => StdRng::seed_from_u64(s),
            None => StdRng::from_seed([0; 32]), // Use fixed seed for StdRng
        };

        Self {
            buffer: VecDeque::with_capacity(max_size),
            max_size,
            rng,
        }
    }

    /// Add experience to buffer
    pub fn add(&mut self, experience: Experience) {
        if self.buffer.len() >= self.max_size {
            self.buffer.pop_front();
        }
        self.buffer.push_back(experience);
    }

    /// Sample a batch of experiences
    pub fn sample(&mut self, batch_size: usize) -> Vec<Experience> {
        let mut samples = Vec::new();
        let buffer_size = self.buffer.len();

        if buffer_size < batch_size {
            return self.buffer.iter().cloned().collect();
        }

        for _ in 0..batch_size {
            let idx = self.rng.random_range(0..buffer_size);
            samples.push(self.buffer[idx].clone());
        }

        samples
    }

    /// Get current buffer size
    pub fn size(&self) -> usize {
        self.buffer.len()
    }

    /// Check if buffer has enough samples
    pub fn can_sample(&self, batch_size: usize) -> bool {
        self.buffer.len() >= batch_size
    }
}

/// Quantum Deep Q-Network (QDQN) agent
pub struct QuantumDQN {
    /// Configuration
    config: QuantumRLConfig,
    /// Q-network (quantum circuit for value function)
    q_network: QuantumValueNetwork,
    /// Target Q-network for stable training
    target_q_network: QuantumValueNetwork,
    /// Policy network for action selection
    policy_network: QuantumPolicyNetwork,
    /// Replay buffer
    replay_buffer: ReplayBuffer,
    /// Training step counter
    training_steps: usize,
    /// Episode counter
    episodes: usize,
    /// Current exploration rate
    current_exploration_rate: f64,
    /// Random number generator
    rng: StdRng,
}

/// Quantum value network for Q-function approximation
pub struct QuantumValueNetwork {
    /// Quantum circuit for value estimation
    circuit: QuantumValueCircuit,
    /// Variational parameters
    parameters: Array1<f64>,
    /// Optimizer for parameter updates
    optimizer: VariationalOptimizer,
}

/// Quantum policy network for action selection
pub struct QuantumPolicyNetwork {
    /// Quantum circuit for policy
    circuit: QuantumPolicyCircuit,
    /// Variational parameters
    parameters: Array1<f64>,
    /// Optimizer for parameter updates
    optimizer: VariationalOptimizer,
}

/// Quantum circuit for value function approximation
#[derive(Debug, Clone)]
pub struct QuantumValueCircuit {
    /// Number of state qubits
    state_qubits: usize,
    /// Number of value qubits
    value_qubits: usize,
    /// Circuit depth
    depth: usize,
    /// Total number of qubits
    total_qubits: usize,
}

/// Quantum circuit for policy network
#[derive(Debug, Clone)]
pub struct QuantumPolicyCircuit {
    /// Number of state qubits
    state_qubits: usize,
    /// Number of action qubits
    action_qubits: usize,
    /// Circuit depth
    depth: usize,
    /// Total number of qubits
    total_qubits: usize,
}

impl QuantumDQN {
    /// Create a new Quantum DQN agent
    pub fn new(config: QuantumRLConfig) -> QuantRS2Result<Self> {
        let rng = match config.random_seed {
            Some(seed) => StdRng::seed_from_u64(seed),
            None => StdRng::from_seed([0; 32]), // Use fixed seed for StdRng
        };

        // Create Q-network
        let q_network = QuantumValueNetwork::new(&config)?;
        let mut target_q_network = QuantumValueNetwork::new(&config)?;

        // Initialize target network with same parameters
        target_q_network
            .parameters
            .clone_from(&q_network.parameters);

        // Create policy network
        let policy_network = QuantumPolicyNetwork::new(&config)?;

        // Create replay buffer
        let replay_buffer = ReplayBuffer::new(config.replay_buffer_size, config.random_seed);

        Ok(Self {
            config: config.clone(),
            q_network,
            target_q_network,
            policy_network,
            replay_buffer,
            training_steps: 0,
            episodes: 0,
            current_exploration_rate: config.exploration_rate,
            rng,
        })
    }

    /// Select action using epsilon-greedy policy with quantum enhancement
    pub fn select_action(&mut self, state: &Array1<f64>) -> QuantRS2Result<usize> {
        // Epsilon-greedy exploration
        if self.rng.random::<f64>() < self.current_exploration_rate {
            // Random action
            let num_actions = 1 << self.config.action_qubits;
            Ok(self.rng.random_range(0..num_actions))
        } else {
            // Greedy action using quantum policy network
            self.policy_network.get_best_action(state)
        }
    }

    /// Store experience in replay buffer
    pub fn store_experience(&mut self, experience: Experience) {
        self.replay_buffer.add(experience);
    }

    /// Train the agent using quantum advantage techniques
    pub fn train(&mut self) -> QuantRS2Result<TrainingMetrics> {
        if !self.replay_buffer.can_sample(self.config.batch_size) {
            return Ok(TrainingMetrics::default());
        }

        // Sample batch from replay buffer
        let experiences = self.replay_buffer.sample(self.config.batch_size);

        // Prepare training data
        let (states, actions, rewards, next_states, dones) =
            self.prepare_training_data(&experiences);

        // Compute target Q-values using quantum advantage
        let target_q_values = self.compute_target_q_values(&next_states, &rewards, &dones)?;

        // Train Q-network
        let q_loss = self.train_q_network(&states, &actions, &target_q_values)?;

        // Train policy network
        let policy_loss = self.train_policy_network(&states)?;

        // Update target network periodically
        if self.training_steps % 100 == 0 {
            self.update_target_network();
        }

        // Update exploration rate
        self.update_exploration_rate();

        self.training_steps += 1;

        Ok(TrainingMetrics {
            q_loss,
            policy_loss,
            exploration_rate: self.current_exploration_rate,
            training_steps: self.training_steps,
        })
    }

    /// Update target network parameters
    fn update_target_network(&mut self) {
        self.target_q_network.parameters = self.q_network.parameters.clone();
    }

    /// Update exploration rate with decay
    fn update_exploration_rate(&mut self) {
        self.current_exploration_rate = (self.current_exploration_rate
            * self.config.exploration_decay)
            .max(self.config.min_exploration_rate);
    }

    /// Prepare training data from experiences
    fn prepare_training_data(
        &self,
        experiences: &[Experience],
    ) -> (
        Array2<f64>,
        Array1<usize>,
        Array1<f64>,
        Array2<f64>,
        Array1<bool>,
    ) {
        let batch_size = experiences.len();
        let state_dim = experiences[0].state.len();

        let mut states = Array2::zeros((batch_size, state_dim));
        let mut actions = Array1::zeros(batch_size);
        let mut rewards = Array1::zeros(batch_size);
        let mut next_states = Array2::zeros((batch_size, state_dim));
        let mut dones = Array1::from_elem(batch_size, false);

        for (i, exp) in experiences.iter().enumerate() {
            states.row_mut(i).assign(&exp.state);
            actions[i] = exp.action;
            rewards[i] = exp.reward;
            next_states.row_mut(i).assign(&exp.next_state);
            dones[i] = exp.done;
        }

        (states, actions, rewards, next_states, dones)
    }

    /// Compute target Q-values using quantum advantage
    fn compute_target_q_values(
        &self,
        next_states: &Array2<f64>,
        rewards: &Array1<f64>,
        dones: &Array1<bool>,
    ) -> QuantRS2Result<Array1<f64>> {
        let batch_size = next_states.nrows();
        let mut target_q_values = Array1::zeros(batch_size);

        for i in 0..batch_size {
            if dones[i] {
                target_q_values[i] = rewards[i];
            } else {
                let next_state = next_states.row(i).to_owned();
                let max_next_q = self.target_q_network.get_max_q_value(&next_state)?;
                target_q_values[i] = self.config.discount_factor.mul_add(max_next_q, rewards[i]);
            }
        }

        Ok(target_q_values)
    }

    /// Train Q-network using quantum gradients
    fn train_q_network(
        &mut self,
        states: &Array2<f64>,
        actions: &Array1<usize>,
        target_q_values: &Array1<f64>,
    ) -> QuantRS2Result<f64> {
        let batch_size = states.nrows();
        let mut total_loss = 0.0;

        for i in 0..batch_size {
            let state = states.row(i).to_owned();
            let action = actions[i];
            let target = target_q_values[i];

            // Compute current Q-value
            let current_q = self.q_network.get_q_value(&state, action)?;

            // Compute loss (squared error)
            let loss = (current_q - target).powi(2);
            total_loss += loss;

            // Compute quantum gradients and update parameters
            let gradients = self.q_network.compute_gradients(&state, action, target)?;
            self.q_network
                .update_parameters(&gradients, self.config.learning_rate)?;
        }

        Ok(total_loss / batch_size as f64)
    }

    /// Train policy network using quantum policy gradients
    fn train_policy_network(&mut self, states: &Array2<f64>) -> QuantRS2Result<f64> {
        let batch_size = states.nrows();
        let mut total_loss = 0.0;

        for i in 0..batch_size {
            let state = states.row(i).to_owned();

            // Compute policy loss using quantum advantage
            let policy_loss = self
                .policy_network
                .compute_policy_loss(&state, &self.q_network)?;
            total_loss += policy_loss;

            // Update policy parameters
            let gradients = self
                .policy_network
                .compute_policy_gradients(&state, &self.q_network)?;
            self.policy_network
                .update_parameters(&gradients, self.config.learning_rate)?;
        }

        Ok(total_loss / batch_size as f64)
    }

    /// End episode and update statistics
    pub const fn end_episode(&mut self, _total_reward: f64) {
        self.episodes += 1;
    }

    /// Get training statistics
    pub fn get_statistics(&self) -> QLearningStats {
        QLearningStats {
            episodes: self.episodes,
            training_steps: self.training_steps,
            exploration_rate: self.current_exploration_rate,
            replay_buffer_size: self.replay_buffer.size(),
        }
    }
}

impl QuantumValueNetwork {
    /// Create a new quantum value network
    fn new(config: &QuantumRLConfig) -> QuantRS2Result<Self> {
        let circuit = QuantumValueCircuit::new(
            config.state_qubits,
            config.value_qubits,
            config.circuit_depth,
        )?;

        let num_parameters = circuit.get_parameter_count();
        let mut parameters = Array1::zeros(num_parameters);

        // Initialize parameters randomly
        let mut rng = match config.random_seed {
            Some(seed) => StdRng::seed_from_u64(seed),
            None => StdRng::from_seed([0; 32]),
        };

        for param in &mut parameters {
            *param = rng.random_range(-std::f64::consts::PI..std::f64::consts::PI);
        }

        let optimizer = VariationalOptimizer::new(0.01, 0.9);

        Ok(Self {
            circuit,
            parameters,
            optimizer,
        })
    }

    /// Get Q-value for a specific state-action pair
    fn get_q_value(&self, state: &Array1<f64>, action: usize) -> QuantRS2Result<f64> {
        self.circuit
            .evaluate_q_value(state, action, &self.parameters)
    }

    /// Get maximum Q-value for a state over all actions
    fn get_max_q_value(&self, state: &Array1<f64>) -> QuantRS2Result<f64> {
        let num_actions = 1 << self.circuit.get_action_qubits();
        let mut max_q = f64::NEG_INFINITY;

        for action in 0..num_actions {
            let q_value = self.get_q_value(state, action)?;
            max_q = max_q.max(q_value);
        }

        Ok(max_q)
    }

    /// Compute gradients using quantum parameter-shift rule
    fn compute_gradients(
        &self,
        state: &Array1<f64>,
        action: usize,
        target: f64,
    ) -> QuantRS2Result<Array1<f64>> {
        self.circuit
            .compute_parameter_gradients(state, action, target, &self.parameters)
    }

    /// Update parameters using gradients
    fn update_parameters(
        &mut self,
        gradients: &Array1<f64>,
        learning_rate: f64,
    ) -> QuantRS2Result<()> {
        for (param, &grad) in self.parameters.iter_mut().zip(gradients.iter()) {
            *param -= learning_rate * grad;
        }
        Ok(())
    }
}

impl QuantumPolicyNetwork {
    /// Create a new quantum policy network
    fn new(config: &QuantumRLConfig) -> QuantRS2Result<Self> {
        let circuit = QuantumPolicyCircuit::new(
            config.state_qubits,
            config.action_qubits,
            config.circuit_depth,
        )?;

        let num_parameters = circuit.get_parameter_count();
        let mut parameters = Array1::zeros(num_parameters);

        // Initialize parameters randomly
        let mut rng = match config.random_seed {
            Some(seed) => StdRng::seed_from_u64(seed),
            None => StdRng::from_seed([0; 32]),
        };

        for param in &mut parameters {
            *param = rng.random_range(-std::f64::consts::PI..std::f64::consts::PI);
        }

        let optimizer = VariationalOptimizer::new(0.01, 0.9);

        Ok(Self {
            circuit,
            parameters,
            optimizer,
        })
    }

    /// Get best action for a state
    fn get_best_action(&self, state: &Array1<f64>) -> QuantRS2Result<usize> {
        self.circuit.get_best_action(state, &self.parameters)
    }

    /// Compute policy loss
    fn compute_policy_loss(
        &self,
        state: &Array1<f64>,
        q_network: &QuantumValueNetwork,
    ) -> QuantRS2Result<f64> {
        // Use expected Q-value as policy loss (to maximize)
        let action_probs = self
            .circuit
            .get_action_probabilities(state, &self.parameters)?;
        let num_actions = action_probs.len();

        let mut expected_q = 0.0;
        for action in 0..num_actions {
            let q_value = q_network.get_q_value(state, action)?;
            expected_q += action_probs[action] * q_value;
        }

        // Negative because we want to maximize (minimize negative)
        Ok(-expected_q)
    }

    /// Compute policy gradients
    fn compute_policy_gradients(
        &self,
        state: &Array1<f64>,
        q_network: &QuantumValueNetwork,
    ) -> QuantRS2Result<Array1<f64>> {
        self.circuit
            .compute_policy_gradients(state, q_network, &self.parameters)
    }

    /// Update parameters
    fn update_parameters(
        &mut self,
        gradients: &Array1<f64>,
        learning_rate: f64,
    ) -> QuantRS2Result<()> {
        for (param, &grad) in self.parameters.iter_mut().zip(gradients.iter()) {
            *param -= learning_rate * grad;
        }
        Ok(())
    }
}

impl QuantumValueCircuit {
    /// Create a new quantum value circuit
    const fn new(state_qubits: usize, value_qubits: usize, depth: usize) -> QuantRS2Result<Self> {
        let total_qubits = state_qubits + value_qubits;

        Ok(Self {
            state_qubits,
            value_qubits,
            depth,
            total_qubits,
        })
    }

    /// Get number of parameters in the circuit
    const fn get_parameter_count(&self) -> usize {
        // Each layer has rotation gates on each qubit (3 parameters each) plus entangling gates
        let rotations_per_layer = self.get_total_qubits() * 3;
        let entangling_per_layer = self.get_total_qubits(); // Simplified estimate
        self.depth * (rotations_per_layer + entangling_per_layer)
    }

    /// Get total number of qubits
    const fn get_total_qubits(&self) -> usize {
        self.state_qubits + self.value_qubits
    }

    /// Get number of action qubits (for external interface)
    const fn get_action_qubits(&self) -> usize {
        // This is a bit of a hack - in a real implementation,
        // the value circuit wouldn't directly know about actions
        2 // Default action qubits
    }

    /// Evaluate Q-value using quantum circuit
    fn evaluate_q_value(
        &self,
        state: &Array1<f64>,
        action: usize,
        parameters: &Array1<f64>,
    ) -> QuantRS2Result<f64> {
        // Encode state into quantum circuit
        let mut gates = Vec::new();

        // State encoding
        for i in 0..self.state_qubits {
            let state_value = if i < state.len() { state[i] } else { 0.0 };
            gates.push(Box::new(RotationY {
                target: QubitId(i as u32),
                theta: state_value * std::f64::consts::PI,
            }) as Box<dyn GateOp>);
        }

        // Action encoding (simple binary encoding). Only encode action bits
        // that fit within the value register so the gate targets stay inside the
        // `total_qubits`-qubit state vector.
        let action_qubits = self.total_qubits.saturating_sub(self.state_qubits);
        for i in 0..action_qubits {
            if (action >> i) & 1 == 1 {
                gates.push(Box::new(PauliX {
                    target: QubitId((self.state_qubits + i) as u32),
                }) as Box<dyn GateOp>);
            }
        }

        // Variational circuit layers
        let mut param_idx = 0;
        for _layer in 0..self.depth {
            // Rotation layer
            for qubit in 0..self.get_total_qubits() {
                if param_idx + 2 < parameters.len() {
                    gates.push(Box::new(RotationX {
                        target: QubitId(qubit as u32),
                        theta: parameters[param_idx],
                    }) as Box<dyn GateOp>);
                    param_idx += 1;

                    gates.push(Box::new(RotationY {
                        target: QubitId(qubit as u32),
                        theta: parameters[param_idx],
                    }) as Box<dyn GateOp>);
                    param_idx += 1;

                    gates.push(Box::new(RotationZ {
                        target: QubitId(qubit as u32),
                        theta: parameters[param_idx],
                    }) as Box<dyn GateOp>);
                    param_idx += 1;
                }
            }

            // Entangling layer
            for qubit in 0..self.get_total_qubits() - 1 {
                if param_idx < parameters.len() {
                    gates.push(Box::new(CRZ {
                        control: QubitId(qubit as u32),
                        target: QubitId((qubit + 1) as u32),
                        theta: parameters[param_idx],
                    }) as Box<dyn GateOp>);
                    param_idx += 1;
                }
            }
        }

        // Exact state-vector evaluation of the Q-value circuit.
        let q_value = self.simulate_circuit_expectation(&gates)?;

        Ok(q_value)
    }

    /// Simulate the value circuit exactly and return the Pauli-Z expectation of
    /// the readout (last) qubit, the scalar Q-value estimate in `[-1, 1]`.
    fn simulate_circuit_expectation(&self, gates: &[Box<dyn GateOp>]) -> QuantRS2Result<f64> {
        let num_qubits = self.get_total_qubits();
        let state = crate::qml::simulator::simulate(num_qubits, gates)?;
        // Read out the Q-value from the highest-index (value) qubit.
        let readout = num_qubits.saturating_sub(1);
        Ok(crate::qml::simulator::expectation_z(&state, readout))
    }

    /// Compute parameter gradients using parameter-shift rule
    fn compute_parameter_gradients(
        &self,
        state: &Array1<f64>,
        action: usize,
        target: f64,
        parameters: &Array1<f64>,
    ) -> QuantRS2Result<Array1<f64>> {
        let mut gradients = Array1::zeros(parameters.len());
        let shift = std::f64::consts::PI / 2.0;

        for i in 0..parameters.len() {
            // Forward shift
            let mut params_plus = parameters.clone();
            params_plus[i] += shift;
            let q_plus = self.evaluate_q_value(state, action, &params_plus)?;

            // Backward shift
            let mut params_minus = parameters.clone();
            params_minus[i] -= shift;
            let q_minus = self.evaluate_q_value(state, action, &params_minus)?;

            // Parameter-shift rule gradient
            let current_q = self.evaluate_q_value(state, action, parameters)?;
            let loss_gradient = 2.0 * (current_q - target); // d/dθ (q - target)²

            gradients[i] = loss_gradient * (q_plus - q_minus) / 2.0;
        }

        Ok(gradients)
    }
}

impl QuantumPolicyCircuit {
    /// Create a new quantum policy circuit
    const fn new(state_qubits: usize, action_qubits: usize, depth: usize) -> QuantRS2Result<Self> {
        let total_qubits = state_qubits + action_qubits;

        Ok(Self {
            state_qubits,
            action_qubits,
            depth,
            total_qubits,
        })
    }

    /// Get number of parameters
    const fn get_parameter_count(&self) -> usize {
        let total_qubits = self.state_qubits + self.action_qubits;
        let rotations_per_layer = total_qubits * 3;
        let entangling_per_layer = total_qubits;
        self.depth * (rotations_per_layer + entangling_per_layer)
    }

    /// Get best action for a state
    fn get_best_action(
        &self,
        state: &Array1<f64>,
        parameters: &Array1<f64>,
    ) -> QuantRS2Result<usize> {
        let action_probs = self.get_action_probabilities(state, parameters)?;

        // Find action with highest probability
        let mut best_action = 0;
        let mut best_prob = action_probs[0];

        for (action, &prob) in action_probs.iter().enumerate() {
            if prob > best_prob {
                best_prob = prob;
                best_action = action;
            }
        }

        Ok(best_action)
    }

    /// Get action probabilities by exactly simulating the policy circuit and
    /// marginalising the amplitude vector over the action register.
    ///
    /// The action register occupies qubits `state_qubits .. total_qubits`. The
    /// probability of an action `a` is the sum of `|amplitude|²` over every
    /// basis state whose action bits equal `a` (the state register is traced
    /// out). The result is a genuine, parameter-dependent Born-rule
    /// distribution, not a heuristic.
    fn get_action_probabilities(
        &self,
        state: &Array1<f64>,
        parameters: &Array1<f64>,
    ) -> QuantRS2Result<Vec<f64>> {
        let num_actions = 1usize << self.action_qubits;
        let gates = self.build_policy_circuit(state, parameters);
        let amplitudes = crate::qml::simulator::simulate(self.total_qubits, &gates)?;

        let mut probabilities = vec![0.0; num_actions];
        for (basis_index, amplitude) in amplitudes.iter().enumerate() {
            // Extract the action sub-register (qubits state_qubits..total_qubits).
            let action = (basis_index >> self.state_qubits) & (num_actions - 1);
            probabilities[action] += amplitude.norm_sqr();
        }

        // Renormalise defensively against floating point drift.
        let sum: f64 = probabilities.iter().sum();
        if sum > 0.0 {
            for prob in &mut probabilities {
                *prob /= sum;
            }
        }

        Ok(probabilities)
    }

    /// Build the parameterised policy circuit: amplitude-style state encoding on
    /// the state register followed by `depth` rotation+entangling layers across
    /// all qubits. Mirrors the structure of the value circuit so the
    /// parameter-shift gradients are exact.
    fn build_policy_circuit(
        &self,
        state: &Array1<f64>,
        parameters: &Array1<f64>,
    ) -> Vec<Box<dyn GateOp>> {
        let mut gates: Vec<Box<dyn GateOp>> = Vec::new();

        // State encoding on the state register.
        for i in 0..self.state_qubits {
            let state_value = if i < state.len() { state[i] } else { 0.0 };
            gates.push(Box::new(RotationY {
                target: QubitId(i as u32),
                theta: state_value * std::f64::consts::PI,
            }) as Box<dyn GateOp>);
        }

        // Variational layers over the full register.
        let mut param_idx = 0;
        for _layer in 0..self.depth {
            for qubit in 0..self.total_qubits {
                if param_idx + 2 < parameters.len() {
                    gates.push(Box::new(RotationX {
                        target: QubitId(qubit as u32),
                        theta: parameters[param_idx],
                    }) as Box<dyn GateOp>);
                    param_idx += 1;
                    gates.push(Box::new(RotationY {
                        target: QubitId(qubit as u32),
                        theta: parameters[param_idx],
                    }) as Box<dyn GateOp>);
                    param_idx += 1;
                    gates.push(Box::new(RotationZ {
                        target: QubitId(qubit as u32),
                        theta: parameters[param_idx],
                    }) as Box<dyn GateOp>);
                    param_idx += 1;
                }
            }

            for qubit in 0..self.total_qubits.saturating_sub(1) {
                if param_idx < parameters.len() {
                    gates.push(Box::new(CRZ {
                        control: QubitId(qubit as u32),
                        target: QubitId((qubit + 1) as u32),
                        theta: parameters[param_idx],
                    }) as Box<dyn GateOp>);
                    param_idx += 1;
                }
            }
        }

        gates
    }

    /// Compute policy gradients
    fn compute_policy_gradients(
        &self,
        state: &Array1<f64>,
        q_network: &QuantumValueNetwork,
        parameters: &Array1<f64>,
    ) -> QuantRS2Result<Array1<f64>> {
        let mut gradients = Array1::zeros(parameters.len());
        let shift = std::f64::consts::PI / 2.0;

        for i in 0..parameters.len() {
            // Forward shift
            let mut params_plus = parameters.clone();
            params_plus[i] += shift;
            let loss_plus = self.compute_policy_loss_with_params(state, q_network, &params_plus)?;

            // Backward shift
            let mut params_minus = parameters.clone();
            params_minus[i] -= shift;
            let loss_minus =
                self.compute_policy_loss_with_params(state, q_network, &params_minus)?;

            // Parameter-shift rule
            gradients[i] = (loss_plus - loss_minus) / 2.0;
        }

        Ok(gradients)
    }

    /// Compute policy loss with specific parameters
    fn compute_policy_loss_with_params(
        &self,
        state: &Array1<f64>,
        q_network: &QuantumValueNetwork,
        parameters: &Array1<f64>,
    ) -> QuantRS2Result<f64> {
        let action_probs = self.get_action_probabilities(state, parameters)?;
        let num_actions = action_probs.len();

        let mut expected_q = 0.0;
        for action in 0..num_actions {
            let q_value = q_network.get_q_value(state, action)?;
            expected_q += action_probs[action] * q_value;
        }

        Ok(-expected_q) // Negative to maximize
    }
}

/// Training metrics for quantum RL
#[derive(Debug, Clone, Default)]
pub struct TrainingMetrics {
    /// Q-network loss
    pub q_loss: f64,
    /// Policy network loss
    pub policy_loss: f64,
    /// Current exploration rate
    pub exploration_rate: f64,
    /// Number of training steps
    pub training_steps: usize,
}

/// Q-learning statistics
#[derive(Debug, Clone)]
pub struct QLearningStats {
    /// Number of episodes completed
    pub episodes: usize,
    /// Number of training steps
    pub training_steps: usize,
    /// Current exploration rate
    pub exploration_rate: f64,
    /// Current replay buffer size
    pub replay_buffer_size: usize,
}

/// Quantum Actor-Critic agent
pub struct QuantumActorCritic {
    /// Configuration
    config: QuantumRLConfig,
    /// Actor network (policy)
    actor: QuantumPolicyNetwork,
    /// Critic network (value function)
    critic: QuantumValueNetwork,
    /// Training metrics
    metrics: TrainingMetrics,
}

impl QuantumActorCritic {
    /// Create a new Quantum Actor-Critic agent
    pub fn new(config: QuantumRLConfig) -> QuantRS2Result<Self> {
        let actor = QuantumPolicyNetwork::new(&config)?;
        let critic = QuantumValueNetwork::new(&config)?;

        Ok(Self {
            config,
            actor,
            critic,
            metrics: TrainingMetrics::default(),
        })
    }

    /// Update networks using actor-critic algorithm
    pub fn update(
        &mut self,
        state: &Array1<f64>,
        _action: usize,
        reward: f64,
        next_state: &Array1<f64>,
        done: bool,
    ) -> QuantRS2Result<()> {
        // Compute TD error
        let current_value = self.critic.get_q_value(state, 0)?; // Use first action for state value
        let next_value = if done {
            0.0
        } else {
            self.critic.get_max_q_value(next_state)?
        };

        let target_value = self.config.discount_factor.mul_add(next_value, reward);
        let td_error = target_value - current_value;

        // Update critic
        let critic_gradients = self.critic.compute_gradients(state, 0, target_value)?;
        self.critic
            .update_parameters(&critic_gradients, self.config.learning_rate)?;

        // Update actor using policy gradient scaled by TD error
        let actor_gradients = self.actor.compute_policy_gradients(state, &self.critic)?;
        let scaled_gradients = actor_gradients * td_error; // Scale by advantage
        self.actor
            .update_parameters(&scaled_gradients, self.config.learning_rate)?;

        // Update metrics
        self.metrics.q_loss = td_error.abs();
        self.metrics.policy_loss = -td_error; // Negative because we want to maximize

        Ok(())
    }

    /// Select action using current policy
    pub fn select_action(&self, state: &Array1<f64>) -> QuantRS2Result<usize> {
        self.actor.get_best_action(state)
    }

    /// Get training metrics
    pub const fn get_metrics(&self) -> &TrainingMetrics {
        &self.metrics
    }
}

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

    #[test]
    fn test_quantum_dqn_creation() {
        let config = QuantumRLConfig::default();
        let agent = QuantumDQN::new(config).expect("Failed to create QuantumDQN agent");

        let stats = agent.get_statistics();
        assert_eq!(stats.episodes, 0);
        assert_eq!(stats.training_steps, 0);
    }

    #[test]
    fn test_replay_buffer() {
        let mut buffer = ReplayBuffer::new(10, Some(42));

        let experience = Experience {
            state: Array1::from_vec(vec![1.0, 0.0, -1.0]),
            action: 1,
            reward: 1.0,
            next_state: Array1::from_vec(vec![0.0, 1.0, 0.0]),
            done: false,
        };

        buffer.add(experience);
        assert_eq!(buffer.size(), 1);

        let samples = buffer.sample(1);
        assert_eq!(samples.len(), 1);
    }

    #[test]
    fn test_quantum_value_circuit() {
        let circuit =
            QuantumValueCircuit::new(3, 2, 4).expect("Failed to create QuantumValueCircuit");
        let param_count = circuit.get_parameter_count();
        assert!(param_count > 0);

        let state = Array1::from_vec(vec![0.5, -0.5, 0.0]);
        let parameters = Array1::zeros(param_count);

        let q_value = circuit
            .evaluate_q_value(&state, 1, &parameters)
            .expect("Failed to evaluate Q-value");
        assert!(q_value.is_finite());
    }

    #[test]
    fn test_quantum_actor_critic() {
        let config = QuantumRLConfig::default();
        let mut agent =
            QuantumActorCritic::new(config).expect("Failed to create QuantumActorCritic agent");

        let state = Array1::from_vec(vec![0.5, -0.5]);
        let next_state = Array1::from_vec(vec![0.0, 1.0]);

        let action = agent
            .select_action(&state)
            .expect("Failed to select action");
        assert!(action < 4); // 2^2 actions for 2 action qubits

        agent
            .update(&state, action, 1.0, &next_state, false)
            .expect("Failed to update agent");

        let metrics = agent.get_metrics();
        assert!(metrics.q_loss >= 0.0);
    }

    #[test]
    fn test_quantum_rl_config_default() {
        let config = QuantumRLConfig::default();
        assert_eq!(config.state_qubits, 4);
        assert_eq!(config.action_qubits, 2);
        assert!(config.learning_rate > 0.0);
        assert!(config.discount_factor < 1.0);
    }

    #[test]
    fn test_q_value_is_real_simulation_not_hash() {
        // The old fabrication returned a hash of the gate matrices, independent
        // of the parameter *values* in a smooth way. A real simulation gives an
        // expectation in [-1, 1] that varies smoothly and continuously with the
        // parameters. Verify both bounds and parameter-dependence.
        let circuit =
            QuantumValueCircuit::new(2, 1, 2).expect("Failed to create QuantumValueCircuit");
        let param_count = circuit.get_parameter_count();
        let state = Array1::from_vec(vec![0.3, -0.7]);

        let q0 = circuit
            .evaluate_q_value(&state, 0, &Array1::zeros(param_count))
            .expect("q0");
        assert!((-1.0..=1.0).contains(&q0), "q value must be in [-1,1]");

        // Perturb all parameters; a real simulation's readout expectation must
        // change (the old hash-based fabrication had no smooth dependence).
        let mut params = Array1::zeros(param_count);
        for (i, p) in params.iter_mut().enumerate() {
            *p = 0.3 * (i as f64 + 1.0);
        }
        let q1 = circuit.evaluate_q_value(&state, 0, &params).expect("q1");
        assert!((-1.0..=1.0).contains(&q1));
        assert!(
            (q0 - q1).abs() > 1e-6,
            "Q-value must respond to parameters (real simulation), got {q0} vs {q1}"
        );
    }

    #[test]
    fn test_action_probabilities_are_real_born_distribution() {
        // The old fabrication returned (uniform + 0.1*sin(...)) noise. A real
        // simulation yields a normalised Born distribution that responds to the
        // parameters and is NOT the near-uniform sin pattern.
        let circuit = QuantumPolicyCircuit::new(2, 2, 2).expect("policy circuit");
        let param_count = circuit.get_parameter_count();
        let state = Array1::from_vec(vec![0.5, 0.5]);

        let probs_zero = circuit
            .get_action_probabilities(&state, &Array1::zeros(param_count))
            .expect("probs zero");
        // Must be a valid distribution.
        let sum: f64 = probs_zero.iter().sum();
        assert!((sum - 1.0).abs() < 1e-9, "probabilities must sum to 1");
        assert!(probs_zero
            .iter()
            .all(|&p| (-1e-12..=1.0 + 1e-9).contains(&p)));
        assert_eq!(probs_zero.len(), 1 << 2);

        // With non-trivial parameters the distribution must change.
        let mut params = Array1::zeros(param_count);
        for (i, p) in params.iter_mut().enumerate() {
            *p = 0.4 * (i as f64 + 1.0);
        }
        let probs_param = circuit
            .get_action_probabilities(&state, &params)
            .expect("probs param");
        let max_diff = probs_zero
            .iter()
            .zip(probs_param.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0_f64, f64::max);
        assert!(
            max_diff > 1e-6,
            "action distribution must respond to parameters (real simulation)"
        );
    }
}