# numpy, for random number generation
import numpy as np

# Qiskit, for transpiler-related functions, the IBMQ provider, and the Aer simulator
from qiskit import IBMQ, Aer, QuantumRegister
from qiskit.transpiler import transpile, transpile_dag, PassManager
from qiskit.converters import circuit_to_dag, dag_to_circuit

# pytket, for optimization
import pytket
from pytket.qiskit import TketPass

# Qiskit Aqua, for chemistry
from qiskit_chemistry.drivers import PySCFDriver, UnitsType
from qiskit_chemistry import FermionicOperator
from qiskit_chemistry.aqua_extensions.components.initial_states import HartreeFock
from qiskit_chemistry.aqua_extensions.components.variational_forms import UCCSD

IBMQ.load_accounts()

# Choose a particular bond length
# NOTE: Units are in Angstroms

bond_length = 0.7

# Set up molecule

# base_molecule_str = 'Li .0 .0 .0; H .0 .0 {}'
base_molecule_str = 'H .0 .0 .0; H .0 .0 {}'

# Specify other molecular properties
charge = 0
spin = 0
basis = 'sto3g'

# Using driver to get fermionic Hamiltonian
# PySCF example

driver =  PySCFDriver(atom=base_molecule_str.format(bond_length),
                      unit=UnitsType.ANGSTROM,
                      charge=charge,
                      spin=spin,
                      basis=basis)

molecule = driver.run()

print("Molecular repulsion energy: ", molecule.nuclear_repulsion_energy)

n_qubits = molecule.one_body_integrals.shape[0]
n_electrons = molecule.num_alpha + molecule.num_beta - molecule.molecular_charge

# get fermionic operator and mapping to qubit operator
ferOp = FermionicOperator(h1=molecule.one_body_integrals, h2=molecule.two_body_integrals)

qubitOp = ferOp.mapping(map_type='JORDAN_WIGNER', threshold=0.00000001)
qubitOp.chop(10**-10)

# Instantiate the initial state as a Hartree-Fock state
initial_hf = HartreeFock(num_qubits=n_qubits, num_orbitals=n_qubits, 
                    qubit_mapping='jordan_wigner', two_qubit_reduction=False, num_particles= n_electrons)

# Create the variational form
var_form = UCCSD(num_qubits=n_qubits, num_orbitals=n_qubits, 
                num_particles=n_electrons, depth=1, initial_state=initial_hf, qubit_mapping='jordan_wigner')

# How many qubits do we need?
print('Number of qubits: {0}'.format(n_qubits))

# Query the variational form for the number of parameters, and the parameter bounds.
var_form.num_parameters, var_form.parameter_bounds[0]

# Instantiate a concrete instance of the VQE ansatz by setting all the parameters to the
# arbitrarily-chosen value of 0.
var_circ = var_form.construct_circuit(np.zeros(var_form.num_parameters))

# Use Terra to convert the circuit to its directed, acyclic graph (DAG) representation.
var_circ_dag = circuit_to_dag(var_circ)

# The .properties() method of the DAG to get circuit properties.
var_circ_dag.properties()

# Grab an Aer backend
aer_backend = Aer.get_backend('qasm_simulator')

# Choose a random set of parameters
seed = 0
np.random.seed(seed)
params = np.random.uniform(low=-3.1, high=3.1, size=var_form.num_parameters)

# Construct a random instance of the variational circuit
var_circuit = var_form.construct_circuit(params)

# Turn the circuit into a DAG
var_dag = circuit_to_dag(var_circuit)

var_dag.properties()

# Create a Terra PassManager object
tk_pass_manager = PassManager()

# Set up the TketPass
tk_pass = TketPass(aer_backend)

# Add the TketPass to the PassManager
tk_pass_manager.append(tk_pass)

var_dag_transpiled = transpile_dag(var_dag, pass_manager=tk_pass_manager)

var_dag_transpiled.properties()

# Grab only backends that have at least 12 qubits
IBMQ.backends(filters=lambda x: x.configuration().n_qubits >= 12)

real_backend = IBMQ.get_backend('ibmq_16_melbourne')

# Create a Terra PassManager object
tk_pass_manager = PassManager()

# Set up the TketPass
tk_pass = TketPass(real_backend)

# Add the TketPass to the PassManager
tk_pass_manager.append(tk_pass)

blank_qubits = QuantumRegister(len(real_backend.properties().qubits) - var_dag.width())
var_dag.add_qreg(blank_qubits)
var_dag_transpiled = transpile_dag(var_dag, pass_manager=tk_pass_manager)

var_dag_transpiled.properties()

transpile_dag(var_dag, coupling_map=real_backend.configuration().coupling_map).properties()

# Code imports

# From Aqua, we need 
from qiskit_aqua import QuantumInstance

from qiskit_aqua.algorithms.adaptive import VQE
from qiskit_aqua.components.optimizers import L_BFGS_B

# From pytket, we need QSE functions

from pytket.chemistry import QseMatrices, QSE

backend = Aer.get_backend('statevector_simulator')

pass_manager = PassManager()
tk_pass = TketPass(backend)
pass_manager.append(tk_pass)

quantum_instance = QuantumInstance(backend, pass_manager=pass_manager)

# Temporary code for Aer on Macbook
import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'

# Set initial values of parameters
number_amplitudes = len(var_form._single_excitations)+ len(var_form._double_excitations)

amplitudes_0 = []
for i in range(number_amplitudes):
    amplitudes_0.append(0.00001)

optimizer = L_BFGS_B()
optimizer.set_options(maxfun=1000, factr=10, iprint=10)

# setup VQE with operator, variation form, and optimzer
vqe_algorithm = VQE(operator=qubitOp, operator_mode='matrix', 
                    var_form=var_form, optimizer=optimizer, initial_point=amplitudes_0)

results = vqe_algorithm.run(quantum_instance)

eigval = results['eigvals'][0]
gs_energy = eigval.real + molecule.nuclear_repulsion_energy

print("GS Minimum value: {}".format(gs_energy))
print("GS Parameters: {}".format(results['opt_params']))

# store ground state amplitudes for subsequent steps
opti_amplitudes = results['opt_params']

qubitOp = ferOp.mapping(map_type='JORDAN_WIGNER', threshold=0.00000001)
n_qubits = qubitOp.num_qubits
qubitOp.chop(10**-10)

# Use matrix term helper class
matrix_terms = QseMatrices(qubitOp, n_qubits)

# Instantiate an instance of the QSE algorithm
qse_algorithm = QSE(matrix_terms, 'matrix', var_form, opt_init_point=opti_amplitudes)

# Run the algorithm
energies = qse_algorithm.run(quantum_instance)['eigvals']

# The excited state energies are the energies from above,
# plus the nuclear repulsion energy.
print("Excited State Energies: ", energies+molecule.nuclear_repulsion_energy)



