quantrs2-device 0.1.3

use super::{
    AtomStateEncoding, NeutralAtomCircuitResult, NeutralAtomClient, NeutralAtomDeviceConfig,
    NeutralAtomExecutionMetadata, NeutralAtomMeasurementData, NeutralAtomQuantumDevice,
    NeutralAtomSystemType,
};
use crate::{Circuit, CircuitExecutor, CircuitResult, DeviceError, DeviceResult, QuantumDevice};
use async_trait::async_trait;
use scirs2_core::random::prelude::*;
use std::collections::HashMap;
use std::time::Duration;

#[derive(Debug, Clone)]
pub struct NeutralAtomDevice {
    pub client: NeutralAtomClient,
    pub device_id: String,
    pub config: NeutralAtomDeviceConfig,
}

impl NeutralAtomDevice {
    pub const fn new(
        client: NeutralAtomClient,
        device_id: String,
        config: NeutralAtomDeviceConfig,
    ) -> Self {
        Self {
            client,
            device_id,
            config,
        }
    }

    #[must_use]
    pub fn with_config(mut self, config: NeutralAtomDeviceConfig) -> Self {
        self.config = config;
        self
    }
}

#[async_trait]
impl QuantumDevice for NeutralAtomDevice {
    async fn is_available(&self) -> DeviceResult<bool> {
        // Simplified implementation - would query device status
        Ok(true)
    }

    async fn qubit_count(&self) -> DeviceResult<usize> {
        Ok(self.config.atom_count)
    }

    async fn properties(&self) -> DeviceResult<HashMap<String, String>> {
        let mut props = HashMap::new();

        props.insert("device_id".to_string(), self.device_id.clone());
        props.insert(
            "system_type".to_string(),
            format!("{:?}", self.config.system_type),
        );
        props.insert("atom_count".to_string(), self.config.atom_count.to_string());
        props.insert(
            "atom_spacing".to_string(),
            self.config.atom_spacing.to_string(),
        );
        props.insert(
            "state_encoding".to_string(),
            format!("{:?}", self.config.state_encoding),
        );

        if let Some(blockade_radius) = self.config.blockade_radius {
            props.insert("blockade_radius".to_string(), blockade_radius.to_string());
        }

        if let Some(loading_efficiency) = self.config.loading_efficiency {
            props.insert(
                "loading_efficiency".to_string(),
                loading_efficiency.to_string(),
            );
        }

        if let Some(gate_fidelity) = self.config.gate_fidelity {
            props.insert("gate_fidelity".to_string(), gate_fidelity.to_string());
        }

        Ok(props)
    }

    async fn is_simulator(&self) -> DeviceResult<bool> {
        Ok(self.device_id.to_lowercase().contains("simulator")
            || self.device_id.to_lowercase().contains("emulator"))
    }
}

#[async_trait]
impl CircuitExecutor for NeutralAtomDevice {
    async fn execute_circuit<const N: usize>(
        &self,
        circuit: &Circuit<N>,
        shots: usize,
    ) -> DeviceResult<CircuitResult> {
        let result = self
            .execute_neutral_atom_circuit(circuit, shots, None)
            .await?;

        Ok(result.circuit_result)
    }

    async fn execute_circuits<const N: usize>(
        &self,
        circuits: Vec<&Circuit<N>>,
        shots: usize,
    ) -> DeviceResult<Vec<CircuitResult>> {
        let mut results = Vec::new();

        for circuit in circuits {
            let result = self.execute_circuit(circuit, shots).await?;
            results.push(result);
        }

        Ok(results)
    }

    async fn can_execute_circuit<const N: usize>(
        &self,
        circuit: &Circuit<N>,
    ) -> DeviceResult<bool> {
        // Check if the circuit can fit on this device
        let required_qubits = N;
        let available_qubits = self.config.atom_count;

        if required_qubits > available_qubits {
            return Ok(false);
        }

        // Check if the device supports the required operations
        // For now, assume all neutral atom devices can execute basic circuits
        Ok(true)
    }

    async fn estimated_queue_time<const N: usize>(
        &self,
        _circuit: &Circuit<N>,
    ) -> DeviceResult<std::time::Duration> {
        // Simplified implementation - would normally query device queue
        Ok(std::time::Duration::from_secs(30))
    }
}

#[async_trait]
impl NeutralAtomQuantumDevice for NeutralAtomDevice {
    async fn system_type(&self) -> DeviceResult<NeutralAtomSystemType> {
        Ok(self.config.system_type)
    }

    async fn atom_count(&self) -> DeviceResult<usize> {
        Ok(self.config.atom_count)
    }

    async fn atom_spacing(&self) -> DeviceResult<f64> {
        Ok(self.config.atom_spacing)
    }

