# !pip install qiskit
# !pip install pylatexenc

from qiskit import *
import numpy as np
from matplotlib import pyplot as plt
from qiskit.visualization import plot_histogram
from qiskit.circuit.library import QFT
import pylatexenc

#Number of qubits
n=6

#Circuit 
phi=0.1 #time step
iterations=30 #iterations

#Potential
V=1.75 

#Spatial Grid
N=2**n #grid points
xMin=-4
xMax=4
x = np.linspace(xMin,xMax, num=N)

#Wavefunction
mu=-2.25
psi= (1/(np.sqrt(2*np.pi)*2*0.4)*np.exp(-(0.5*(x-mu)**2)/(2*0.016)))*np.exp(-300j*x)
psi/= np.sqrt(np.sum(np.abs(psi**2)))
plt.plot(x,psi)

qubits=QuantumRegister(n+1)
splitOp=QuantumCircuit(qubits,name="splitOp")
'''potential operator''' 
splitOp.ccx(3,4,0)
splitOp.cp(V*phi,0,5)
splitOp.ccx(3,4,0)
splitOp.barrier()
'''momentum operator'''
splitOp.append(QFT(n),qubits[1:])
for i in range(n):
  splitOp.p(phi/2**(n-3+i), n-i)
for i in range(n):
  for j in range(i+1,n):
    splitOp.cx(n-j,0)
    splitOp.cx(n-i,0)
    splitOp.p(phi*(2**(2-i-j)),0)
    splitOp.cx(n-i,0)
    splitOp.cx(n-j,0)
splitOp.append(QFT(n,inverse=True),qubits[1:])
splitOp.barrier()
splitOp.draw('mpl')

q=QuantumRegister(n+1)
c=ClassicalRegister(n)
circuit=QuantumCircuit(q,c)
'''initialization'''
circuit.initialize(psi,q[1:])
circuit.barrier()
'''time steps'''
for i in range(iterations):
  circuit.append(splitOp,q[:])
'''measurement'''
for i in range(1,n+1):
  circuit.measure(i,i-1)
#circuit.draw('mpl')
answer = execute(circuit, backend=BasicAer.get_backend('qasm_simulator'), shots=5000).result().get_counts()

for a in range(N):
    s = bin(a)[2:]
    while len(s)!=6:
        s='0'+s
    if s not in answer.keys():
        answer[s] = 0
plot_histogram(answer, figsize = [20,10])

new_psi=np.zeros(N)
count=0
for key in sorted(answer.keys()):
  new_psi[count]=answer[key]
  count+=1
new_psi=new_psi/np.sqrt(np.sum(psi**2))
barrier=np.zeros(N)
barrier[14:16]=20000
print(np.sum(np.abs(psi**2)))
plt.plot(x,new_psi)
plt.plot(x,barrier)
plt.show()

