//! # Quantum Gates
//!
//! This module provides a set of quantum programming operations and utilities
//! for easy use in Rust.
//! 
//! Note that you need a Quantum Simulator to execute the code generated by this library.
//! [KBW](https://gitlab.com/quantum-ket/kbw) is the recommended open-source option. 
//!
//! ## Features
//!
//! This library provides the following features:
//!
//! - Creation and manipulation of quantum states through the `Quant` struct.
//! - Application of various quantum gates such as Pauli-X, Pauli-Y, Pauli-Z, Hadamard, and more.
//! - Control operations, adjoint operations, and composite operations.
//! - Measurement and state dumping of quantum states.
//!
//! ## Example
//!
//! Example demonstrating the usage of the module:
//!
//! ```rust
//! use ket::*;
//!
//! fn main() -> Result<()> {
//!     let p = Process::new_ptr();
//!     let q = Quant::new(&p, 2)?;
//!
//!     h(&q.at(0))?;
//!     ctrl(&q.at(0), || x(&q.at(1)))??;
//!
//!     let d = dump(&q)?;
//!
//!     p.borrow_mut().prepare_for_execution()?;
//!
//!     let result = "{\n\"future\": [\n0\n],\n\"dump\": [\n{\n\"Vector\": {\n\"basis_states\": [\n[\n0\n],\n[\n3\n]\n],\n\"amplitudes_real\": [\n0.7071067811865476,\n0.7071067811865476\n],\n\"amplitudes_imag\": [\n0.0,\n0.0\n]\n}\n}\n],\n\"exec_time\": 0.00103683\n}";
//!
//!     p.borrow_mut().set_result(serde_json::from_str(result).unwrap())?;
//!         
//!     println!("{:#?}", d);
//!
//!     Ok(())
//! }
//! ```

use std::{cell::RefCell, ops::RangeBounds, rc::Rc, result};

use crate::{error::KetError, Dump, Future, Process, Qubit};

#[derive(Debug)]

/// Represents a quantum state.
pub struct Quant {
    pub(crate) qubits: Vec<Rc<RefCell<Qubit>>>,
    pub(crate) process: Rc<RefCell<Process>>,
}

/// Alias for the result type used in the module.
pub type Result<T> = result::Result<T, KetError>;

impl Quant {
    /// Creates a new `Quant` object with the given process and number of qubits.
    ///
    /// # Arguments
    ///
    /// * `process` - A reference to the process associated with the quantum state.
    /// * `size` - The number of qubits to allocate for the quantum state.
    ///     
    /// # Returns
    ///
    /// A `Result` containing the created `Quant` object if successful, or a `KetError` if an error occurred.
    pub fn new(process: &Rc<RefCell<Process>>, size: usize) -> Result<Quant> {
        let mut qubits = Vec::with_capacity(size);
        for _ in 0..size {
            qubits.push(Rc::new(RefCell::new(
                process.borrow_mut().allocate_qubit(false)?,
            )));
        }

        Ok(Quant {
            qubits,
            process: Rc::clone(process),
        })
    }

    /// Returns a new `Quant` object that represents a slice of the original `Quant` object.
    ///
    /// # Arguments
    ///
    /// * `range` - The range of indices to include in the slice.
    ///
    /// # Returns
    ///
    /// A new `Quant` object that represents the specified slice of the current `Quant` object.
    pub fn slice<R: RangeBounds<usize>>(&self, range: R) -> Quant {
        let start = match range.start_bound() {
            std::ops::Bound::Included(start) => *start,
            std::ops::Bound::Excluded(_) => panic!(),
            std::ops::Bound::Unbounded => 0,
        };

        let end = match range.end_bound() {
            std::ops::Bound::Included(_) => panic!(),
            std::ops::Bound::Excluded(end) => *end,
            std::ops::Bound::Unbounded => self.qubits.len(),
        };

        Quant {
            process: Rc::clone(&self.process),
            qubits: Vec::from_iter(self.qubits[start..end].iter().map(Rc::clone)),
        }
    }

    /// Concatenates multiple `Quant` objects into a single `Quant` object.
    ///
    /// # Arguments
    ///
    /// * `quants` - An array of references to `Quant` objects to concatenate.
    ///
    /// # Returns
    ///
    /// A `Result` containing the concatenated `Quant` object if successful, or a `KetError` if an error occurred.
    pub fn cat(quants: &[&Quant]) -> Result<Quant> {
        let process = &quants
            .first()
            .expect("\"quants\" must have length >= 1")
            .process;

        let pid = process.borrow().pid();

        let mut qubits = Vec::new();
        for quant in quants {
            if quant.process.borrow().pid() != pid {
                return Err(KetError::PidMismatch);
            }
            qubits.extend(quant.qubits.iter().cloned());
        }

        Ok(Quant {
            qubits,
            process: Rc::clone(process),
        })
    }

