quantrs2-circuit 0.2.1

//! Enhanced Quantum Pulse Control with Advanced `SciRS2` Signal Processing
//!
//! This module provides state-of-the-art pulse-level control for quantum devices
//! with ML-based pulse optimization, real-time calibration, advanced waveform
//! synthesis, and comprehensive error mitigation powered by `SciRS2`.

pub mod config;
pub mod pulses;

#[cfg(test)]
mod tests;

// Re-export main types
pub use config::*;
pub use pulses::*;

use quantrs2_core::error::{QuantRS2Error, QuantRS2Result};
use scirs2_core::ndarray::Array1;
use scirs2_core::Complex64;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};
use std::fmt;
use std::sync::Arc;

/// Enhanced pulse controller
pub struct EnhancedPulseController {
    config: EnhancedPulseConfig,
    signal_processor: SignalProcessor,
    pub ml_optimizer: Option<Arc<dyn PulseOptimizationModel>>,
    calibration_data: CalibrationData,
}

impl EnhancedPulseController {
    /// Create a new enhanced pulse controller
    #[must_use]
    pub fn new(config: EnhancedPulseConfig) -> Self {
        Self {
            config,
            signal_processor: SignalProcessor::new(),
            ml_optimizer: Some(Arc::new(DefaultPulseOptimizer::new())),
            calibration_data: CalibrationData::default(),
        }
    }
}

/// Signal processor for pulse waveforms
pub struct SignalProcessor {
    pub config: SignalProcessorConfig,
    buffer_manager: PulseSignalBufferManager,
    fft_engine: FFTEngine,
    filter_bank: FilterBank,
    adaptive_processor: AdaptiveSignalProcessor,
}

impl SignalProcessor {
    #[must_use]
    pub fn new() -> Self {
        Self {
            config: SignalProcessorConfig::default(),
            buffer_manager: PulseSignalBufferManager::new(),
            fft_engine: FFTEngine::new(),
            filter_bank: FilterBank::new(),
            adaptive_processor: AdaptiveSignalProcessor::new(),
        }
    }
}

impl Default for SignalProcessor {
    fn default() -> Self {
        Self::new()
    }
}

/// Signal processor configuration
#[derive(Debug, Clone)]
pub struct SignalProcessorConfig {
    pub window_size: usize,
    pub overlap: usize,
    pub enable_simd: bool,
    pub max_frequency: f64,
}

impl Default for SignalProcessorConfig {
    fn default() -> Self {
        Self {
            window_size: 1024,
            overlap: 512,
            enable_simd: true,
            max_frequency: 500e6,
        }
    }
}

/// Buffer manager for signal processing
struct PulseSignalBufferManager {
    complex_buffers: Vec<Vec<Complex64>>,
    real_buffers: Vec<Vec<f64>>,
    fft_workspace: Vec<Complex64>,
    filter_states: HashMap<String, FilterState>,
}

impl PulseSignalBufferManager {
    fn new() -> Self {
        Self {
            complex_buffers: Vec::new(),
            real_buffers: Vec::new(),
            fft_workspace: Vec::new(),
            filter_states: HashMap::new(),
        }
    }
}

/// Filter state for signal processing
pub struct FilterState {
    pub delay_line: VecDeque<f64>,
    pub coefficients: Vec<f64>,
    pub history: Vec<f64>,
}

impl FilterState {
    pub fn new(order: usize) -> Self {
        Self {
            delay_line: VecDeque::with_capacity(order),
            coefficients: Vec::new(),
            history: Vec::with_capacity(order),
        }
    }
}

/// FFT engine for spectral analysis
struct FFTEngine {
    fft_plans: HashMap<usize, FFTPlan>,
    buffer_pool: Vec<Vec<Complex64>>,
}

impl FFTEngine {
    fn new() -> Self {
        Self {
            fft_plans: HashMap::new(),
            buffer_pool: Vec::new(),
        }
    }
}

#[derive(Debug, Clone)]
struct FFTPlan {
    size: usize,
    direction: FFTDirection,
}

