quantum-sim 0.0.1

use ndarray::{Array2, arr2};
use num_complex::Complex;
use num_traits::{Float, FromPrimitive, One, Zero};

/// Represents a quantum gate with a name and unitary matrix representation.
///
/// # Type Parameters
/// - `T`: Floating-point type for the matrix elements to generate `Complex<T>`
///
/// # Fields
/// - `name`: String identifier for the gate
/// - `matrix`: Unitary matrix representing the gate operation (2D array of complex numbers)
///
/// # Examples
/// ```
/// use quantum_sim::gates::Gate;
/// use num_complex::Complex;
/// use num_traits::{One, Zero};
/// use ndarray::arr2;
///
/// let gate = Gate::<f64>::new("Custom".to_string(), arr2(&[[Complex::one(), Complex::zero()],
///                                                  [Complex::zero(), Complex::one()]]));
/// ```
#[derive(Debug, Clone)]
pub struct Gate<T> {
    pub name: String,
    pub matrix: Array2<Complex<T>>,
}

impl<T> Gate<T>
where
    T: Float + FromPrimitive + Copy + PartialOrd + FromPrimitive + std::fmt::LowerExp + 'static,
{
    /// Creates a new quantum gate with the given name and matrix.
    ///
    /// # Arguments
    /// * `name` - String identifier for the gate
    /// * `matrix` - Unitary matrix representing the gate operation
    ///
    /// # Panics
    /// The caller must ensure the matrix is unitary. This is checked internally.
    ///
    /// # Examples
    /// ```
    /// use quantum_sim::gates::Gate;
    /// use ndarray::arr2;
    /// use num_complex::Complex;
    /// use num_traits::{One, Zero};
    ///
    /// let matrix = arr2(&[[Complex::one(), Complex::zero()],
    ///                     [Complex::zero(), Complex::one()]]);
    /// let identity = Gate::<f64>::new("Identity".to_string(), matrix);
    /// ```
    pub fn new(name: String, matrix: Array2<Complex<T>>) -> Result<Self, String> {
        let rows = matrix.shape()[0];
        let cols = matrix.shape()[1];

        // 1. Check if the matrix is square
        if rows != cols {
            return Err(format!(
                "Matrix for gate '{}' must be square, but has dimensions {}x{}.",
                name, rows, cols
            ));
        }
        // 2. Check if the matrix is unitary
        if !is_unitary(&matrix) {
            Err(format!(
                "Matrix for gate '{}' is not unitary. Max difference from identity:",
                name
            ))
        } else {
            Ok(Self { name, matrix })
        }
    }

    /// Creates an Identity gate (I).
    ///
    /// # Examples
    /// ```
    /// use quantum_sim::gates::Gate;
    ///
    /// let i_gate = Gate::<f64>::i();
    /// ```
    pub fn i() -> Self {
        Self::new(
            "I".to_string(),
            arr2(&[
                [Complex::one(), Complex::zero()],
                [Complex::zero(), Complex::one()],
            ]),
        )
        .unwrap()
    }

    /// Creates a Pauli-X gate (NOT gate).
    /// # Examples
    /// ```
    /// use quantum_sim::gates::Gate;
    ///
    /// let x_gate = Gate::<f64>::x();
    /// assert_eq!(x_gate.name, "X");
    /// ```
    pub fn x() -> Self {
        Self::new(
            "X".to_string(),
            arr2(&[
                [Complex::zero(), Complex::one()],
                [Complex::one(), Complex::zero()],
            ]),
        )
        .unwrap()
    }

    /// Creates a Pauli-Y gate.
    /// # Examples
    /// ```
    /// use quantum_sim::gates::Gate;
    ///
    /// let y_gate = Gate::<f64>::y();
    /// assert_eq!(y_gate.name, "Y");
    /// ```
    pub fn y() -> Self {
        Self::new(
            "Y".to_string(),
            arr2(&[
                [Complex::zero(), -Complex::i()],
                [Complex::i(), Complex::zero()],
            ]),
        )
        .unwrap()
    }

    /// Creates a Pauli-Z gate.
    /// # Examples
    /// ```
    /// use quantum_sim::gates::Gate;
    ///
    /// let z_gate = Gate::<f64>::z();
    /// assert_eq!(z_gate.name, "Z");
    /// ```
    pub fn z() -> Self {
        Self::new(
            "Z".to_string(),
            arr2(&[
                [Complex::one(), Complex::zero()],
                [Complex::zero(), -Complex::one()],
            ]),
        )
        .unwrap()
    }

    /// Creates a Hadamard gate (H).
    ///
    /// The Hadamard gate creates superposition states.
    /// Matrix representation:
    /// ```text
    /// [1/sqrt(2)  1/sqrt(2)]
    /// [1/sqrt(2) -1/sqrt(2)]
    /// ```
    ///
    /// # Examples
    /// ```
    /// use quantum_sim::gates::Gate;
    ///
    /// let h_gate = Gate::<f64>::h();
    /// assert_eq!(h_gate.name, "H");
    /// ```
    pub fn h() -> Self {
        let factor = T::one() / T::from(2.0).unwrap().sqrt();
        Self::new(
            "H".to_string(),
            arr2(&[
                [Complex::one() * factor, Complex::one() * factor],
                [Complex::one() * factor, -Complex::one() * factor],
            ]),
        )
        .unwrap()
    }

    /// Creates a Phase gate (S gate).
    ///
    /// The Phase gate introduces a π/2 phase shift.
    /// Matrix representation:
    /// ```text
    /// [1 0]
    /// [0 i]
    /// ```
    ///
    /// # Examples
    /// ```
    /// use quantum_sim::gates::Gate;
    ///
    /// let s_gate = Gate::<f64>::s();
    /// assert_eq!(s_gate.name, "S");
    /// ```
    pub fn s() -> Self {
        Self::new(
            "S".to_string(),
            arr2(&[
                [Complex::one(), Complex::zero()],
                [Complex::zero(), Complex::i()],
            ]),
        )
        .unwrap()
    }

    /// Creates a T gate (π/8 gate).
    ///
    /// The T gate introduces a π/4 phase shift.
    /// Matrix representation:
    /// ```text
    /// [1 0]
    /// [0 e^(iπ/4)]
    /// ```
    ///
    /// # Examples
    /// ```
    /// use quantum_sim::gates::Gate;
    ///
    /// let t_gate = Gate::<f64>::t();
    /// assert_eq!(t_gate.name, "T");
    /// ```
    pub fn t() -> Self {
        let pi = T::from(std::f64::consts::PI).unwrap();
        let angle = pi / T::from(4.0).unwrap();
        Self::new(
            "T".to_string(),
            arr2(&[
                [Complex::one(), Complex::zero()],
                [Complex::zero(), Complex::new(angle.cos(), angle.sin())],
            ]),
        )
        .unwrap()
    }

    /// Creates a Controlled-NOT gate (CNOT).
    ///
    /// The CNOT gate flips the target qubit if the control qubit is |1⟩.
    /// Matrix representation:
    /// ```text
    /// [1 0 0 0]
    /// [0 1 0 0]
    /// [0 0 0 1]
    /// [0 0 1 0]
    /// ```
    ///
    /// # Examples
    /// ```
    /// use quantum_sim::gates::Gate;
    ///
    /// let cnot_gate = Gate::<f64>::cnot();
    /// assert_eq!(cnot_gate.name, "CNOT");
    /// ```
    pub fn cnot() -> Self {
        Self::new(
            "CNOT".to_string(),
            arr2(&[
                [
                    Complex::one(),
                    Complex::zero(),
                    Complex::zero(),
                    Complex::zero(),
                ],
                [
                    Complex::zero(),
                    Complex::one(),
                    Complex::zero(),
                    Complex::zero(),
                ],
                [
                    Complex::zero(),
                    Complex::zero(),
                    Complex::zero(),
                    Complex::one(),
                ],
                [
                    Complex::zero(),
                    Complex::zero(),
                    Complex::one(),
                    Complex::zero(),
                ],
            ]),
        )
        .unwrap()
    }
}

fn is_unitary<T>(matrix: &Array2<Complex<T>>) -> bool
where
    T: Float + 'static,
{
    let product = matrix.dot(&matrix.t().mapv(|c| c.conj()));
    let identity = Array2::<Complex<T>>::eye(matrix.shape()[0]);
    let diff = &product - &identity;
    let max_diff_norm = diff.iter().map(|c| c.norm()).fold(T::zero(), T::max);
    let epsilon = T::from(1e-6).unwrap(); // Use a small epsilon for floating point comparison
    max_diff_norm <= epsilon
}