# Imports needed for basic no noise VQE simulation
import numpy as np
import pylab
import copy
from qiskit import BasicAer
from qiskit.aqua import aqua_globals, QuantumInstance
from qiskit.aqua.algorithms import NumPyMinimumEigensolver, VQE
from qiskit.aqua.components.optimizers import SLSQP, SPSA
from qiskit.chemistry.components.initial_states import HartreeFock
from qiskit.chemistry.components.variational_forms import UCCSD
from qiskit.chemistry.drivers import PySCFDriver
from qiskit.chemistry.core import Hamiltonian, QubitMappingType
from qiskit.circuit.library import EfficientSU2


# Some deprecated package issue; one warning is enough. 
import warnings
warnings.filterwarnings(action='once')

# Defining a noise model! 

from qiskit import Aer
from qiskit.providers.aer.noise import NoiseModel
from qiskit.providers.aer import QasmSimulator
from qiskit.test.mock import FakeVigo
from qiskit.ignis.mitigation.measurement import CompleteMeasFitter

device_backend = FakeVigo()

# NOTE the key difference that the noisy simulations need to use 
# 'qasm_simulator' rather than 'statevector_simulator'

backend = Aer.get_backend('qasm_simulator')

device = QasmSimulator.from_backend(device_backend)
coupling_map = device.configuration().coupling_map
noise_model = NoiseModel.from_backend(device)
basis_gates = noise_model.basis_gates

print(noise_model)
print()

molecule = 'H .0 .0 -{0}; Li .0 .0 {0}' # Define molecule; could be H2 or LiH, etc.
distances = np.arange(0.5, 3.75, 0.25) # Distances in Angstrom
vqe_energies = [] 
hf_energies = [] # Hartree-Fock energies
exact_energies = [] # Exact energies from numerical diagonalization

for i, d in enumerate(distances):
    print('step', i)
    
    # initialize driver
    driver = PySCFDriver(molecule.format(d/2), basis='sto3g')
    qmolecule = driver.run()
    operator = Hamiltonian(qubit_mapping=QubitMappingType.PARITY, 
                           two_qubit_reduction=True, freeze_core=True,
                           orbital_reduction=[-3, -2])
    
    qubit_op, aux_ops = operator.run(qmolecule)
    
    # exact diagonalization
    exact_result = NumPyMinimumEigensolver(qubit_op, aux_operators=aux_ops).run()
    exact_result = operator.process_algorithm_result(exact_result)
    
    # VQE OPTION 1: SLSQP optimizer + UCCSD var_form; NO NOISE.
    optimizer = SLSQP(maxiter=100)
    initial_state = HartreeFock(operator.molecule_info['num_orbitals'],
                                operator.molecule_info['num_particles'],
                                qubit_mapping=operator._qubit_mapping,
                                two_qubit_reduction=operator._two_qubit_reduction)
                                # Chemical approxns to make this problem tractable.
    
    
    var_form = UCCSD(num_orbitals=operator.molecule_info['num_orbitals'],
                     num_particles=operator.molecule_info['num_particles'],
                     initial_state=initial_state,
                     qubit_mapping=operator._qubit_mapping,
                     two_qubit_reduction=operator._two_qubit_reduction)
        
    algo = VQE(qubit_op, var_form, optimizer, aux_operators=aux_ops)
    
    quantum_instance = QuantumInstance(backend=backend, 
                                       shots=8192, 
                                       noise_model=None)
    
    ###########################################################################
    
    # VQE OPTION 2: SLSA optimizer + SU(2) [i.e. RYRZ] var_form; w/ NOISE + Mitigation
#     optimizer = SLSA(maxiter=100)
#     initial_state = HartreeFock(operator.molecule_info['num_orbitals'],
#                                 operator.molecule_info['num_particles'],
#                                 qubit_mapping=operator._qubit_mapping,
#                                 two_qubit_reduction=operator._two_qubit_reduction)
#                                 # Chemical approxns to make this problem tractable.
    
    
#     var_form = EfficientSU2(qubit_op.num_qubits, entanglement="linear")
        
#     algo = VQE(qubit_op, var_form, optimizer, aux_operators=aux_ops)
    