#[derive(Debug, Clone, Copy)]
enum FFTDirection {
    Forward,
    Inverse,
}

/// Filter bank for different filter types
struct FilterBank {
    butterworth_filters: HashMap<usize, ButterworthFilter>,
    chebyshev_filters: HashMap<usize, ChebyshevFilter>,
    fir_filters: HashMap<usize, FIRFilter>,
    adaptive_filters: Vec<AdaptiveFilter>,
}

impl FilterBank {
    fn new() -> Self {
        Self {
            butterworth_filters: HashMap::new(),
            chebyshev_filters: HashMap::new(),
            fir_filters: HashMap::new(),
            adaptive_filters: Vec::new(),
        }
    }
}

struct ButterworthFilter {
    order: usize,
    cutoff: f64,
}

struct ChebyshevFilter {
    order: usize,
    ripple: f64,
}

struct FIRFilter {
    taps: Vec<f64>,
}

struct AdaptiveFilter {
    weights: Vec<f64>,
    step_size: f64,
}

/// Adaptive signal processor
struct AdaptiveSignalProcessor {
    noise_estimator: NoiseEstimator,
    distortion_corrector: DistortionCorrector,
    interference_canceller: InterferenceCanceller,
    channel_equalizer: ChannelEqualizer,
}

impl AdaptiveSignalProcessor {
    fn new() -> Self {
        Self {
            noise_estimator: NoiseEstimator {
                noise_floor: -80.0,
                noise_profile: Array1::zeros(1024),
                estimation_window: 1024,
                update_rate: 0.01,
            },
            distortion_corrector: DistortionCorrector {
                correction_model: PredistortionModel::Linear,
                model_parameters: vec![1.0, 0.0],
                adaptation_enabled: true,
                correction_strength: 1.0,
            },
            interference_canceller: InterferenceCanceller {
                reference_signals: Vec::new(),
                cancellation_filters: Vec::new(),
                threshold: 0.1,
            },
            channel_equalizer: ChannelEqualizer {
                frequency_response: Array1::ones(1024),
                target_response: Array1::ones(1024),
                equalization_filter: vec![1.0],
                adaptation_rate: 0.01,
            },
        }
    }
}

struct NoiseEstimator {
    noise_floor: f64,
    noise_profile: Array1<f64>,
    estimation_window: usize,
    update_rate: f64,
}

struct DistortionCorrector {
    correction_model: PredistortionModel,
    model_parameters: Vec<f64>,
    adaptation_enabled: bool,
    correction_strength: f64,
}

struct InterferenceCanceller {
    reference_signals: Vec<Array1<Complex64>>,
    cancellation_filters: Vec<Vec<f64>>,
    threshold: f64,
}

struct ChannelEqualizer {
    frequency_response: Array1<f64>,
    target_response: Array1<f64>,
    equalization_filter: Vec<f64>,
    adaptation_rate: f64,
}

/// Predistortion models
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PredistortionModel {
    Linear,
    Polynomial,
    MemoryPolynomial,
}

/// Calibration data
#[derive(Debug, Clone, Default)]
struct CalibrationData {
    qubit_frequencies: HashMap<usize, f64>,
    anharmonicities: HashMap<usize, f64>,
    coupling_strengths: HashMap<(usize, usize), f64>,
}

/// Pulse sequence
#[derive(Debug, Clone)]
pub struct PulseSequence {
    pub channels: Vec<PulseChannel>,
    pub duration: f64,
    pub metadata: PulseMetadata,
}

/// Pulse channel
#[derive(Debug, Clone)]
pub struct PulseChannel {
    pub channel_id: usize,
    pub waveform: Waveform,
    pub frequency: f64,
    pub phase: f64,
    pub frame_change: Option<f64>,
}

/// Waveform data
#[derive(Debug, Clone)]
pub struct Waveform {
    pub samples: Vec<Complex64>,
    pub sample_rate: f64,
}

