quantum-sim 0.0.1

//! # Random Operator Generator
//!
//! This module provides tools to generate random **unitary matrices**,
//! **density matrices**, and **state vectors**....
//! It supports different distributions (Haar, Bures) and both real and complex
//! cases. The API follows a *builder pattern* for flexible, keyword-like usage.
//!
//! # Example
//! ```
//! use quantum_sim::random_gate::{RandomDensityMat, RandomUnitary};
//! use num_complex::Complex;
//!
//! let u = RandomUnitary::<f64>::new(4).is_real(false).build();
//! let rho = RandomDensityMat::<f64>::new(4).build();
//! let rho_bures = RandomDensityMat::<f64>::new(4).bures().is_real(true).build();
//! ```
//!
//! ## Features
//! - Haar and Bures random state generation
//! - Real or complex support
//! - Generic over f32/f64
//! - Type-safe and builder-based interface

use ndarray::{Array1, Array2};
use ndarray_linalg::{Lapack, QR};
use ndarray_rand::{
    RandomExt,
    rand::distributions::Distribution,
    rand::thread_rng,
    rand_distr::{Normal, StandardNormal},
};
use num_complex::Complex;
use num_traits::Float;

/// Internal: generate a random Ginibre matrix.
///
/// # Parameters
/// - `m`: number of columns
/// - `n`: number of rows
/// - `is_real`: whether to generate real entries
///
/// # Returns
/// Random matrix sampled from real or complex Gaussian distribution.
fn ginibre<T>(m: usize, n: usize, is_real: bool) -> Array2<Complex<T>>
where
    T: Float + ndarray_linalg::Scalar,
    StandardNormal: Distribution<T>,
{
    let mut rng = thread_rng();
    let normal = Normal::new(T::zero(), T::one()).unwrap();

    if is_real {
        Array2::random_using((m, n), normal, &mut rng).mapv(|x| Complex::new(x, T::zero()))
    } else {
        let re = Array2::random_using((m, n), normal, &mut rng);
        let im = Array2::random_using((m, n), normal, &mut rng);
        re.mapv(|x| Complex::new(x, T::zero())) + im.mapv(|y| Complex::new(T::zero(), y))
    }
}

/// A struct for generating random unitary matrices.
///
/// This struct creates a random matrix distributed according to the Haar measure,
/// with complex case and an orthogonal matrix in the real case.
///
/// # Example
/// ```rust
/// use quantum_sim::random_gate::RandomUnitary;
///
/// let ru = RandomUnitary::<f64>::new(4).build();
/// let ru1 = RandomUnitary::<f64>::new(4).is_real(true).build();
/// ```
pub struct RandomUnitary<T> {
    dim_col: usize,
    dim_row: Option<usize>,
    is_real: Option<bool>,
    _marker: std::marker::PhantomData<T>,
}

impl<T> RandomUnitary<T>
where
    T: Float + ndarray_linalg::Scalar,
    StandardNormal: Distribution<T>,
    Complex<T>: Lapack,
{
    /// Creates a new `RandomUnitary` builder.
    ///
    /// # Arguments
    ///
    /// * `dim_col` - The number of columns for the matrix. By default, the
    ///   matrix will be square (`dim_row` = `dim_col`).
    pub fn new(dim_col: usize) -> Self {
        Self {
            dim_col,
            dim_row: None,
            is_real: None,
            _marker: std::marker::PhantomData,
        }
    }

    /// Set row dimension (non-square matrices).
    pub fn dim(mut self, dim_row: usize) -> Self {
        self.dim_row = Some(dim_row);
        self
    }

    /// Set to real matrix.
    pub fn is_real(mut self, flag: bool) -> Self {
        self.is_real = Some(flag);
        self
    }

    /// Builds the random unitary (or orthogonal) matrix.
    ///
    /// # Returns
    ///
    /// An `Array2<Complex<T>>` representing the random matrix.
    pub fn build(self) -> Array2<Complex<T>> {
        let dim_row = self.dim_row.unwrap_or(self.dim_col);
        let is_real = self.is_real.unwrap_or(false);
        let gin: Array2<Complex<T>> = ginibre::<T>(self.dim_col, dim_row, is_real);
        let (mut q, r) = gin.qr().expect("QR decomposition failed");

        let r_diag_signs: Array1<T> = r.diag().mapv(|c| {
            let re = c.re.signum();
            if re == T::zero() { T::one() } else { re }
        });

        for (mut col, &sign) in q.columns_mut().into_iter().zip(&r_diag_signs) {
            col.mapv_inplace(|z| z * Complex::new(sign, T::zero()));
        }

        q
    }
}

/// A struct for generating random density matrices.
///
/// It can generate random density matrices from different statistical ensembles
/// (Haar or Bures), both real and complex matrices.
///
///  # Example
/// ```rust
/// use quantum_sim::random_gate::{RandomDensityMat, RandomUnitary};
///
/// let rho = RandomDensityMat::<f64>::new(4).build();
/// let rho_bures = RandomDensityMat::<f64>::new(4).bures().is_real(true).build();
/// ```

pub struct RandomDensityMat<T: Float + ndarray_linalg::Scalar> {
    dim_col: usize,
    dim_row: Option<usize>,
    is_real: Option<bool>,
    is_bures: bool,
    _marker: std::marker::PhantomData<T>,
}

impl<T> RandomDensityMat<T>
where
    T: Float + ndarray_linalg::Scalar,
    StandardNormal: Distribution<T>,
    Complex<T>: Lapack,
{
    pub fn new(dim_col: usize) -> Self {
        Self {
            dim_col,
            dim_row: None,
            is_real: None,
            is_bures: false,
            _marker: std::marker::PhantomData,
        }
    }

    /// Set row dimension (non-square).
    pub fn dim(mut self, dim_row: usize) -> Self {
        self.dim_row = Some(dim_row);
        self
    }

    /// Set to real density matrix.
    pub fn is_real(mut self, flag: bool) -> Self {
        self.is_real = Some(flag);
        self
    }

    /// Use Bures distribution instead of Haar.
    pub fn bures(mut self) -> Self {
        self.is_bures = true;
        self
    }

    /// Build and return a normalized density matrix.
    ///
    /// # Returns
    /// An `Array2<Complex<T>>` representing the random, trace-normalized density matrix.
    pub fn build(self) -> Array2<Complex<T>> {
        let dim_row = self.dim_row.unwrap_or(self.dim_col);
        let is_real = self.is_real.unwrap_or(false);

        let mut gin = ginibre::<T>(self.dim_col, dim_row, is_real);

        // Bures case: (U + I)G
        if self.is_bures {
            let u = RandomUnitary::<T>::new(self.dim_col)
                .is_real(is_real)
                .build();
            let eye = Array2::<Complex<T>>::eye(self.dim_col);
            gin = (u + eye).dot(&gin);
        }

        // ρ = GG† / Tr(GG†)
        let mut rho = gin.dot(&gin.t().mapv(|c| c.conj()));
        let trace = rho.diag().sum();
        rho.mapv_inplace(|x| x / trace);
        rho
    }
}