#     quantum_instance = QuantumInstance(backend=backend, 
#                                    shots=8192, 
#                                    noise_model=noise_model, 
#                                    coupling_map=coupling_map,
#                                    measurement_error_mitigation_cls=CompleteMeasFitter,
#                                    cals_matrix_refresh_period=30)
                                   
    # Nb: The last two inputs incorporate mitigation. Remove to get unmitigated result.
    
    ###########################################################################
    
    vqe_result = algo.run(quantum_instance)
    vqe_result = operator.process_algorithm_result(vqe_result)
    
    exact_energies.append(exact_result.energy)
    vqe_energies.append(vqe_result.energy)
    hf_energies.append(vqe_result.hartree_fock_energy)

pylab.plot(distances, hf_energies, label='Hartree Fock')
pylab.plot(distances, vqe_energies, 'o', label='VQE')
pylab.plot(distances, exact_energies, 'x', label='Exact')

pylab.xlabel('Interatomic distance')
pylab.ylabel('Energy')
pylab.title('LiH Ground State Energy')
pylab.legend(loc='upper right')

# Uncomment to save plot! 
#pylab.savefig('vqe.png', dpi=300)

pylab.plot(distances, hf_energies, label='Hartree Fock')
pylab.plot(distances, vqe_energies, 'o', label='VQE')
pylab.plot(distances, exact_energies, 'x', label='Exact')

pylab.xlabel('Interatomic distance')
pylab.ylabel('Energy')
pylab.title('LiH Ground State Energy')
pylab.legend(loc='upper right')

# Uncomment to save plot! 
#pylab.savefig('vqe.png', dpi=300)





from qiskit import execute, QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.quantum_info import Kraus, SuperOp
from qiskit.providers.aer import QasmSimulator
from qiskit.tools.visualization import plot_histogram

# Import from Qiskit Aer noise module
from qiskit.providers.aer.noise import NoiseModel
from qiskit.providers.aer.noise import QuantumError, ReadoutError
from qiskit.providers.aer.noise import pauli_error
from qiskit.providers.aer.noise import depolarizing_error
from qiskit.providers.aer.noise import thermal_relaxation_error


# T1 and T2 values for qubits 0-3
T1s = np.random.normal(50e3, 10e3, 4) # Sampled from normal distribution mean 50 microsec
T2s = np.random.normal(70e3, 10e3, 4)  # Sampled from normal distribution mean 50 microsec

# Truncate random T2s <= T1s
T2s = np.array([min(T2s[j], 2 * T1s[j]) for j in range(4)])

# Instruction times (in nanoseconds)
time_u1 = 0   # virtual gate
time_u2 = 50  # (single X90 pulse)
time_u3 = 100 # (two X90 pulses)
time_cx = 300
time_reset = 1000  # 1 microsecond
time_measure = 1000 # 1 microsecond

# QuantumError objects
errors_reset = [thermal_relaxation_error(t1, t2, time_reset)
                for t1, t2 in zip(T1s, T2s)]
errors_measure = [thermal_relaxation_error(t1, t2, time_measure)
                  for t1, t2 in zip(T1s, T2s)]
errors_u1  = [thermal_relaxation_error(t1, t2, time_u1)
              for t1, t2 in zip(T1s, T2s)]
errors_u2  = [thermal_relaxation_error(t1, t2, time_u2)
              for t1, t2 in zip(T1s, T2s)]
errors_u3  = [thermal_relaxation_error(t1, t2, time_u3)
              for t1, t2 in zip(T1s, T2s)]
errors_cx = [[thermal_relaxation_error(t1a, t2a, time_cx).expand(
             thermal_relaxation_error(t1b, t2b, time_cx))
              for t1a, t2a in zip(T1s, T2s)]
               for t1b, t2b in zip(T1s, T2s)]

# Add errors to noise model
noise_thermal = NoiseModel()
for j in range(4):
    noise_thermal.add_quantum_error(errors_reset[j], "reset", [j])
    noise_thermal.add_quantum_error(errors_measure[j], "measure", [j])
    noise_thermal.add_quantum_error(errors_u1[j], "u1", [j])
    noise_thermal.add_quantum_error(errors_u2[j], "u2", [j])
    noise_thermal.add_quantum_error(errors_u3[j], "u3", [j])
    for k in range(4):
        noise_thermal.add_quantum_error(errors_cx[j][k], "cx", [j, k])

print(noise_thermal)