//! Fast, reliable and easy optimisation of fermion-qubit encodings.
//!
//! To simulate fermionic Hamiltonians with gate-based quantum computers,
//! it is necessary to encode the fermionic operators to qubit operators
//! which obey commutation fermionic relations.
//!
//! This file contains the PyO3 interop layer which wraps rust functions and exposes
//! these to a python API

use ::core::panic;
use log::debug;
use numpy::ndarray::Array1;
use numpy::{
    Complex64, IntoPyArray, PyArray1, PyArray2, PyArray3, PyReadonlyArray1, PyReadonlyArray2,
    PyReadonlyArrayDyn,
};
use pyo3::types::{IntoPyDict, PyComplex, PyDict, PyInt, PyString};
use pyo3::{prelude::*, pymodule, Bound};
use std::collections::HashMap;
use tinyvec::ArrayVec;
pub mod operators;
mod utils;
use crate::operators::{FermionProduct, LadderOperator, MajoranaSparse};
use crate::optimise::topphatt;
use crate::utils::*;
mod hamiltonians;
use crate::hamiltonians::QubitHamiltonian;
mod encoding;
use crate::encoding::{Encode, MajoranaEncoding};
mod optimise;
use crate::optimise::anneal_enumerations;
pub mod ternarytree;
use crate::ternarytree::{TTFlatPack, TernaryTree};