/// Pulse metadata
#[derive(Debug, Clone)]
pub struct PulseMetadata {
    pub gate_name: String,
    pub target_qubits: Vec<usize>,
    pub fidelity_estimate: Option<f64>,
    pub optimization_history: Vec<OptimizationStep>,
}

/// Optimization step
#[derive(Debug, Clone)]
pub struct OptimizationStep {
    pub iteration: usize,
    pub cost: f64,
    pub parameters: Vec<f64>,
}

/// Gate analysis
#[derive(Debug, Clone)]
pub struct GateAnalysis {
    pub target_unitary: Vec<Vec<Complex64>>,
    pub qubit_indices: Vec<usize>,
}

/// Pulse optimization model trait
pub trait PulseOptimizationModel: Send + Sync {
    fn optimize(
        &self,
        pulse: &PulseSequence,
        target: &GateAnalysis,
        constraints: &PulseConstraints,
    ) -> QuantRS2Result<PulseSequence>;

    fn update(&mut self, feedback: &OptimizationFeedback);
}

/// Default pulse optimizer
struct DefaultPulseOptimizer {
    // ML model placeholder
}

impl DefaultPulseOptimizer {
    const fn new() -> Self {
        Self {}
    }
}

impl PulseOptimizationModel for DefaultPulseOptimizer {
    /// Real (non-ML, closed-form) default pulse optimizer.
    ///
    /// For a single-qubit target gate, decomposes `target.target_unitary`
    /// (up to an unobservable global phase) into Z-X-Z Euler angles
    /// `Rz(alpha) * Rx(theta) * Rz(beta)`, then applies the standard
    /// virtual-Z identity `Rz(alpha) * Rx(theta) = X_alpha(theta) * Rz(alpha)`
    /// (a Z rotation commutes through a resonant drive by shifting its
    /// phase) to rewrite the target as `X_alpha(theta) * Rz(alpha + beta)`:
    /// a single drive at phase `alpha` realizing rotation angle `theta`,
    /// followed by a trailing virtual-Z frame correction of `alpha + beta`.
    /// The channel matching the target qubit has its envelope rescaled (in
    /// both magnitude and phase) so its integrated area exactly equals
    /// `theta * e^{i alpha}`, and its `frame_change` updated to carry the
    /// trailing virtual-Z — both real, physically meaningful pulse-level
    /// operations, not a no-op passthrough.
    ///
    /// Multi-qubit targets are out of scope for this closed-form model (a
    /// genuine multi-qubit GRAPE-style optimizer would need a full
    /// multi-qubit Hamiltonian/time-evolution simulation); for those the
    /// pulse is honestly returned unchanged rather than fabricating a
    /// multi-qubit solution.
    fn optimize(
        &self,
        pulse: &PulseSequence,
        target: &GateAnalysis,
        constraints: &PulseConstraints,
    ) -> QuantRS2Result<PulseSequence> {
        let is_single_qubit_target = target.qubit_indices.len() == 1
            && target.target_unitary.len() == 2
            && target.target_unitary.iter().all(|row| row.len() == 2);

        if !is_single_qubit_target {
            return Ok(pulse.clone());
        }

        let (theta, alpha, beta) = zxz_euler_angles(&target.target_unitary)?;
        let target_qubit = target.qubit_indices[0];

        let mut optimized = pulse.clone();
        for channel in &mut optimized.channels {
            if channel.channel_id == target_qubit {
                reshape_channel_for_rotation(channel, theta, alpha, beta, constraints);
            }
        }

        Ok(optimized)
    }

    /// This default optimizer is a deterministic closed-form solver (not a
    /// trainable/adaptive model), so there is no learned state for
    /// measurement feedback to update; adaptive re-optimization from
    /// [`OptimizationFeedback`] is left to real ML-backed
    /// [`PulseOptimizationModel`] implementations.
    fn update(&mut self, _feedback: &OptimizationFeedback) {}
}

