###### Import Libraries
import numpy as np

from qiskit import Aer
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

import time

molecule = Molecule(
    geometry=[["H", [0.0, 0.0, 0.0]], ["F", [0.0, 0.0, 0.9554]]], charge=0, multiplicity=1
)

print(molecule.geometry)

###### Define the electronic structure driver

driver = ElectronicStructureMoleculeDriver(
    molecule, basis="sto3g", driver_type=ElectronicStructureDriverType.PYSCF,
)

# define the electronic structure problem as a function of the driver
es_problem = ElectronicStructureProblem(driver) 

# build the second quantized Hamiltonian of the molecule
second_q_op = es_problem.second_q_ops()


# print the Hamiltonian, no. of spin orbitals or qubits, and the no. of paticles
print('\n',"The second quantized form of the electronic Hamiltonian is" , second_q_op[0])
print('\n',"The no. of spin orbitals is", es_problem.num_spin_orbitals) 
print('\n', "The total no. of electrons is", es_problem.num_particles)

qubit_converter = QubitConverter(JordanWignerMapper())

%%time
from qiskit.algorithms import NumPyMinimumEigensolver
from qiskit_nature.algorithms import GroundStateEigensolver

##############################
# Provide your code here

numpy_solver = NumPyMinimumEigensolver()
numpy_ground_state_solver = GroundStateEigensolver(qubit_converter=qubit_converter, solver=numpy_solver)
numpy_results = numpy_ground_state_solver.solve(problem=es_problem)

##############################

exact_energy = np.real(numpy_results.electronic_energies + numpy_results.nuclear_repulsion_energy)[0]
print(f"Exact electronic energy: {exact_energy:.6f} Hartree\n")
print(numpy_results)

%%time
from qiskit.utils import QuantumInstance
from qiskit_nature.algorithms import VQEUCCFactory

#backend
backend=Aer.get_backend("aer_simulator_statevector")

# Select the simulator to run the calculation
quantum_instance = QuantumInstance(backend=backend)

# define the VQE simulator
vqe_solver = VQEUCCFactory(quantum_instance)

%%time
from qiskit_nature.algorithms import GroundStateEigensolver

# do the calculation
calc = GroundStateEigensolver(qubit_converter, vqe_solver) 

# save the result
res = calc.solve(es_problem) 

# print the result
print(res) 

from qiskit.algorithms.optimizers import COBYLA
from qiskit_nature.drivers.second_quantization import PySCFDriver

driver_reduced = PySCFDriver(atom='H 0.0, 0.0, 0.0; F 0.0, 0.0,  0.9554',
                     unit=UnitsType.ANGSTROM, charge=0, spin=0, basis='STO3G')

es_problem = ElectronicStructureProblem(driver_reduced) 

properties = driver_reduced.run()

from qiskit_nature.circuit.library import HartreeFock
from qiskit.circuit.library import EfficientSU2
from qiskit_nature.properties.second_quantization.electronic import ParticleNumber


particle_number = properties.get_property(ParticleNumber)
init_state = HartreeFock(num_spin_orbitals=particle_number.num_spin_orbitals, 
                         num_particles=particle_number.num_particles,
                         qubit_converter=qubit_converter)

init_state.barrier()

##############################

ansatz = EfficientSU2(num_qubits=particle_number.num_spin_orbitals,
                      reps=1,
                      entanglement='linear',
                      initial_state=init_state,
                      insert_barriers=True)

##############################

ansatz.decompose().draw('mpl')

from qiskit.providers.aer import StatevectorSimulator
from qiskit.algorithms.optimizers import  L_BFGS_B
from qiskit.algorithms import VQE

##############################
# Provide your code here

backend1 = StatevectorSimulator()
optimizer = L_BFGS_B(maxfun=400)

##############################

from qiskit_nature.algorithms import VQEUCCFactory, GroundStateEigensolver
from jupyterplot import ProgressPlot

error_threshold = 10 # mHartree

np.random.seed(5)  # fix seed for reproducibility
initial_point = np.random.random(ansatz.num_parameters)

from qiskit.providers.aer import StatevectorSimulator
from qiskit.algorithms.optimizers import  L_BFGS_B
from qiskit.algorithms import VQE