    /// Returns a new `Quant` object that represents a single qubit at the given index.
    ///
    /// # Arguments
    ///
    /// * `index` - The index of the qubit to select.
    ///
    /// # Returns
    ///
    /// A new `Quant` object containing only the qubit at the specified index.
    pub fn at(&self, index: usize) -> Quant {
        Quant {
            process: Rc::clone(&self.process),
            qubits: vec![self.qubits[index].clone()],
        }
    }
}

/// Applies the Pauli-X gate to the given qubits.
///
/// # Arguments
///
/// * `qubits` - The quantum state on which to apply the gate.
///
/// # Returns
///
/// A `Result` indicating success or failure.
pub fn x(qubits: &Quant) -> Result<()> {
    for target in &qubits.qubits {
        qubits
            .process
            .borrow_mut()
            .apply_gate(crate::QuantumGate::PauliX, &target.borrow())?
    }

    Ok(())
}

/// Applies the Pauli-Y gate to the given qubits.
///
/// # Arguments
///
/// * `qubits` - The quantum state on which to apply the gate.
///
/// # Returns
///
/// A `Result` indicating success or failure.
pub fn y(qubits: &Quant) -> Result<()> {
    for target in &qubits.qubits {
        qubits
            .process
            .borrow_mut()
            .apply_gate(crate::QuantumGate::PauliY, &target.borrow())?
    }

    Ok(())
}

/// Applies the Pauli-Z gate to the given qubits.
///
/// # Arguments
///
/// * `qubits` - The quantum state on which to apply the gate.
///
/// # Returns
///
/// A `Result` indicating success or failure.
pub fn z(qubits: &Quant) -> Result<()> {
    for target in &qubits.qubits {
        qubits
            .process
            .borrow_mut()
            .apply_gate(crate::QuantumGate::PauliZ, &target.borrow())?
    }

    Ok(())
}

/// Applies the Hadamard gate to the given qubits.
///
/// # Arguments
///
/// * `qubits` - The quantum state on which to apply the gate.
///
/// # Returns
///
/// A `Result` indicating success or failure.
pub fn h(qubits: &Quant) -> Result<()> {
    for target in &qubits.qubits {
        qubits
            .process
            .borrow_mut()
            .apply_gate(crate::QuantumGate::Hadamard, &target.borrow())?
    }

    Ok(())
}

/// Applies the phase gate with the given lambda parameter to the given qubits.
///
/// # Arguments
///
/// * `lambda` - The lambda parameter of the phase gate.
/// * `qubits` - The quantum state on which to apply the gate.
///
/// # Returns
///
/// A `Result` indicating success or failure.
pub fn phase(lambda: f64, qubits: &Quant) -> Result<()> {
    for target in &qubits.qubits {
        qubits
            .process
            .borrow_mut()
            .apply_gate(crate::QuantumGate::Phase(lambda), &target.borrow())?
    }

    Ok(())
}

/// Applies the rotation gate around the X axis with the given theta parameter to the given qubits.
///
/// # Arguments
///
/// * `theta` - The theta parameter of the rotation gate.
/// * `qubits` - The quantum state on which to apply the gate.
///
/// # Returns
///
/// A `Result` indicating success or failure.
pub fn rx(theta: f64, qubits: &Quant) -> Result<()> {
    for target in &qubits.qubits {
        qubits
            .process
            .borrow_mut()
            .apply_gate(crate::QuantumGate::RX(theta), &target.borrow())?
    }

    Ok(())
}

/// Applies the rotation gate around the Y axis with the given theta parameter to the given qubits.
///
/// # Arguments
///
/// * `theta` - The theta parameter of the rotation gate.
/// * `qubits` - The quantum state on which to apply the gate.
///
/// # Returns
///
/// A `Result` indicating success or failure.
pub fn ry(theta: f64, qubits: &Quant) -> Result<()> {
    for target in &qubits.qubits {
        qubits
            .process
            .borrow_mut()
            .apply_gate(crate::QuantumGate::RY(theta), &target.borrow())?
    }

    Ok(())
}

/// Applies the rotation gate around the Z axis with the given theta parameter to the given qubits.
///
/// # Arguments
///
/// * `theta` - The theta parameter of the rotation gate.
/// * `qubits` - The quantum state on which to apply the gate.
///
/// # Returns
///
/// A `Result` indicating success or failure.
pub fn rz(theta: f64, qubits: &Quant) -> Result<()> {
    for target in &qubits.qubits {
        qubits
            .process
            .borrow_mut()
            .apply_gate(crate::QuantumGate::RZ(theta), &target.borrow())?
    }

    Ok(())
}

