import sys
import numpy as np
import multiprocessing
from time import time

from qiskit_nature.transformers import FreezeCoreTransformer

mol_strings = {
            'H2': 'H .0 .0 .0; H .0 .0 0.735',
            'LiH': 'H .0 .0 .0; Li .0 .0 2.5',
            'HF': 'H .0 .0 .0; F .0 .0 1.25',
}

mol_string = mol_strings['LiH']
#mol_string = mol_strings['H2']
method = 'qasm_simulator'
#method = 'statevector_simulator'
threads = 32

from qiskit_nature.drivers import PySCFDriver, UnitsType
driver = PySCFDriver(atom=mol_string, unit=UnitsType.ANGSTROM, basis="sto3g")

from qiskit_nature.problems.second_quantization import ElectronicStructureProblem
problem = ElectronicStructureProblem(driver)

second_q_ops = problem.second_q_ops()

main_op = second_q_ops[0]

num_particles = (
    problem.molecule_data_transformed.num_alpha,
    problem.molecule_data_transformed.num_beta,
    )

num_spin_orbitals = 2 * problem.molecule_data.num_molecular_orbitals

# from qiskit.algorithms.optimizers import L_BFGS_B
# optimizer = L_BFGS_B()

from qiskit.algorithms.optimizers import SLSQP
optimizer = SLSQP(maxiter=5000)

from qiskit_nature.mappers.second_quantization import ParityMapper
from qiskit_nature.converters.second_quantization import QubitConverter

converter = QubitConverter(mapper=ParityMapper(), two_qubit_reduction=False)

qubit_op = converter.convert(main_op, num_particles=num_particles)

from qiskit_nature.circuit.library import HartreeFock

init_state = HartreeFock(num_spin_orbitals, num_particles, converter)

from qiskit_nature.circuit.library import UCCSD
#from qiskit_nature.circuit.library import SUCCD

ansatz = UCCSD(
    converter,
    num_particles,
    num_spin_orbitals,
    initial_state=init_state,
)

# set the backend for the quantum computation
from qiskit.utils import QuantumInstance
from qiskit_nature.algorithms import VQEUCCFactory
from qiskit import Aer

backend = Aer.get_backend(method)
quantum_instance = QuantumInstance(backend = backend)
quantum_instance.backend_options['backend_options'] = {'max_parallel_experiments':threads} 

vqe_solver = VQEUCCFactory(quantum_instance, optimizer=optimizer, include_custom=False, ansatz=ansatz, initial_state=init_state)
vqe_solver._vqe._max_evals_grouped = 128
optimizer.set_max_evals_grouped(vqe_solver._vqe._max_evals_grouped)

from qiskit_nature.algorithms import GroundStateEigensolver
calc = GroundStateEigensolver(converter, vqe_solver)
res = calc.solve(problem)

print(res)