#Create the random probability Vector 
import numpy as np
np.random.seed(999999)
target_distr = np.random.rand(2)
#print (target_distr)
# We now convert the random vector into a valid probability vector
target_distr /= sum(target_distr)
print (target_distr) #This is our original state, we want our result close to this



#This block of code is to make a function that takes the parameters of our single U3 variational form as argument 
#and returns the corresponding quantum circuit
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister
def get_var_form(params):
    qr = QuantumRegister(1, name="q")
    cr = ClassicalRegister(1, name='c')
    qc = QuantumCircuit(qr, cr)
   
    qc.u3(params[0], params[1], params[2], qr[0])
    qc.measure(qr, cr[0])
    
    return qc

#Now specify the objective function which takes as input a list of the variational form's parameter and returns 
#the cost associated with those parameters(how much the final state is differ from intial state)
from qiskit import Aer, execute
backend = Aer.get_backend("qasm_simulator")
NUM_SHOTS = 10000

def get_probability_distribution(counts):
    output_distr = [v / NUM_SHOTS for v in counts.values()]
    #print("Counts",counts)
    print ("Count Values",counts.values())
    #print("Output_distr",output_distr)
    #print("Target_distr",target_distr)
    if len(output_distr) == 1:
        output_distr.append(0)
    return output_distr

def objective_function(params):
    # Obtain a quantum circuit instance from the paramters
    qc = get_var_form(params)
    
    # Execute the quantum circuit to obtain the probability distribution associated with the current parameters
    result = execute(qc, backend, shots=NUM_SHOTS).result()
    print("Results",result.get_counts(qc))
    
    # Obtain the counts for each measured state, and convert those counts into a probability vector
    output_distr = get_probability_distribution(result.get_counts(qc))
    
    # Calculate the cost as the distance between the output distribution and the target distribution
    cost = sum([np.abs(output_distr[i] - target_distr[i]) for i in range(2)])
    #display(qc.draw('mpl'))
    print('Input_distr:',target_distr)
    print('Output_distr:',output_distr)
    print('Cost',cost)
    print('-------------')
   
    return cost #This step is similar as to get the minimum energy state using variational principle

#Finally COBYLA optimizer is used and run the algorithm. 
from qiskit.aqua.components.optimizers import COBYLA 
# Constrained Optimization by Linear Approximation optimizer (COBYLA)

# Initialize the COBYLA optimizer

optimizer = COBYLA(maxiter=500, disp=True,tol=0.0001)

# Create the initial parameters (noting that our single qubit variational form has 3 parameters)
params = np.random.rand(3)
print ('Input parameters:->',params)
ret = optimizer.optimize(num_vars=3, objective_function=objective_function, initial_point=params)

#print("ret==>",ret)
# Obtain the output distribution using the final parameters

qc = get_var_form(ret[0])
counts = execute(qc, backend, shots=NUM_SHOTS).result().get_counts(qc)
output_distr = get_probability_distribution(counts)
print("Target Distribution:", target_distr)
print("Obtained Distribution:", output_distr)
print("Output Error (Manhattan Distance):", ret[1])
print("Parameters Found:", ret[0])

from qiskit.aqua.algorithms import VQE, NumPyEigensolver
import matplotlib.pyplot as plt
import numpy as np
from qiskit.chemistry.components.variational_forms import UCCSD
from qiskit.chemistry.components.initial_states import HartreeFock
from qiskit.circuit.library import EfficientSU2
from qiskit.aqua.components.optimizers import COBYLA, SPSA, SLSQP
from qiskit.aqua.operators import Z2Symmetries
from qiskit import IBMQ, BasicAer, Aer
from qiskit.chemistry.drivers import PySCFDriver, UnitsType
from qiskit.chemistry import FermionicOperator
from qiskit import IBMQ
from qiskit.aqua import QuantumInstance
from qiskit.ignis.mitigation.measurement import CompleteMeasFitter
from qiskit.providers.aer.noise import NoiseModel