    async fn state_encoding(&self) -> DeviceResult<AtomStateEncoding> {
        Ok(self.config.state_encoding)
    }

    async fn blockade_radius(&self) -> DeviceResult<Option<f64>> {
        Ok(self.config.blockade_radius)
    }

    async fn supports_rydberg_gates(&self) -> DeviceResult<bool> {
        Ok(matches!(
            self.config.system_type,
            NeutralAtomSystemType::Rydberg | NeutralAtomSystemType::Hybrid
        ))
    }

    async fn supports_tweezer_manipulation(&self) -> DeviceResult<bool> {
        Ok(matches!(
            self.config.system_type,
            NeutralAtomSystemType::OpticalTweezer | NeutralAtomSystemType::Hybrid
        ))
    }

    async fn loading_efficiency(&self) -> DeviceResult<f64> {
        Ok(self.config.loading_efficiency.unwrap_or(0.95))
    }

    async fn gate_fidelity(&self) -> DeviceResult<f64> {
        Ok(self.config.gate_fidelity.unwrap_or(0.995))
    }

    async fn execute_neutral_atom_circuit<const N: usize>(
        &self,
        circuit: &Circuit<N>,
        shots: usize,
        config: Option<NeutralAtomDeviceConfig>,
    ) -> DeviceResult<NeutralAtomCircuitResult> {
        let _job_config = config.unwrap_or_else(|| self.config.clone());

        // Simplified implementation - would normally submit to quantum hardware
        let measurement_data = NeutralAtomMeasurementData::default();
        let execution_metadata = NeutralAtomExecutionMetadata {
            system_type: self.config.system_type,
            atoms_used: self.config.atom_count,
            execution_time: Duration::from_millis(100),
            gate_sequence: vec!["X".to_string(), "CNOT".to_string()],
            optimizations_applied: vec!["rydberg_blockade".to_string()],
            temperature: Some(1e-6),
            laser_power: Some(10.0),
        };

        let circuit_result = CircuitResult {
            counts: {
                let mut counts = HashMap::new();
                // Generate mock measurement results
                let all_zeros = "0".repeat(N);
                counts.insert(all_zeros, shots);
                counts
            },
            shots,
            metadata: {
                let mut metadata = HashMap::new();
                metadata.insert("execution_time_ms".to_string(), "100".to_string());
                metadata.insert("success".to_string(), "true".to_string());
                metadata.insert(
                    "system_type".to_string(),
                    format!("{:?}", self.config.system_type),
                );
                metadata
            },
        };

        Ok(NeutralAtomCircuitResult {
            circuit_result,
            neutral_atom_data: measurement_data,
            execution_metadata,
        })
    }

    async fn load_atoms(&self, _positions: &[(f64, f64, f64)]) -> DeviceResult<Vec<bool>> {
        // Simplified implementation - would normally control optical tweezers
        let success_rate = self.config.loading_efficiency.unwrap_or(0.95);
        let loading_results = (0..self.config.atom_count)
            .map(|_| thread_rng().random::<f64>() < success_rate)
            .collect();
        Ok(loading_results)
    }

    async fn move_atoms(
        &self,
        _atom_indices: &[usize],
        _new_positions: &[(f64, f64, f64)],
    ) -> DeviceResult<()> {
        // Simplified implementation - would normally control optical tweezers
        if !self.supports_tweezer_manipulation().await? {
            return Err(DeviceError::UnsupportedOperation(
                "Tweezer manipulation not supported by this system".to_string(),
            ));
        }
        Ok(())
    }

    async fn rydberg_excitation(
        &self,
        atom_indices: &[usize],
        _excitation_time: Duration,
        _laser_power: f64,
    ) -> DeviceResult<Vec<bool>> {
        // Simplified implementation - would normally control Rydberg lasers
        if !self.supports_rydberg_gates().await? {
            return Err(DeviceError::UnsupportedOperation(
                "Rydberg excitation not supported by this system".to_string(),
            ));
        }

        let success_rate = 0.99;
        let excitation_results = atom_indices
            .iter()
            .map(|_| thread_rng().random::<f64>() < success_rate)
            .collect();
        Ok(excitation_results)
    }

    async fn global_rydberg_operation(
        &self,
        _operation: &str,
        _parameters: &HashMap<String, f64>,
    ) -> DeviceResult<()> {
        // Simplified implementation - would normally perform global Rydberg operations
        if !self.supports_rydberg_gates().await? {
            return Err(DeviceError::UnsupportedOperation(
                "Global Rydberg operations not supported by this system".to_string(),
            ));
        }
        Ok(())
    }

