from qiskit_nature.drivers import Molecule
from qiskit_nature.drivers.second_quantization import (
    ElectronicStructureDriverType,
    ElectronicStructureMoleculeDriver)
from qiskit_nature.problems.second_quantization import ElectronicStructureProblem

from qiskit_nature.transformers.second_quantization.electronic import FreezeCoreTransformer

molecule = Molecule([["O",[0.0,0.0,0.0]],["H", [0.757,0.586,0.0]], ["H", [-0.757,0.586,0.0]]], charge = 0, multiplicity = 1)
driver = ElectronicStructureMoleculeDriver(molecule, basis = "sto3g", driver_type = ElectronicStructureDriverType.PYSCF)
es_problem = ElectronicStructureProblem(driver, transformers=[FreezeCoreTransformer(freeze_core=True,
                               remove_orbitals=[])])
second_q_ops = es_problem.second_q_ops()
from qiskit_nature.converters.second_quantization import QubitConverter
from qiskit_nature.mappers.second_quantization import JordanWignerMapper, ParityMapper
mapper = ParityMapper()
converter = QubitConverter(mapper = mapper, two_qubit_reduction = True)
qubit_op = converter.convert(second_q_ops["ElectronicEnergy"],num_particles = es_problem.num_particles)
print(qubit_op)

from qiskit.circuit.library import PauliTwoDesign, TwoLocal, EfficientSU2
from qiskit_nature.circuit.library import HartreeFock
init_state = HartreeFock(es_problem.num_spin_orbitals, es_problem.num_particles, converter)
#init_state = None
rotation_blocks = ['ry','rx']
entanglement_blocks = ['crx']
ansatz = TwoLocal(qubit_op.num_qubits, rotation_blocks, entanglement_blocks, 'sca', insert_barriers = False,initial_state = init_state)

from qiskit.algorithms.optimizers import COBYLA, SPSA, SLSQP, NELDER_MEAD
optimizer = COBYLA(maxiter=50)

from qiskit_ibm_runtime import Estimator, QiskitRuntimeService, Session, Options

service = QiskitRuntimeService(channel = "ibm_quantum")

options = Options()
options.execution.shots = 2000
options.resilience_level = 0
options.optimization_level = 0
estimator = Estimator(session=Session(service, backend = "ibmq_qasm_simulator"), options=options)

print(Options._DEFAULT_RESILIENCE_LEVEL,Options._DEFAULT_OPTIMIZATION_LEVEL)



counts = []
values = []

def store_intermediate_result(eval_counts, parameters, mean, std):
    counts.append(eval_counts)
    values.append(mean)

from qiskit.algorithms.minimum_eigensolvers import VQE

vqe = VQE(estimator, ansatz, optimizer,callback=store_intermediate_result)
result = vqe.compute_minimum_eigenvalue(qubit_op)
    
print(result)

from qiskit.algorithms import NumPyEigensolver
exact = NumPyEigensolver()
exact_run = exact.compute_eigenvalues(qubit_op)
print(exact_run)

print(es_problem.grouped_property_transformed._properties["ElectronicEnergy"])

result_freeze=result
values_freeze = values
for i in range(len(values)):
    values_freeze[i] = values[i]-60.66177520599246
counts_freeze = counts

import matplotlib.pyplot as plt

plt.rcParams["font.size"] = 14

# Plot energy and reference value
plt.figure(figsize=(12, 6))
plt.plot(counts_freeze, values_freeze, label='spsa')
plt.axhline(y=-84.206, color="tab:red", ls="--", label="Target")

plt.legend(loc="best")
plt.xlabel("Iteration")
plt.ylabel("Energy [H]")
plt.title("VQE energy")
plt.show()

#Not to be shown in code walkthrough; just for me to have.
import numpy
f = open("qasm_spsa_H2O", "w+")
f.write("# optimization_step Energy\n")        # column names
numpy.savetxt(f, numpy.array([counts_freeze, values_freeze]).T)
#x, y = numpy.loadtxt("trialdata", unpack=True)

from qiskit import IBMQ
provider = IBMQ.load_account()

from qiskit_ibm_runtime import Estimator, QiskitRuntimeService, Session, Options

optimizer = SPSA(maxiter = 50)

counts = []
values = []

#def store_intermediate_result(eval_counts, parameters, mean, std):
#    counts.append(eval_counts)
#    values.append(mean)


service = QiskitRuntimeService()

backend = "ibm_washington"

with Session(service=service, backend=backend) as session:
    options = Options()
    options.execution.shots = 2000
    options.resilience_level = 1
    options.optimization_level = 3

    vqe = VQE(Estimator(session=session, options=options),
              ansatz, optimizer, callback=store_intermediate_result)
    result = vqe.compute_minimum_eigenvalue(qubit_op)
    print(result)

session.close()

result_real=result
counts_real=counts
values_real=values

#Not to be shown in code walkthrough. Just for me to have.
f = open("wash_spsa_H2O", "w+")
f.write("# optimization_step Energy\n")        # column names
numpy.savetxt(f, numpy.array([counts_real, values_real]).T)
#x, y = numpy.loadtxt("trialdata", unpack=True)

for i in range(len(values)):
    values_real[i] = values[i]-60.66177520599246

plt.rcParams["font.size"] = 14

# Plot energy and reference value
plt.figure(figsize=(12, 6))
plt.plot(counts_real, values_real, label='spsa on ibm_washington')
plt.axhline(y=-84.206, color="tab:red", ls="--", label="Target")

plt.legend(loc="best")
plt.xlabel("Iteration")
plt.ylabel("Energy [H]")
plt.title("VQE energy")
plt.show()

counts = []
values = []

with Session(service=service, backend=backend) as session:
    options = Options()
    options.execution.shots = 2000
    options.resilience_level = 0
    options.optimization_level = 0

    vqe = VQE(Estimator(session=session, options=options),
              ansatz, optimizer, callback=store_intermediate_result)
    result = vqe.compute_minimum_eigenvalue(qubit_op)
    print(result)

session.close()

result_zero=result
values_zero = values
for i in range(len(values)):
    values_zero[i] = values[i]-60.66177520599246
counts_zero = counts

#Not to be shown in walkthrough. Just for me to have.
f = open("qasm_spsa_H2O", "w+")
f.write("# optimization_step Energy\n")        # column names
numpy.savetxt(f, numpy.array([counts_zero, values_zero]).T)

plt.rcParams["font.size"] = 14

# Plot energy and reference value
plt.figure(figsize=(12, 6))
plt.plot(counts_real, values_real, label='Default')
plt.plot(counts_zero, values_zero, label='Res=0,Opt=0')
plt.axhline(y=-84.206, color="tab:red", ls="--", label="Target")

plt.legend(loc="best")
plt.xlabel("Iteration")
plt.ylabel("Energy [H]")
plt.title("VQE energy")
plt.show()