def get_qubit_op(dist):
    #Step 1: Define a molecule
    #Here, we use LiH in the sto3g basis with the PySCF driver as an example. 
    #The molecule object records the information from the PySCF driver.
    driver = PySCFDriver(atom="Li .0 .0 .0; H .0 .0 " + str(dist), unit=UnitsType.ANGSTROM, 
                         charge=0, spin=0, basis='sto3g')
    #sto3g-> small toy model
    molecule = driver.run()
    #ferOp_nofreeze=FermionicOperator(h1=molecule.one_body_integrals, h2=molecule.two_body_integrals)

    freeze_list = [0]
    remove_list = [-3, -2]
    repulsion_energy = molecule.nuclear_repulsion_energy 
    #under Born Opperhiemer Approximation all the nuclei are stationary 
    #so their contribution to the ground state energy is only nuclear repulsion. 
    #print(repulsion_energy) #in Hatree Unit
    
    num_particles = molecule.num_alpha + molecule.num_beta #number of electrons in the system
    #print("number of particles==>",num_particles) #4
    num_spin_orbitals = molecule.num_orbitals * 2 #number of spin orbital is twice of that space orbitals
    
    
    print("number of spin orbitals",num_spin_orbitals)
    #as spin degeneracy is 2
    remove_list = [x % molecule.num_orbitals for x in remove_list] 
    #% in interpreted language: Reminder=dividend-divisor*floor(quotient)
    #==> R_0 = -3 -6*(-1) = +3 and R_1= -2-6*(-1) = +4 
    
    #print ('freeze_list=>',freeze_list) #[0]
    #print ('Remove list=>',remove_list) #[3,4]
    freeze_list = [x % molecule.num_orbitals for x in freeze_list]
    remove_list = [x - len(freeze_list) for x in remove_list]
    #print ('freeze_list=>',freeze_list) #[0]
    #print ('Remove list=>',remove_list) #[2,3]
    remove_list += [x + molecule.num_orbitals - len(freeze_list)  for x in remove_list]
    freeze_list += [x + molecule.num_orbitals for x in freeze_list]
    print ('Remove list=>',remove_list) #[2,3,7,8]
    print ('freeze_list=>',freeze_list) #[0,6]
    
    
    #Define Fermionic Operator
    ferOp = FermionicOperator(h1=molecule.one_body_integrals, h2=molecule.two_body_integrals)
    ferOp, energy_shift = ferOp.fermion_mode_freezing(freeze_list)# the frozen electron(2 inner electron from Li+) 
    #do not interact with other electrons but contribute by a constant term to the total ground state energy
    num_spin_orbitals -= len(freeze_list) #rescale num_spin_orbiatls after freeze list
    num_particles -= len(freeze_list) 
    print("number of particles==>",num_particles) #print=>2 
    
    ferOp = ferOp.fermion_mode_elimination(remove_list)
    h1=molecule.one_body_integrals
    h2=molecule.two_body_integrals
    #print('one body integral=>')
    #print (ferOp_nofreeze.h1) #intialially it has 12 
    #print('one body integral after freezeing and removing')
    #print(ferOp.h1) # after applying freeze and remove it drop down to 6 
    #print('-------------------------------------------')
    #print('two body integrals=>')
    #print(ferOp.h2)
    #print ('------------------------------')
    num_spin_orbitals -= len(remove_list)
    #print("========================Fermionic Operator After applying freezing and removing list==================")
    #print (ferOp)
    
    #Create a Qubit Operator 
    qubitOp = ferOp.mapping(map_type='parity', threshold=0.00000001)
    #print ("=====Qubit Operator berfore applying parity========")
    
    print(qubitOp) # qubits is 6 which is same as fermionic operator after applying freezing and removing list
    
    print(Z2Symmetries.find_Z2_symmetries(qubitOp))#Show spin and charge symmetries, so two conserved quantities
    #which will further helps to reduce the qubitOp
    qubitOp = Z2Symmetries.two_qubit_reduction(qubitOp, num_particles)
    shift = energy_shift + repulsion_energy 
    #shift is the enegy comming from repulsion energy of two nuclei under Born–Oppenheimer approximation and
    #and freezing two electrons
    #print("-----------------Qubit Operator after appying parity---------------------")
    #print (qubitOp) # give type of operator(paulis) working on number of qubits(6) and number of paulis operator 
    #print(qubitOp.print_details()) 
    return qubitOp, num_particles, num_spin_orbitals, shift






backend = BasicAer.get_backend("statevector_simulator")
#distances=[1.6]
distances = np.arange(0.5, 4.0, 0.1)
exact_energies = []
vqe_energies = []
optimizer = SLSQP(maxiter=5)
#Sequential Least SQuares Programming optimizer (SLSQP)
for dist in distances:
    qubitOp, num_particles, num_spin_orbitals, shift = get_qubit_op(dist) #Returing multiple values in python function
    result = NumPyEigensolver(qubitOp).run() #Just to check with result obtained from VQE
    exact_energies.append(np.real(result.eigenvalues) + shift)
    #Create initial Guess state(Ansatz)
    initial_state = HartreeFock(
        num_spin_orbitals,
        num_particles,
        qubit_mapping='parity'
    ) 
    print(initial_state.bitstr)
    HF_circuit=initial_state.construct_circuit('circuit')
    display(HF_circuit.decompose().draw(output='mpl')) #inital Hatree Fock state
    
    #UCCSD is used as problem has one and two body integral
    var_form = UCCSD(
        num_orbitals=num_spin_orbitals,
        num_particles=num_particles,
        initial_state=initial_state,
        qubit_mapping='parity'
    )
    #print(var_form.single_excitations)
    print("Number of Parameters",var_form.num_parameters)
    #var_circuit=var_form.construct_circuit([2])
    #display(var_circuit.decompose().draw('mpl'))
    vqe = VQE(qubitOp, var_form, optimizer)
    vqe_result = np.real(vqe.run(backend)['eigenvalue'] + shift)
    vqe_energies.append(vqe_result)
    print("Interatomic Distance:", np.round(dist, 2), "VQE Result:", vqe_result, "Exact Energy:", exact_energies[-1])
    
print("All energies have been calculated")

plt.plot(distances, exact_energies, label="Exact Energy")
#plt.plot(distances, vqe_energies, label="VQE Energy")
plt.xlabel('Atomic distance (Angstrom)')
plt.ylabel('Energy')
plt.legend()
plt.show()

init_state=HartreeFock(num_orbitals=4,num_particles=2,qubit_mapping='jordan_wigner')
print(init_state.bitstr) #give fermionic mode on which particle are occupied
HF_circuit=init_state.construct_circuit('circuit')
display(HF_circuit.decompose().draw('mpl'))
UCCSD_var_form=UCCSD(num_orbitals=4,num_particles=2,qubit_mapping='jordan_wigner',excitation_type='s',
                    method_singles='beta',initial_state=init_state,two_qubit_reduction=False,reps=1)
var_circuit=UCCSD_var_form.construct_circuit([1])
var_circuit.decompose().draw(output='mpl')
#print(UCCSD_var_form.single_excitations)