##############################
# Provide your code here

backend1 = StatevectorSimulator()
optimizer = L_BFGS_B(maxfun=400)

##############################

from qiskit_nature.algorithms import VQEUCCFactory, GroundStateEigensolver
from jupyterplot import ProgressPlot

error_threshold = 10 # mHartree

np.random.seed(5)  # fix seed for reproducibility
initial_point = np.random.random(ansatz.num_parameters)

%%time
# for live plotting
pp = ProgressPlot(plot_names=['Energy'],
                  line_names=['Runtime VQE', f'Target + {error_threshold}mH', 'Target']) 

intermediate_info = {
    'nfev': [],
    'parameters': [],
    'energy': [],
    'stddev': []
}

def callback(nfev, parameters, energy, stddev):
    intermediate_info['nfev'].append(nfev)
    intermediate_info['parameters'].append(parameters)
    intermediate_info['energy'].append(energy)
    intermediate_info['stddev'].append(stddev)
    pp.update([[energy, exact_energy+error_threshold/1000, exact_energy]])

##############################
# Provide your code here

vqe = VQE(ansatz=ansatz,
          optimizer=optimizer,
          quantum_instance=backend1,
          callback=callback,
          initial_point=initial_point)

vqe_ground_state_solver = GroundStateEigensolver(qubit_converter=qubit_converter, solver=vqe)

vqe_results = vqe_ground_state_solver.solve(problem=es_problem)

##############################

print(vqe_results)

error = (vqe_results.computed_energies[0] - exact_energy) * 1000 # mHartree
print(f'Error is: {error:.3f} Hartree')

%%time

WITHOUT__ActiveSpace = ElectronicStructureProblem(driver_reduced)
second_q_ops = WITHOUT__ActiveSpace.second_q_ops()
main_op = second_q_ops[0]

# particle_number = problem.molecule_data_transformed
# num_particles = (particle_number.num_alpha, particle_numbaer.num_beta)
# num_spin_orbitals = particle_number.num_molecular_orbitals*2

molecule = driver.run()
simulation = QuantumInstance(backend=backend)

mapper = JordanWignerMapper()
qubit_converter = QubitConverter(mapper=mapper, two_qubit_reduction=True)

qubit_op = qubit_converter.convert(main_op, num_particles=2)

%%time
shots=10

begin = time.time()

VQE = VQEUCCFactory(simulation, optimizer=COBYLA(maxiter=shots))
calculation = GroundStateEigensolver(qubit_converter, VQE)
result0 = calculation.solve(WITHOUT__ActiveSpace)

end = time.time()
print(f"Total runtime of the program is {end - begin}")

np.real(result0.electronic_energies + result0.nuclear_repulsion_energy)[0]

# from qiskit_nature.transformers import ActiveSpaceTransformer
from qiskit_nature.transformers.second_quantization.electronic import ActiveSpaceTransformer
transformer = ActiveSpaceTransformer(2, 2, [4,5]) 

%%time

WITH_ActiveSpace = ElectronicStructureProblem(driver_reduced, transformers=[transformer])
second_q_ops = WITH_ActiveSpace.second_q_ops()
main_op = second_q_ops[0]

# particle_number = problem.molecule_data_transformed
# num_particles = (particle_number.num_alpha, particle_number.num_beta)
# num_spin_orbitals = particle_number.num_molecular_orbitals*2

molecule = driver.run()
simulation = QuantumInstance(backend=backend)

mapper = JordanWignerMapper()
qubit_converter = QubitConverter(mapper=mapper, two_qubit_reduction=True)

qubit_op = qubit_converter.convert(main_op, num_particles=2)

shots=10

simulation = QuantumInstance(backend=backend)

begin = time.time()

VQE = VQEUCCFactory(simulation, optimizer=COBYLA(maxiter=shots))
calculation = GroundStateEigensolver(qubit_converter, VQE)
result1 = calculation.solve(WITH_ActiveSpace)

end = time.time()
print(f"Total runtime of the program is {end - begin}")

np.real(result1.electronic_energies + result1.nuclear_repulsion_energy)[0]



