rqism 0.4.1

use std::{borrow::Cow, f32::consts::PI};

use ndarray::{array, Array, Array2};
use num::complex::Complex;

use lazy_static::lazy_static;

use crate::{circuit::Circuit, sparse_mat::SMat};

const ONE_SQR_TWO: f32 = 0.707_106_77;

//TODO: efficient measure_into instruction

pub mod gate_matrices {
    use super::*;

    lazy_static! {
        pub static ref PAULI_X: Array2<Complex<f32>> = Array::from_shape_vec(
            (2, 2),
            vec![
                Complex::new(0.0, 0.0),
                Complex::new(1.0, 0.0),
                Complex::new(1.0, 0.0),
                Complex::new(0.0, 0.0),
            ],
        )
        .unwrap();
        pub static ref PAULI_Y: Array2<Complex<f32>> = Array::from_shape_vec(
            (2, 2),
            vec![
                Complex::new(0.0, 0.0),
                Complex::new(0.0, -1.0),
                Complex::new(0.0, 1.0),
                Complex::new(0.0, 0.0),
            ],
        )
        .unwrap();
        pub static ref PAULI_Z: Array2<Complex<f32>> = Array::from_shape_vec(
            (2, 2),
            vec![
                Complex::new(-1.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 1.0),
            ],
        )
        .unwrap();
        pub static ref CLIFFORD_H: Array2<Complex<f32>> = Complex::new(ONE_SQR_TWO, 0.0)
            * Array::from_shape_vec(
                (2, 2),
                vec![
                    Complex::new(1.0, 0.0),
                    Complex::new(1.0, 0.0),
                    Complex::new(1.0, 0.0),
                    Complex::new(-1.0, 0.0),
                ],
            )
            .unwrap();
        pub static ref CLIFFORD_S: Array2<Complex<f32>> = Array::from_shape_vec(
            (2, 2),
            vec![
                Complex::new(1.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 1.0),
            ],
        )
        .unwrap();
        pub static ref CLIFFORD_CNOT: Array2<Complex<f32>> = Array::from_shape_vec(
            (4, 4),
            vec![
                Complex::new(1.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(1.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(1.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(1.0, 0.0),
                Complex::new(0.0, 0.0),
            ],
        )
        .unwrap();
        pub static ref SWAP: Array2<Complex<f32>> = array![
            [
                Complex::<f32>::from(1.0),
                Complex::<f32>::from(0.0),
                Complex::<f32>::from(0.0),
                Complex::<f32>::from(0.0)
            ],
            [
                Complex::<f32>::from(0.0),
                Complex::<f32>::from(0.0),
                Complex::<f32>::from(1.0),
                Complex::<f32>::from(0.0)
            ],
            [
                Complex::<f32>::from(0.0),
                Complex::<f32>::from(1.0),
                Complex::<f32>::from(0.0),
                Complex::<f32>::from(0.0)
            ],
            [
                Complex::<f32>::from(0.0),
                Complex::<f32>::from(0.0),
                Complex::<f32>::from(0.0),
                Complex::<f32>::from(1.0)
            ]
        ];
        pub static ref T_GATE: Array2<Complex<f32>> = Array::from_shape_vec(
            (2, 2),
            vec![
                Complex::new(1.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(0.0, 0.0),
                Complex::new(ONE_SQR_TWO + ONE_SQR_TWO, 1.0),
            ],
        )
        .unwrap();
        pub static ref CZ: Array2<Complex<f32>> = {
            let mut e: Array2<Complex<f32>> = Array2::eye(4);

            *e.get_mut((3, 3)).unwrap() = Complex::new(-1.0, 0.0);

            e
        };
    }
}

#[derive(Copy, Clone, Debug)]
pub enum PauliKind {
    X,
    Y,
    Z,
}

impl PauliKind {
    pub fn to_matrix(&self) -> &Array2<Complex<f32>> {
        match self {
            Self::X => &gate_matrices::PAULI_X,

            Self::Y => &gate_matrices::PAULI_Y,
            Self::Z => &gate_matrices::PAULI_Z,
        }
    }
}

#[derive(Copy, Clone, Debug)]
pub enum CliffordKind {
    H,
    S,
    CNOT,
}

impl CliffordKind {
    pub fn to_matrix(&self) -> &Array2<Complex<f32>> {
        match self {
            Self::H => &gate_matrices::CLIFFORD_H,

            Self::S => &gate_matrices::CLIFFORD_S,

            Self::CNOT => &gate_matrices::CLIFFORD_CNOT,
        }
    }
}

#[derive(Copy, Clone)]
pub enum CliffordTarget {
    Single { i: usize },
    Two { i: usize, j: usize },
}

impl CliffordTarget {
    pub fn to_vec(&self) -> Vec<usize> {
        match self {
            Self::Single { i } => vec![*i],
            Self::Two { i, j } => vec![*i, *j],
        }
    }

    pub fn to_single(&self) -> usize {
        match self {
            Self::Single { i } => *i,
            Self::Two { .. } => panic!("called to_singe on two clifford indices"),
        }
    }

    pub fn to_two(&self) -> (usize, usize) {
        match self {
            Self::Single { .. } => panic!("called to_singe on two clifford indices"),
            Self::Two { i, j } => (*i, *j),
        }
    }
}

/// Optional gate kind used for optimization.
#[derive(Copy, Clone)]
pub enum GateKind {
    /// 1-qubit gates that only change phase.
    SQPhase,
    /// n-qubit permutation gates.
    /// There is no kind for sq permuatation gates because the only sq
    /// permutation gates are identity, wich shouldnt be there anyway and the
    /// X gate, which is [implemented seperately](Instruction::not).
    Perm,
    /// 1-qubit superposition gate (e.g. hadamard)
    SQSuperpos,
}

#[derive(Clone)]
pub enum Instruction<'a> {
    Pauli {
        kind: PauliKind,
        target: usize,
    },

    Clifford {
        kind: CliffordKind,
        target: CliffordTarget,
    },

    /// Custon n-qubit gate
    Gate {
        matrix: Cow<'a, Array2<Complex<f32>>>,
        indices: Vec<usize>,
        kind: Option<GateKind>,
    },

    Measure {
        indices: Vec<usize>,
    },

    MeasureNoisy {
        indices: Vec<usize>,
        p: f32,
    },

    /// Measure the entire state
    MeasureState,

    Identity,

    /// Matrix that is directly dotted with the statevector with no modifications
    StateVecMat {
        mat: Cow<'a, SMat>,
    },

    /// Classical instruction that applies a quantum instruction if a boolean
    /// expression evaluates to true. The inputs to the expression are the
    /// states of the qubits specified in the indices field. This instruction
    /// does not measure the qubits, so you generally want to measure the
    /// qubits before this instruction
    ConditionalInstruction {
        indices: Vec<usize>,
        condition: fn(Vec<bool>) -> bool,
        inst: Box<[Instruction<'a>]>,
    },
}

impl<'a> Instruction<'a> {
    pub fn hadamard(i: usize) -> Self {
        Self::Clifford {
            kind: CliffordKind::H,
            target: CliffordTarget::Single { i },
        }
    }

    pub fn not(i: usize) -> Self {
        Self::Pauli {
            kind: PauliKind::X,
            target: i,
        }
    }

    pub fn swap(i: [usize; 2]) -> Self {
        Self::Gate {
            matrix: Cow::Borrowed(&gate_matrices::SWAP),
            indices: i.to_vec(),
            kind: None,
        }
    }

    pub fn cnot(t: [usize; 2]) -> Self {
        Self::Clifford {
            kind: CliffordKind::CNOT,
            target: CliffordTarget::Two { i: t[0], j: t[1] },
        }
    }

    pub fn cp(theta: f32, i: [usize; 2]) -> Self {
        let mut gate = Array2::<Complex<f32>>::eye(4);

        *gate.get_mut((3, 3)).unwrap() =
            Complex::new(std::f32::consts::E, 0.0).powc(Complex::new(0.0, 1.0) * theta);

        Self::Gate {
            matrix: Cow::Owned(gate),
            indices: i.to_vec(),
            kind: None,
        }
    }

    pub fn t_gate(i: usize) -> Self {
        Self::Gate {
            matrix: Cow::Borrowed(&gate_matrices::T_GATE),
            indices: vec![i],
            kind: Some(GateKind::SQPhase),
        }
    }

    pub fn s_gate(i: usize) -> Self {
        Self::Clifford {
            kind: CliffordKind::S,
            target: CliffordTarget::Single { i },
        }
    }

    pub fn rx(i: usize, theta: f32) -> Self {
        let (c, s) = ((theta * 0.5).cos(), (theta * 0.5).sin());
        let gate = Array2::<Complex<f32>>::from_shape_vec(
            (2, 2),
            vec![
                Complex::new(c, 0.0),
                Complex::new(s, -1.0),
                Complex::new(s, -1.0),
                Complex::new(c, 0.0),
            ],
        )
        .unwrap();

        Self::Gate {
            matrix: Cow::Owned(gate),
            indices: vec![i],
            kind: None,
        }
    }

    pub fn ry(i: usize, theta: f32) -> Self {
        let (c, s) = ((theta * 0.5).cos(), (theta * 0.5).sin());

        let gate = Array2::<Complex<f32>>::from_shape_vec(
            (2, 2),
            vec![
                Complex::new(c, 0.0),
                Complex::new(-s, 0.0),
                Complex::new(s, 0.0),
                Complex::new(c, 0.0),
            ],
        )
        .unwrap();

        Self::Gate {
            matrix: Cow::Owned(gate),
            indices: vec![i],
            kind: None,
        }
    }

    pub fn cz(a: usize, b: usize) -> Self {
        Self::Gate {
            matrix: Cow::Borrowed(&gate_matrices::CZ),
            indices: vec![a, b],
            kind: None,
        }
    }
}

fn qft_r(n: usize, ins: &mut Vec<Instruction>) {
    if n == 0 {
        return;
    }

    let n = n - 1;
    ins.push(Instruction::hadamard(n));

    for q in 0..n {
        ins.push(Instruction::cp(PI / 2.0_f32.powf((n - q) as f32), [q, n]));
    }
}

fn swap_reg(n: usize, ins: &mut Vec<Instruction>) {
    for q in 0..n / 2 {
        ins.push(Instruction::swap([q, n - q - 1]));
    }
}

/// Generates an n-qubit QFT circuit
pub fn quantum_fourier_transform(n: usize) -> Circuit<'static> {
    let mut ins = vec![];

    qft_r(n, &mut ins);
    swap_reg(n, &mut ins);

    Circuit { ins, n }
}