/// Decompose a 2x2 unitary (up to global phase) into Z-X-Z Euler angles
/// `(theta, alpha, beta)` such that `U == e^{i*phase} * Rz(alpha) * Rx(theta) * Rz(beta)`,
/// where `Rz(phi) = diag(e^{-i phi/2}, e^{i phi/2})` and
/// `Rx(theta) = [[cos(theta/2), -i sin(theta/2)], [-i sin(theta/2), cos(theta/2)]]`.
fn zxz_euler_angles(target_unitary: &[Vec<Complex64>]) -> QuantRS2Result<(f64, f64, f64)> {
    if target_unitary.len() != 2 || target_unitary.iter().any(|row| row.len() != 2) {
        return Err(QuantRS2Error::InvalidInput(
            "zxz_euler_angles requires a 2x2 unitary matrix".to_string(),
        ));
    }

    let u00 = target_unitary[0][0];
    let u01 = target_unitary[0][1];
    let u10 = target_unitary[1][0];
    let u11 = target_unitary[1][1];

    // Normalize to SU(2) by dividing out the global phase / determinant.
    let determinant = u00 * u11 - u01 * u10;
    if determinant.norm_sqr() < 1e-24 {
        return Err(QuantRS2Error::InvalidInput(
            "target unitary is singular; cannot extract Euler angles".to_string(),
        ));
    }
    let det_sqrt = determinant.sqrt();
    let v00 = u00 / det_sqrt;
    let v10 = u10 / det_sqrt;

    let theta = 2.0 * v10.norm().atan2(v00.norm());

    // (alpha+beta)/2 = -arg(v00); undetermined (gauge freedom) when v00 ~ 0.
    let half_sum = if v00.norm() > 1e-9 { -v00.arg() } else { 0.0 };
    // (alpha-beta)/2 = arg(v10) + pi/2; undetermined when v10 ~ 0.
    let half_diff = if v10.norm() > 1e-9 {
        v10.arg() + std::f64::consts::FRAC_PI_2
    } else {
        0.0
    };

    let alpha = half_sum + half_diff;
    let beta = half_sum - half_diff;

    Ok((theta, alpha, beta))
}

/// Upper bound on how many samples [`reshape_channel_for_rotation`] will
/// synthesize when extending a pulse's duration to respect
/// [`PulseConstraints::max_amplitude`]. This keeps the (physically
/// legitimate) duration-extension strategy from ever attempting to
/// allocate an unbounded number of samples for pathological inputs (e.g. a
/// target rotation combined with an extremely low amplitude ceiling and an
/// extremely high sample rate, which would imply a pulse lasting many
/// physical seconds at nanosecond sampling).
const MAX_PULSE_EXTENSION_SAMPLES: usize = 1_000_000;