    async fn measure_atom_states(&self, atom_indices: &[usize]) -> DeviceResult<Vec<String>> {
        // Simplified implementation - would normally perform state detection
        let states = atom_indices
            .iter()
            .map(|_| {
                if thread_rng().random::<f64>() < 0.5 {
                    "ground".to_string()
                } else {
                    "excited".to_string()
                }
            })
            .collect();
        Ok(states)
    }

    async fn calculate_atom_correlations(
        &self,
        atom_pairs: &[(usize, usize)],
        _correlation_type: &str,
    ) -> DeviceResult<HashMap<String, f64>> {
        // Simplified implementation - would normally calculate quantum correlations
        let mut correlations = HashMap::new();
        for (i, (atom1, atom2)) in atom_pairs.iter().enumerate() {
            let correlation_key = format!("{atom1}_{atom2}");
            correlations.insert(
                correlation_key,
                thread_rng().random::<f64>().mul_add(2.0, -1.0),
            );
        }
        Ok(correlations)
    }

    async fn estimate_fidelity(
        &self,
        _target_state: &str,
        _measurement_data: &NeutralAtomMeasurementData,
    ) -> DeviceResult<f64> {
        // Simplified implementation - would normally estimate state fidelity
        Ok(self.config.gate_fidelity.unwrap_or(0.995))
    }
}

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

    fn create_test_device() -> NeutralAtomDevice {
        let client = NeutralAtomClient::new(
            "https://test-neutral-atom-api.example.com".to_string(),
            "test-token".to_string(),
        )
        .expect("Failed to create neutral atom client");
        let config = NeutralAtomDeviceConfig::default();
        NeutralAtomDevice::new(client, "test-device-1".to_string(), config)
    }

    #[tokio::test]
    async fn test_device_creation() {
        let device = create_test_device();
        assert_eq!(device.device_id, "test-device-1");
        assert_eq!(
            device.client.base_url,
            "https://test-neutral-atom-api.example.com"
        );
    }

    #[tokio::test]
    async fn test_device_properties() {
        let device = create_test_device();
        let properties = device.properties().await.expect("Failed to get properties");

        assert_eq!(
            properties.get("device_id").expect("Missing device_id"),
            "test-device-1"
        );
        assert_eq!(
            properties.get("system_type").expect("Missing system_type"),
            "Rydberg"
        );
        assert_eq!(
            properties.get("atom_count").expect("Missing atom_count"),
            "100"
        );
    }

    #[tokio::test]
    async fn test_quantum_device_traits() {
        let device = create_test_device();

        assert!(device
            .is_available()
            .await
            .expect("Failed to check availability"));
        assert_eq!(
            device
                .qubit_count()
                .await
                .expect("Failed to get qubit count"),
            100
        );
        assert!(!device
            .is_simulator()
            .await
            .expect("Failed to check is_simulator"));
    }

    #[tokio::test]
    async fn test_neutral_atom_capabilities() {
        let device = create_test_device();

        assert_eq!(
            device
                .system_type()
                .await
                .expect("Failed to get system type"),
            NeutralAtomSystemType::Rydberg
        );
        assert_eq!(
            device.atom_count().await.expect("Failed to get atom count"),
            100
        );
        assert_eq!(
            device
                .atom_spacing()
                .await
                .expect("Failed to get atom spacing"),
            5.0
        );
        assert_eq!(
            device
                .state_encoding()
                .await
                .expect("Failed to get state encoding"),
            AtomStateEncoding::GroundExcited
        );
        assert!(device
            .supports_rydberg_gates()
            .await
            .expect("Failed to check Rydberg gates support"));
        assert!(!device
            .supports_tweezer_manipulation()
            .await
            .expect("Failed to check tweezer manipulation support"));
    }

    #[tokio::test]
    async fn test_atom_operations() {
        let device = create_test_device();

        // Test atom loading
        let positions = vec![(0.0, 0.0, 0.0), (5.0, 0.0, 0.0), (10.0, 0.0, 0.0)];
        let loading_results = device
            .load_atoms(&positions)
            .await
            .expect("Failed to load atoms");
        assert_eq!(loading_results.len(), 100); // Should match atom_count

        // Test Rydberg excitation
        let atom_indices = vec![0, 1, 2];
        let excitation_results = device
            .rydberg_excitation(&atom_indices, Duration::from_nanos(1000), 10.0)
            .await
            .expect("Failed to perform Rydberg excitation");
        assert_eq!(excitation_results.len(), 3);

        // Test state measurement
        let states = device
            .measure_atom_states(&atom_indices)
            .await
            .expect("Failed to measure atom states");
        assert_eq!(states.len(), 3);
        assert!(states.iter().all(|s| s == "ground" || s == "excited"));
    }
}