import numpy as np
from qiskit.aqua.algorithms import VQE, ExactEigensolver
from qiskit.chemistry.aqua_extensions.components.variational_forms import UCCSD
from qiskit.aqua.components.variational_forms import TASP
from qiskit.chemistry.aqua_extensions.components.initial_states import HartreeFock
from qiskit.aqua.components.optimizers import COBYLA
from qiskit import BasicAer
from qiskit.chemistry.drivers import PySCFDriver, UnitsType
from qiskit.chemistry import FermionicOperator
from qiskit.aqua import QuantumInstance


#Classical part: set up molecular data, map to WeightedPauliOperator and find exact solution
dist=0.735
driver = PySCFDriver(atom="H .0 .0 .0; H .0 .0 " + str(dist), unit=UnitsType.ANGSTROM, 
                     charge=0, spin=0, basis='sto3g')
molecule = driver.run()
repulsion_energy = molecule.nuclear_repulsion_energy
num_spin_orbitals=molecule.num_orbitals*2

map_type='jordan_wigner'

ferOp = FermionicOperator(h1=molecule.one_body_integrals, h2=molecule.two_body_integrals)
qubitOp = ferOp.mapping(map_type=map_type, threshold=0.00000001)

exact_solution = ExactEigensolver(qubitOp).run()['energy']
exact_solution=exact_solution+repulsion_energy


#Set up initial state and UCCSD variational form with standard parameters
HF_state=HartreeFock(qubitOp.num_qubits, num_spin_orbitals, num_particles, map_type)
varform_UCCSD = UCCSD(qubitOp.num_qubits, depth=1, num_orbitals=num_spin_orbitals, num_particles=num_particles,
                 active_occupied=None, active_unoccupied=None, initial_state=HF_state,
                 qubit_mapping='jordan_wigner', two_qubit_reduction=False, num_time_slices=1,
                 shallow_circuit_concat=True, z2_symmetries=None)




#Run VQE using UCCSD variational form and compare with exact solution
print('Molecule H2 at bond length ' + str(dist))
print('Running VQE with UCCSD. This may take a few minutes. Some intermediary energies while you wait:')
opt_iter=15
optimizer = COBYLA(maxiter=opt_iter)
n_shots=1000
instance=QuantumInstance(backend=BasicAer.get_backend("qasm_simulator"), shots=n_shots)
vqe_instance = VQE(qubitOp, varform_UCCSD, optimizer=optimizer)
result=vqe_instance.run(instance)
energy=result['energy']+repulsion_energy 
print('Optimal energy is ' + str(energy) +'. Exact result is ' + str(exact_solution))



#Split the Hamiltonian into the abovementioned pieces
diag_one_body=np.diag(np.diagonal(molecule.one_body_integrals))
offdiag_one_body=molecule.one_body_integrals-diag_one_body

diag_two_body=np.zeros((num_spin_orbitals,num_spin_orbitals,num_spin_orbitals,num_spin_orbitals))
hop_two_body=np.zeros((num_spin_orbitals,num_spin_orbitals,num_spin_orbitals,num_spin_orbitals))
ex_two_body=np.zeros((num_spin_orbitals,num_spin_orbitals,num_spin_orbitals,num_spin_orbitals))

for i1 in range(num_spin_orbitals):
    for i2 in range(num_spin_orbitals):
        for i3 in range(num_spin_orbitals):
            for i4 in range(num_spin_orbitals):
                if i2==i3:
                    if i1==i4:
                        diag_two_body[i1,i2,i3,i4]=molecule.two_body_integrals[i1,i2,i3,i4]
                    else:
                        hop_two_body[i1,i2,i3,i4]=molecule.two_body_integrals[i1,i2,i3,i4]
                else:
                    ex_two_body[i1,i2,i3,i4]=molecule.two_body_integrals[i1,i2,i3,i4]
                    
H_diag=FermionicOperator(h1=diag_one_body,h2=diag_two_body)
H_hop=FermionicOperator(h1=offdiag_one_body,h2=hop_two_body)
H_ex=FermionicOperator(h1=np.zeros((num_spin_orbitals,num_spin_orbitals)),h2=ex_two_body)
H_diag=H_diag.mapping(map_type=map_type, threshold=0.00000001)
H_hop=H_hop.mapping(map_type=map_type, threshold=0.00000001)
H_ex=H_ex.mapping(map_type=map_type, threshold=0.00000001)
h_list=[H_ex, H_hop, H_diag]


#Define TASP variational form (the TASP class file is attached to original comment)
S=1
varform_TASP=TASP(num_spin_orbitals, depth=S, h_list=h_list, initial_state=HF_state)

print('Running VQE with TASP. This may take a few minutes. Some intermediary energies while you wait:')
vqe_instance = VQE(qubitOp, varform_TASP, optimizer=optimizer)
result=vqe_instance.run(instance)
energy=result['energy']+repulsion_energy 

print('Optimal energy is ' + str(energy) +'. Exact result is ' + str(exact_solution))

n_param_UCCSD=varform_UCCSD.num_parameters
qc_UCCSD=varform_UCCSD.construct_circuit(np.random.rand(n_param_UCCSD))
print('Number of parameters for UCCSD:', n_param_UCCSD)
print('Depth for UCCSD:', qc_UCCSD.depth())
print('Gate counts:', qc_UCCSD.count_ops())

n_param_TASP=varform_TASP.num_parameters
qc_TASP=varform_TASP.construct_circuit(np.random.rand(n_param_TASP))
print('Number of parameters for TASP:', n_param_TASP)
print('Depth for TASP:', qc_TASP.depth())
print('Gate counts:', qc_TASP.count_ops())