/// Applies a controlled operation on the given qubits.
///
/// The `control` qubits control the execution of the `gate` function.
///
/// # Arguments
///
/// * `control` - The control qubits for the controlled operation.
/// * `gate` - The gate function to be executed when the control qubits are in the desired state.
///
/// # Returns
///
/// A `Result` containing the result of the controlled operation if successful, or a `KetError` if an error occurred.
pub fn ctrl<F, T>(control: &Quant, gate: F) -> Result<T>
where
    F: FnOnce() -> T,
{
    control
        .process
        .borrow_mut()
        .ctrl_push_ref(&control.qubits)?;

    let result = gate();

    control.process.borrow_mut().ctrl_pop()?;

    Ok(result)
}

/// Applies an adjoint operation on the given process.
///
/// The `gate` function represents the operation to be executed.
///
/// # Arguments
///
/// * `process` - A reference to the process on which to apply the adjoint operation.
/// * `gate` - The gate function representing the operation to be executed.
///
/// # Returns
///
/// A `Result` containing the result of the adjoint operation if successful, or a `KetError` if an error occurred.
pub fn adj<F, T>(process: &Rc<RefCell<Process>>, gate: F) -> Result<T>
where
    F: FnOnce() -> T,
{
    process.borrow_mut().adj_begin()?;

    let result = gate();

    process.borrow_mut().adj_end()?;

    Ok(result)
}

/// Executes a function `inner` within a function `outer`.
///
/// This is typically used to apply an operation before and after the `inner` function.
///
/// # Arguments
///
/// * `process` - A reference to the process on which to apply the adjoint operation.
/// * `outer` - The outer function to be executed before and after (adj) the `inner` function.
/// * `inner` - The inner function to be executed within the context of the `outer` function.
///
/// # Returns
///
/// A `Result` containing the result of the inner operation if successful, or a `KetError` if an error occurred.
pub fn around<Outer, Inner, T>(
    process: &Rc<RefCell<Process>>,
    outer: Outer,
    inner: Inner,
) -> Result<T>
where
    Outer: Fn(),
    Inner: FnOnce() -> T,
{
    outer();

    let result = inner();

    adj(process, outer)?;

    Ok(result)
}

/// Measures the given qubits, collapsing them to a classical state.
///
/// # Arguments
///
/// * `qubits` - The quantum state to be measured.
///
/// # Returns
///
/// A `Result` containing the future object representing the measurement outcome if successful, or a `KetError` if an error occurred.
pub fn measure(qubits: &mut Quant) -> Result<Future> {
    qubits.process.borrow_mut().measure_ref(&mut qubits.qubits)
}

/// Applies a plugin operation on the given qubits with the specified plugin name and arguments.
///
/// # Arguments
///
/// * `name` - The name of the plugin to be applied.
/// * `args` - The arguments to be passed to the plugin.
/// * `qubits` - The quantum state on which to apply the plugin operation.
///
/// # Returns
///
/// A `Result` indicating success or failure.
pub fn plugin(name: &str, args: &str, qubits: &Quant) -> Result<()> {
    qubits
        .process
        .borrow_mut()
        .apply_plugin_ref(name, &qubits.qubits, args)
}

/// Dumps the quantum state of the given qubits.
///
/// # Arguments
///
/// * `qubits` - The quantum state to be dumped.
///
/// # Returns
///
/// A `Result` containing the dump object representing the quantum state if successful, or a `KetError` if an error occurred.
pub fn dump(qubits: &Quant) -> Result<Dump> {
    qubits.process.borrow_mut().dump_ref(&qubits.qubits)
}

#[cfg(test)]
mod tests {

    use crate::{dump, Process};

    use super::{ctrl, h, x, Quant};

    #[test]
    fn test_slice() {
        let p = Process::new_ptr();

        let q = Quant::new(&p, 10).unwrap();

        let tail = q.slice(1..);
        let head = q.at(0);
        print!("{:#?}", Quant::cat(&[&tail, &head]).unwrap().qubits)
    }

    #[test]
    fn test_bell() {
        let p = Process::new_ptr();

        let q = Quant::new(&p, 2).unwrap();
        h(&q.at(0)).unwrap();
        ctrl(&q.at(0), || x(&q.at(1)).unwrap()).unwrap();

        let d = dump(&q).unwrap();

        p.borrow_mut().prepare_for_execution().unwrap();

        let result = "{\n\"future\": [\n0\n],\n\"dump\": [\n{\n\"Vector\": {\n\"basis_states\": [\n[\n0\n],\n[\n3\n]\n],\n\"amplitudes_real\": [\n0.7071067811865476,\n0.7071067811865476\n],\n\"amplitudes_imag\": [\n0.0,\n0.0\n]\n}\n}\n],\n\"exec_time\": 0.00103683\n}";

        p.borrow_mut()
            .set_result(serde_json::from_str(result).unwrap())
            .unwrap();

        println!("{:#?}", d);
    }
}