/// Rescale a pulse channel's envelope so its integrated (complex) area
/// equals `theta * e^{i alpha}` — i.e. a resonant drive at phase `alpha`
/// realizing rotation angle `theta` — and set its trailing `frame_change`
/// to carry the virtual-Z correction `alpha + beta`.
///
/// If the configured [`PulseConstraints::max_amplitude`] would be
/// exceeded, naively clamping every sample's amplitude down to the ceiling
/// would silently shrink the pulse's integrated area and therefore fail to
/// realize the intended rotation at all — a physically wrong result
/// dressed up as success. Instead, the pulse duration is extended (more,
/// proportionally smaller-amplitude samples at the same sample rate) so
/// the *exact* target area can still be reached without exceeding the
/// amplitude ceiling, exactly as a real hardware pulse generator would
/// trade duration for peak power. Only when the required extension would
/// need an unreasonable number of samples ([`MAX_PULSE_EXTENSION_SAMPLES`])
/// does this fall back to a bounded-duration clamp, honestly accepting a
/// reduced rotation angle rather than allocating unbounded memory.
fn reshape_channel_for_rotation(
    channel: &mut PulseChannel,
    theta: f64,
    alpha: f64,
    beta: f64,
    constraints: &PulseConstraints,
) {
    let dt = if channel.waveform.sample_rate > 0.0 {
        1.0 / channel.waveform.sample_rate
    } else {
        1.0
    };
    let target_area = Complex64::from_polar(theta, alpha);

    if channel.waveform.samples.is_empty() {
        channel.waveform.samples = vec![Complex64::new(1.0, 0.0)];
    }
    let sample_count = channel.waveform.samples.len();

    let current_area: Complex64 =
        channel.waveform.samples.iter().copied().sum::<Complex64>() * Complex64::new(dt, 0.0);

    if current_area.norm() > 1e-12 {
        let correction = target_area / current_area;
        for sample in &mut channel.waveform.samples {
            *sample *= correction;
        }
    } else {
        // Degenerate (zero-area) envelope: fall back to a flat-top pulse of
        // the same length carrying the exact required area.
        let flat_amplitude = target_area / Complex64::new(sample_count as f64 * dt, 0.0);
        channel.waveform.samples = vec![flat_amplitude; sample_count];
    }

    if let Some(max_amplitude) = constraints.max_amplitude {
        let peak = channel
            .waveform
            .samples
            .iter()
            .map(|sample| sample.norm())
            .fold(0.0_f64, f64::max);
        if peak > max_amplitude && peak > 0.0 && max_amplitude > 0.0 && dt > 0.0 {
            // Minimum sample count for a flat-top pulse at exactly
            // `max_amplitude` to still integrate to the full target area:
            // n * max_amplitude * dt >= |target_area|.
            let min_samples_for_target_area = (target_area.norm() / (max_amplitude * dt)).ceil();

            let extended_len = if min_samples_for_target_area.is_finite()
                && min_samples_for_target_area <= MAX_PULSE_EXTENSION_SAMPLES as f64
            {
                sample_count
                    .max(min_samples_for_target_area as usize)
                    .max(1)
            } else {
                usize::MAX
            };

            if extended_len <= MAX_PULSE_EXTENSION_SAMPLES {
                // Physically realizable within a reasonable duration:
                // synthesize an exact flat-top pulse of the required
                // length, hitting the target area exactly while never
                // exceeding the amplitude ceiling.
                let flat_amplitude = target_area / Complex64::new(extended_len as f64 * dt, 0.0);
                channel.waveform.samples = vec![flat_amplitude; extended_len];
            } else {
                // The required duration extension is unreasonably large
                // (e.g. an extremely tight amplitude ceiling paired with an
                // extremely high sample rate). Rather than allocate an
                // unbounded number of samples, honestly clamp the peak
                // amplitude at the existing duration: the requested
                // rotation angle will not be fully achieved, which is a
                // genuine hardware limitation, not a fabricated success.
                let clamp_scale = max_amplitude / peak;
                for sample in &mut channel.waveform.samples {
                    *sample *= clamp_scale;
                }
            }
        }
    }

    channel.frame_change = Some(channel.frame_change.unwrap_or(0.0) + alpha + beta);
}

/// Optimization feedback
#[derive(Debug, Clone)]
pub struct OptimizationFeedback {
    pub measured_fidelity: f64,
    pub execution_time: f64,
    pub success: bool,
}

/// Mitigation strategies
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MitigationStrategy {
    PhaseCorrection,
    AmplitudeStabilization,
    DriftCompensation,
    LeakageReduction,
    CrosstalkCancellation,
}

impl fmt::Display for PulseSequence {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "Pulse Sequence:")?;
        writeln!(f, "  Duration: {:.2} ns", self.duration * 1e9)?;
        writeln!(f, "  Channels: {}", self.channels.len())?;
        for channel in &self.channels {
            writeln!(
                f,
                "    Channel {}: {} samples @ {:.1} GHz",
                channel.channel_id,
                channel.waveform.samples.len(),
                channel.frequency / 1e9
            )?;
        }
        writeln!(f, "  Gate: {}", self.metadata.gate_name)?;
        if let Some(fidelity) = self.metadata.fidelity_estimate {
            writeln!(f, "  Estimated fidelity: {fidelity:.4}")?;
        }
        Ok(())
    }
}