# based on https://github.com/Qiskit-Extensions/qiskit-alt/blob/main/bench/jordan_wigner_nature_time.py

from timeit import timeit

from qiskit import __qiskit_version__
from qiskit_nature.drivers import UnitsType, Molecule
from qiskit_nature.drivers.second_quantization import (
    ElectronicStructureDriverType,
    ElectronicStructureMoleculeDriver,
)
from qiskit_nature.problems.second_quantization import ElectronicStructureProblem
from qiskit_nature.converters.second_quantization import QubitConverter
from qiskit_nature.mappers.second_quantization import JordanWignerMapper
from qiskit_nature.operators.second_quantization import FermionicOp

geometry = {
    "h2": [["H", [0.0, 0.0, 0.0]], ["H", [0.0, 0.0, 0.735]]],
    "h2o": [["O", [0.0, 0.0, 0.0]], ["H", [0.757, 0.586, 0.0]], ["H", [-0.757, 0.586, 0.0]]],
}

def setup(basis, geom):
    molecule = Molecule(geometry=geom, charge=0, multiplicity=1)
    driver = ElectronicStructureMoleculeDriver(
        molecule, basis=basis, driver_type=ElectronicStructureDriverType.PYSCF
    )
    es_problem = ElectronicStructureProblem(driver)
    second_q_op = es_problem.second_q_ops()
    qubit_converter = QubitConverter(mapper=JordanWignerMapper())
    hamiltonian = second_q_op[0].reduce()
    return qubit_converter, hamiltonian


def run_one_basis(basis, molecule, num_repetitions):
    qubit_converter, hamiltonian = setup(basis, geometry[molecule])
    time = timeit(lambda: qubit_converter.convert(hamiltonian), number=num_repetitions)
    t = 1000 * time / num_repetitions
    print(f"geometry={molecule}, basis={basis}, size={len(hamiltonian)}, time={t:0.2f} ms")
    return t

print(__qiskit_version__)
for basis, molecule, num_repetitions in [
    ("sto3g", "h2", 10),
    ("631g", "h2", 10),
    ("631++g", "h2", 5),
    ("sto3g", "h2o", 5),
    ("631g", "h2o", 1),
]:
    run_one_basis(basis, molecule, num_repetitions)

# Pure-Python PoC of unique
# https://github.com/Qiskit/qiskit-terra/pull/7656/commits/7e886dde94a7c2e46345d85b97b8bb4f7c928fc0
def _unique(array):
    # This function corresponds to
    # _, indexes, inverses = np.unique(array, return_index=True, return_inverse=True, axis=0)
    table = {}
    indexes = []
    inverses = np.empty(array.shape[0], dtype=int)
    for i, ary in enumerate(array):
        b = ary.data.tobytes()
        if b in table:
            inverses[i] = table[b]
        else:
            inverses[i] = table[b] = len(table)
            indexes.append(i)
    return indexes, inverses