/// A Python module implemented in Rust.
#[allow(clippy::type_complexity)]
#[pymodule]
#[pyo3(name = "core")]
fn core(m: &Bound<'_, PyModule>) -> PyResult<()> {
    pyo3_log::init();
    debug!("Initializing Python module 'core'");

    #[pyfn(m)]
    #[pyo3(name = "symplectic_product")]
    fn wrap_symplectic_product_py<'py>(
        py: Python<'py>,
        left: PyReadonlyArray1<bool>,
        right: PyReadonlyArray1<bool>,
    ) -> (usize, Bound<'py, PyArray1<bool>>) {
        /*
        Computes the symplectic product between two numpy boolean arrays.

        # Simple example
        ```python
        import ferrmion
        import numpy as np
        a = np.array([True, False, True, False])
        b = np.array([False, True, False, True])
        ipower, product = ferrmion.symplectic_product(a, b)
        ```
        */
        let left = left.as_array();
        let right = right.as_array();
        let (product, ipower) = MajoranaEncoding::symplectic_product(left, right, 0);
        let pyproduct = PyArray1::from_owned_array(py, product);
        (ipower as usize, pyproduct)
    }

    #[pyfn(m)]
    #[pyo3(name = "ternary_tree_hartree_fock_state")]
    fn wrap_ternary_tree_hartree_fock_state<'py>(
        py: Python<'py>,
        fermionic_hf_state: PyReadonlyArray1<bool>,
        mode_op_map: PyReadonlyArray1<usize>,
        ipowers: PyReadonlyArray1<u8>,
        symplectic_matrix: PyReadonlyArray2<bool>,
    ) -> Bound<'py, PyArray1<bool>> {
        /*
        Computes the Hartree-Fock state from Python using numpy arrays.

        # Simple example
        ```python
        import ferrmion
        import numpy as np
        vacuum = np.zeros(6)
        hf = np.array([True, True, False, False, False, False])
        mode_op_map = np.array([0,1,2,3,4,5])
        symplectic = np.eye(6, 12, dtype=bool)
        coeffs, states = ferrmion.ternary_tree_hartree_fock_state(vacuum, hf, mode_op_map, symplectic)
        ```
        */
        let fermionic_hf_state = fermionic_hf_state.as_array();
        let mode_op_map = mode_op_map.as_array();
        let ipowers = ipowers.as_array().to_owned();
        let symplectic_matrix = symplectic_matrix.as_array().to_owned();
        let encoding = MajoranaEncoding::new(ipowers, symplectic_matrix);
        let state = encoding
            .ternary_tree_hartree_fock_state(fermionic_hf_state, mode_op_map)
            .expect("Should be able to get HF state.");
        PyArray1::from_owned_array(py, state)
    }

    #[pyfn(m)]
    #[pyo3(name = "symplectic_to_pauli")]
    fn wrap_symplectic_to_pauli<'py>(
        py: Python<'py>,
        symplectic: PyReadonlyArray1<bool>,
        ipower: u8,
    ) -> (Bound<'py, PyString>, Bound<'py, PyInt>) {
        let symplectic = symplectic.as_array();
        let (pauli, ipower) = MajoranaEncoding::symplectic_to_pauli(symplectic, ipower);
        (PyString::new(py, &pauli), PyInt::new(py, ipower))
    }

    #[pyfn(m)]
    #[pyo3(name = "pauli_to_symplectic")]
    fn wrap_pauli_to_symplectic(
        py: Python<'_>,
        pauli: String,
        ipower: usize,
    ) -> (Bound<'_, PyArray1<bool>>, Bound<'_, PyInt>) {
        // let pauli = pauli.extract();
        let (symplectic, ipower) = pauli_to_symplectic(pauli, ipower);
        (
            PyArray1::from_owned_array(py, symplectic),
            PyInt::new(py, ipower),
        )
    }

    #[pyfn(m)]
    #[pyo3(name = "symplectic_product_map")]
    fn wrap_symplectic_product_map<'py>(
        py: Python<'py>,
        ipowers: PyReadonlyArray1<u8>,
        symplectics: PyReadonlyArray2<bool>,
    ) -> (Bound<'py, PyArray2<u8>>, Bound<'py, PyArray3<bool>>) {
        let encoding = MajoranaEncoding::new(
            ipowers.as_array().to_owned(),
            symplectics.as_array().to_owned(),
        );

        let (power_map, product_map) = encoding.symplectic_product_map();
        (
            PyArray2::from_owned_array(py, power_map),
            PyArray3::from_owned_array(py, product_map),
        )
    }

    #[pyfn(m)]
    #[pyo3(name = "symplectic_to_sparse")]
    fn wrap_symplectic_to_sparse<'py>(
        py: Python<'py>,
        symplectic: PyReadonlyArray1<bool>,
        ipower: usize,
    ) -> (
        Bound<'py, PyString>,
        Bound<'py, PyArray1<usize>>,
        Bound<'py, PyComplex>,
    ) {
        let symplectic = symplectic.as_array();
        let (pauli_string, position_vec, coeff) = symplectic_to_sparse(symplectic, ipower);
        (
            PyString::new(py, &pauli_string),
            PyArray1::from_owned_array(py, position_vec),
            PyComplex::from_complex_bound(py, coeff),
        )
    }

    #[allow(clippy::too_many_arguments)]
    #[pyfn(m)]
    #[pyo3(name = "anneal_enumerations")]
    fn wrap_anneal_enumerations<'py>(
        py: Python<'py>,
        ipowers: PyReadonlyArray1<u8>,
        symplectics: PyReadonlyArray2<bool>,
        signatures: Vec<String>,
        coeffs: Vec<PyReadonlyArrayDyn<f64>>,
        temperature: f64,
        initial_guess: PyReadonlyArray1<usize>,
        coefficient_weighted: bool,
    ) -> PyResult<(Bound<'py, PyArray1<u8>>, Bound<'py, PyArray2<bool>>)> {
        let initial_guess = initial_guess.as_array();

        let msparse = MajoranaSparse::from_signatures_and_coeffs(
            signatures,
            coeffs.iter().map(|v| v.as_array()).collect(),
            0.,
        );
        let encoding = MajoranaEncoding::new(
            ipowers.as_array().to_owned(),
            symplectics.as_array().to_owned(),
        );
        let best_mode_enumeration: Array1<usize>;
        (_, best_mode_enumeration) = anneal_enumerations(
            msparse,
            encoding,
            temperature,
            initial_guess,
            coefficient_weighted,
        )
        .expect("Annealing should have succeeded.");

        let encoding = MajoranaEncoding::new(
            ipowers.as_array().to_owned(),
            symplectics.as_array().to_owned(),
        )
        .apply_mode_enumeration(best_mode_enumeration.to_vec());

        Ok((
            encoding.ipowers.into_pyarray(py),
            encoding.symplectics.into_pyarray(py),
        ))
    }

    #[pyfn(m)]
    #[pyo3(name = "flatpack_symplectic_matrix")]
    fn wrap_flatpack_symplectic_matrix(
        py: Python<'_>,
        flatpack: TTFlatPack,
    ) -> PyResult<(Bound<'_, PyArray1<u8>>, Bound<'_, PyArray2<bool>>)> {
        // ) -> PyResult<()> {
        let n_qubits: &usize = flatpack
            .iter()
            .map(|(v, _)| v)
            .max()
            .expect("Flatpack should have maxiumum qubit index.");

        let tree: TernaryTree = TernaryTree::from_flatpack_naive(&flatpack)
            .expect("Should be able to build tree from flatpack.");

        debug!("Got Tree");
        let encoding = tree
            .build_encoding(*n_qubits + 1)
            .expect("Should be able to crrate encoding from tree.");
        debug!("Got encoding");

        debug!("Got qham");
        Ok((
            encoding.ipowers.into_pyarray(py),
            encoding.symplectics.into_pyarray(py),
        ))
    }

    #[pyfn(m)]
    #[pyo3(name = "standard_symplectic_matrix")]
    fn wrap_standard_symplectic_matrix(
        py: Python<'_>,
        encoding: String,
        n_modes: usize,
    ) -> PyResult<(Bound<'_, PyArray1<u8>>, Bound<'_, PyArray2<bool>>)> {
        // ) -> PyResult<()> {

        let tree: TernaryTree = match encoding.as_str() {
            "Jordan-Wigner" | "JW" => TernaryTree::naive_jordan_wigner(n_modes),
            "Bravyi-Kitaev" | "BK" => TernaryTree::naive_bravyi_kitaev(n_modes),
            "Parity" | "PE" => TernaryTree::naive_parity(n_modes),
            "JKMN" => TernaryTree::naive_jkmn(n_modes),
            _ => panic!("Encoding must be one of JW, PE, BK or JKMN."),
        };
        debug!("Got Tree");
        let encoding = tree.build_encoding(n_modes).unwrap();
        debug!("Got encoding");

        debug!("Got qham");
        Ok((
            encoding.ipowers.into_pyarray(py),
            encoding.symplectics.into_pyarray(py),
        ))
    }

    #[pyfn(m)]
    #[pyo3(name = "fermionic_to_sparse_majorana")]
    fn fermionic_to_sparse_majorana<'py>(
        py: Python<'py>,
        signatures: Vec<String>,
        coeffs: Vec<PyReadonlyArrayDyn<f64>>,
        constant_energy: f64,
    ) -> PyResult<Bound<'py, PyDict>> {
        // ) -> PyResult<()> {
        assert_eq!(
            signatures.len(),
            coeffs.len(),
            "Signatures and coefficients should be same length"
        );

        let hamiltonian = MajoranaSparse::from_signatures_and_coeffs(
            signatures,
            coeffs.iter().map(|v| v.as_array()).collect(),
            constant_energy,
        );

        let mut output: HashMap<
            (Option<u16>, Option<u16>, Option<u16>, Option<u16>),
            numpy::Complex64,
        > = HashMap::new();
        for (key, val) in std::iter::zip(hamiltonian.indices, hamiltonian.coefficients) {
            let key_with_options = key
                .into_inner()
                .iter()
                .enumerate()
                .map(|(ind, &v)| if ind < key.len() { Some(v) } else { None })
                .collect::<ArrayVec<[Option<u16>; 4]>>()
                .into_inner();
            output
                .entry((
                    key_with_options[0],
                    key_with_options[1],
                    key_with_options[2],
                    key_with_options[3],
                ))
                .and_modify(|v| *v += val)
                .or_insert(val);
        }
        output.into_py_dict(py)
    }

    #[pyfn(m)]
    #[pyo3(name = "encode_fermion_product")]
    fn wrap_encode_product<'py>(
        py: Python<'py>,
        ipowers: PyReadonlyArray1<u8>,
        symplectics: PyReadonlyArray2<bool>,
        signatures: String,
        indices: Vec<usize>,
        coefficient: Complex64,
    ) -> PyResult<Bound<'py, PyDict>> {
        // ) -> PyResult<()> {
        assert_eq!(
            signatures.len(),
            indices.len(),
            "Signatures and indices should be same length"
        );
        let ipowers = ipowers.as_array().to_owned();
        let symplectics = symplectics.as_array().to_owned();
        let n_qubits = symplectics.ncols() / 2;
        let n_modes = symplectics.nrows() / 2;
        let vec_sig: Vec<LadderOperator> = signatures
            .chars()
            .map(|v| LadderOperator::try_from(v).expect("Signature components should be + or -"))
            .collect();

        assert!(
            n_qubits >= n_modes,
            "Must have at least as many qubits as modes."
        );

        let fproduct = FermionProduct::new(vec_sig, indices, coefficient)
            .expect("Should be able to create FermionProduct.");
        let encoding = MajoranaEncoding::new(ipowers, symplectics);
        debug!("Got encoding");
        let qham: QubitHamiltonian = encoding.encode(fproduct);
        debug!("Got Hamiltonian");

        debug!("Got qham");
        Ok(qham
            .into_py_dict(py)
            .expect("Should be able to convert QubitHamiltonian to PyDict."))
        // Ok(())
    }

    #[pyfn(m)]
    #[pyo3(name = "encode")]
    fn wrap_encode<'py>(
        py: Python<'py>,
        ipowers: PyReadonlyArray1<u8>,
        symplectics: PyReadonlyArray2<bool>,
        signatures: Vec<String>,
        coeffs: Vec<PyReadonlyArrayDyn<f64>>,
        constant_energy: f64,
    ) -> PyResult<Bound<'py, PyDict>> {
        // ) -> PyResult<()> {
        assert_eq!(
            signatures.len(),
            coeffs.len(),
            "Signatures and coefficients should be same length"
        );
        let ipowers = ipowers.as_array().to_owned();
        let symplectics = symplectics.as_array().to_owned();
        let n_qubits = symplectics.ncols() / 2;
        let n_modes = symplectics.nrows() / 2;

        assert!(
            n_qubits >= n_modes,
            "Must have at least as many qubits as modes."
        );

        let hamiltonian = MajoranaSparse::from_signatures_and_coeffs(
            signatures,
            coeffs.iter().map(|v| v.as_array()).collect(),
            constant_energy,
        );
        let encoding = MajoranaEncoding::new(ipowers, symplectics);
        debug!("Got encoding");
        let qham: QubitHamiltonian = encoding.encode(&hamiltonian);
        debug!("Got Hamiltonian");

        debug!("Got qham");
        Ok(qham
            .into_py_dict(py)
            .expect("Should be able to convert QubitHamiltonian to PyDict."))
        // Ok(())
    }

    #[pyfn(m)]
    #[pyo3(name = "encode_standard")]
    fn wrap_encode_standard<'py>(
        py: Python<'py>,
        encoding: String,
        n_modes: usize,
        n_qubits: usize,
        signatures: Vec<String>,
        coeffs: Vec<PyReadonlyArrayDyn<f64>>,
        constant_energy: f64,
    ) -> PyResult<Bound<'py, PyDict>> {
        // ) -> PyResult<()> {
        assert_eq!(
            signatures.len(),
            coeffs.len(),
            "Signatures and coefficients should be same length"
        );
        assert!(
            n_qubits >= n_modes,
            "Must have at least as many qubits as modes."
        );

        let hamiltonian = MajoranaSparse::from_signatures_and_coeffs(
            signatures,
            coeffs.iter().map(|v| v.as_array()).collect(),
            constant_energy,
        );
        debug!("Got MSparse");
        debug!("Got Hamiltonian");
        let tree: TernaryTree = match encoding.as_str() {
            "Jordan-Wigner" | "JW" => TernaryTree::naive_jordan_wigner(n_modes),
            "Bravyi-Kitaev" | "BK" => TernaryTree::naive_bravyi_kitaev(n_modes),
            "Parity" | "PE" => TernaryTree::naive_parity(n_modes),
            "JKMN" => TernaryTree::naive_jkmn(n_modes),
            _ => panic!("Encoding must be one of JW, PE, BK or JKMN."),
        };
        debug!("Got Tree");
        debug!("Hamiltonian {:?}", hamiltonian);
        debug!("Hamiltonian {:?}", hamiltonian);
        let encoding = tree.build_encoding(n_qubits).unwrap();
        debug!("Got encoding {:?}", encoding);
        debug!("Got encoding {:?}", encoding);

        let qham: QubitHamiltonian = encoding.encode(&hamiltonian);

        debug!("Got qham");
        debug!("Got qham {:?}", qham);
        Ok(qham
            .into_py_dict(py)
            .expect("Should be able to convert QubitHamiltonian to PyDict."))
        // Ok(())
    }

    #[pyfn(m)]
    #[pyo3(name = "topphatt")]
    fn wrap_topphatt<'py>(
        py: Python<'py>,
        flatpack: TTFlatPack,
        n_qubits: usize,
        signatures: Vec<String>,
        coeffs: Vec<PyReadonlyArrayDyn<f64>>,
        parallelize: bool,
    ) -> PyResult<(Bound<'py, PyArray1<u8>>, Bound<'py, PyArray2<bool>>)> {
        // ) -> PyResult<()> {
        debug!("Starting TOPPHATT");
        let flatpack: TTFlatPack = flatpack;
        debug!("Got flatpack");

        let hamiltonian = MajoranaSparse::from_signatures_and_coeffs(
            signatures,
            coeffs.iter().map(|v| v.as_array()).collect(),
            0.,
        );

        debug!("Got MSparse");
        debug!("Got Hamiltonian");
        let mut tree: TernaryTree = TernaryTree::from_flatpack_naive(&flatpack)
            .expect("Ternary tree should build from flatpack");
        debug!("Got Tree");
        debug!("Hamiltonian {:?}", hamiltonian);
        tree = topphatt(hamiltonian, tree, parallelize)
            .expect("TOPPHATT should have failed by now.");

        let encoding = tree.build_encoding(n_qubits).unwrap();
        debug!("Got encoding");
        Ok((
            encoding.ipowers.into_pyarray(py),
            encoding.symplectics.into_pyarray(py),
        ))
    }

    #[pyfn(m)]
    #[pyo3(name = "encode_topphatt")]
    fn wrap_encode_topphatt<'py>(
        py: Python<'py>,
        flatpack: TTFlatPack,
        n_qubits: usize,
        signatures: Vec<String>,
        coeffs: Vec<PyReadonlyArrayDyn<f64>>,
        constant_energy: f64,
        parallelize: bool,
    ) -> PyResult<(Bound<'py, PyArray1<u8>>, Bound<'py, PyArray2<bool>>, Bound<'py, PyDict>)> {
        debug!("Starting TOPPHATT");
        let flatpack: TTFlatPack = flatpack;
        debug!("Got flatpack");

        let hamiltonian = MajoranaSparse::from_signatures_and_coeffs(
            signatures,
            coeffs.iter().map(|v| v.as_array()).collect(),
            constant_energy,
        );

        debug!("Got MSparse");
        debug!("Got Hamiltonian");
        let mut tree: TernaryTree = TernaryTree::from_flatpack_naive(&flatpack)
            .expect("Ternary tree should build from flatpack");
        debug!("Got Tree");
        debug!("Hamiltonian {:?}", hamiltonian);
        tree = topphatt(hamiltonian.clone(), tree, parallelize)
            .expect("TOPPHATT should have failed by now.");

        let encoding = tree.build_encoding(n_qubits).unwrap();
        debug!("Got encoding");
        let qham: QubitHamiltonian = encoding.encode(&hamiltonian);
        debug!("Got qham");

        Ok((encoding.ipowers.into_pyarray(py),
            encoding.symplectics.into_pyarray(py),
            qham
            .into_py_dict(py)
            .expect("Should be able to convert QubitHamiltonian to PyDict.")))
    }

    #[pyfn(m)]
    #[pyo3(name = "topphatt_standard")]
    fn wrap_topphatt_standard<'py>(
        py: Python<'py>,
        encoding: String,
        n_modes: usize,
        n_qubits: usize,
        signatures: Vec<String>,
        coeffs: Vec<PyReadonlyArrayDyn<f64>>,
        parallelize: bool,
    ) -> PyResult<(Bound<'py, PyArray1<u8>>, Bound<'py, PyArray2<bool>>)> {
        // ) -> PyResult<Bound<'py, PyDict>> {
        assert_eq!(
            signatures.len(),
            coeffs.len(),
            "Signatures and coefficients should be same length"
        );
        assert!(
            n_qubits >= n_modes,
            "Must have at least as many qubits as modes."
        );

        debug!("Starting TOPPHATT");
        // let flatpack: TTFlatPack = node_map.extract::<TTFlatPack>()?;
        let hamiltonian = MajoranaSparse::from_signatures_and_coeffs(
            signatures,
            coeffs.iter().map(|v| v.as_array()).collect(),
            0.,
        );
        debug!("Got MSparse");
        debug!("Got Hamiltonian");
        let mut tree: TernaryTree = match encoding.as_str() {
            "Jordan-Wigner" | "JW" => TernaryTree::naive_jordan_wigner(n_modes),
            "Bravyi-Kitaev" | "BK" => TernaryTree::naive_bravyi_kitaev(n_modes),
            "Parity" | "PE" => TernaryTree::naive_parity(n_modes),
            "JKMN" => TernaryTree::naive_jkmn(n_modes),
            _ => panic!("Encoding must be one of JW, PE, BK or JKMN."),
        };
        debug!("Got Tree");
        debug!("Hamiltonian {:?}", hamiltonian);
        tree = topphatt(hamiltonian.clone(), tree, parallelize)
            .expect("TOPPHATT should have failed by now.");
        let encoding = tree.build_encoding(n_qubits).unwrap();
        debug!("Got encoding");
        Ok((
            encoding.ipowers.into_pyarray(py),
            encoding.symplectics.into_pyarray(py),
        ))
        // let qham: QubitHamiltonian = encoding.encode(&hamiltonian);
        // debug!("Got qham");
        // Ok(qham
        //     .into_py_dict(py)
        //     .expect("Should be able to convert QubitHamiltonian to PyDict."))
        // Ok(())
    }
    Ok(())
}
