/*
Utility functions for core functionality.
*/
use ndarray::{concatenate, Axis, Zip};
use numpy::ndarray::{Array1, ArrayView1};
use numpy::Complex64;

/// Converts an imaginary count to a complex sign.
///
/// # Examples
///
/// ```
/// use ferrmion::utils::icount_to_sign;
/// use num_complex::Complex64;
/// assert_eq!(icount_to_sign(0), Complex64::new(1.0, 0.0));
/// ```
pub fn icount_to_sign(icount: usize) -> Complex64 {
    match icount % 4 {
        0 => Complex64::new(1., 0.),
        1 => Complex64::new(0., 1.),
        2 => Complex64::new(-1., 0.),
        3 => Complex64::new(0., -1.),
        _ => unreachable!(),
    }
}

/// Computes the Kronecker product of two vectors.
///
/// # Examples
///
/// ```
/// use ferrmion::utils::vector_kron;
/// use ndarray::arr1;
/// use num_complex::Complex64;
/// let left = arr1(&[Complex64::new(1.0, 0.0)]);
/// let right = arr1(&[Complex64::new(1.0, 0.0), Complex64::new(0.0, 1.0)]);
/// let result = vector_kron(&left, &right);
/// ```
///
// This is not actually dead code, it's used in tests.
#[allow(dead_code)]
pub fn vector_kron(left: &Array1<Complex64>, right: &Array1<Complex64>) -> Array1<Complex64> {
    concatenate![
        Axis(0),
        left.mapv(|l| l * right[0]),
        left.mapv(|l| l * right[1])
    ]
}

/// Converts symplectic representation to sparse form.
///
/// # Examples
///
/// ```
/// use ferrmion::utils::symplectic_to_sparse;
/// use ndarray::arr1;
/// let symplectic = arr1(&[true, false, false, true]);
/// let (pauli, indices, coeff) = symplectic_to_sparse(symplectic.view(), 0);
/// ```
pub fn symplectic_to_sparse(
    symplectic: ArrayView1<bool>,
    ipower: usize,
) -> (String, Array1<usize>, Complex64) {
    let mut ipower: usize = ipower;
    let (x_block, z_block) = symplectic.split_at(Axis(0), symplectic.len_of(Axis(0)) / 2);
    let mut pauli_string = String::new();
    let mut indices: Vec<usize> = Vec::new();
    let mut index: usize = 0;
    Zip::from(x_block).and(z_block).for_each(|&x, &z| {
        let pauli = match (&x, &z) {
            (&false, &false) => 'I',
            (&true, &true) => 'Y',
            (&true, &false) => 'X',
            (&false, &true) => 'Z',
        };
        match pauli {
            'I' => {}
            _ => {
                pauli_string.push(pauli);
                indices.push(index);
                ipower += 3
            }
        };
        index += 1;
    });
    let ipower = ipower % 4;
    (
        pauli_string,
        ndarray::arr1(&indices),
        icount_to_sign(ipower),
    )
}

#[test]
fn test_symplectic_to_sparse() {
    let symplectic: ndarray::ArrayBase<ndarray::OwnedRepr<bool>, ndarray::Dim<[usize; 1]>> =
        ndarray::arr1(&[true, true, false, false, true, false, true, false]);
    assert_eq!(
        symplectic_to_sparse(symplectic.view(), 0),
        (
            "YXZ".to_string(),
            ndarray::arr1(&[0, 1, 2]),
            Complex64::new(0., 1.)
        )
    );
}

fn _valid_pauli_string(pauli: &str) -> bool {
    pauli.chars().all(|c| matches!(c, 'X' | 'Y' | 'Z' | 'I'))
}

#[test]
fn test_valid_pauli_string() {
    assert!(_valid_pauli_string("XYZI"));
    assert!(_valid_pauli_string("X"));
    assert!(_valid_pauli_string("Y"));
    assert!(_valid_pauli_string("Z"));
    assert!(_valid_pauli_string("I"));
    assert!(!_valid_pauli_string("XYZA"));
}

/// Converts a Pauli string to symplectic representation.
///
/// # Examples
///
/// ```
/// use ferrmion::utils::pauli_to_symplectic;
/// let (symplectic, ipower) = pauli_to_symplectic("XY".to_string(), 0);
/// ```
pub fn pauli_to_symplectic(pauli: String, ipower: usize) -> (Array1<bool>, usize) {
    let string_len = pauli.len();
    assert!(_valid_pauli_string(&pauli));

    let mut ipower: usize = ipower;
    let mut x_block: Array1<bool> = Array1::from_elem(string_len, false);
    let mut z_block: Array1<bool> = Array1::from_elem(string_len, false);
    Zip::from(&Array1::from_iter(pauli.chars()))
        .and(&mut x_block)
        .and(&mut z_block)
        .for_each(|c, x, z| match c {
            'X' => *x = true,
            'Z' => *z = true,
            'Y' => {
                *x = true;
                *z = true;
                ipower += 1;
            }
            _ => {
                *x = false;
                *z = false;
            }
        });
    (concatenate![Axis(0), x_block, z_block], ipower % 4)
}

#[test]
fn test_pauli_to_symplectic() {
    let valid_string = String::from("IXZY");
    let valid_symplectic = ndarray::arr1(&[false, true, false, true, false, false, true, true]);
    assert_eq!(pauli_to_symplectic(valid_string, 0), (valid_symplectic, 1));

    let all_y = String::from("YYY");
    assert_eq!(
        pauli_to_symplectic(all_y, 0),
        (ndarray::arr1(&[true, true, true, true, true, true]), 3)
    )
}
