import prepare_MPS as mps
import numpy as np
from qiskit import BasicAer, execute

#Create Random MPS with size 4, bond dimension 4 and physical dimension 2 (qubits)
N=4
d=2
chi=4
phi_final=np.random.rand(chi)
phi_initial=np.random.rand(chi)
A=mps.create_random_tensors(N,chi,d)

#Create the circuit. The 'reg' register corresponds to the 'MPS' register in the picture above
qc, reg = mps.MPS_to_circuit(A, phi_initial, phi_final)

#Run the circuit on the statevector simulator
backend = BasicAer.get_backend("statevector_simulator")
job = execute(qc, backend)
result = job.result()
psi_out=result.get_statevector()

#Contract out the ancilla with the known state
psi_out=psi_out.reshape(d**N,chi)
exp=psi_out.dot(phi_final)

#Prepare the MPS classically
thr,_=mps.create_statevector(A,phi_initial,phi_final,qiskit_ordering=True)

#Compare the resulting vectors (fixing phase and normalization)
exp=mps.normalize(mps.extract_phase(exp))
thr=mps.normalize(mps.extract_phase(thr))

print("The MPS is \n {}".format(thr))
print("The statevector produced by the circuit is \n {}".format(exp))



from qiskit import ClassicalRegister

N=5
chi=2

#The following is the standard representation of a GHZ state in terms of MPS
phi_initial=np.array([1,1])
phi_final=np.array([1,1])
T=np.zeros((d,chi,chi))
T[0,0,0]=1
T[1,1,1]=1
A=[]
for _ in range(N):
    A.append(T)
    

#Create the circuit, store the relevant wavefunction is register 'reg' and measure
qc, reg = mps.MPS_to_circuit(A, phi_initial, phi_final)
creg=ClassicalRegister(N)
qc.add_register(creg)
qc.measure(reg,creg)

#Run on a simulator
backend = BasicAer.get_backend("qasm_simulator")
job = execute(qc, backend)
result = job.result()
counts = result.get_counts(qc)
print("\nTotal counts are:",counts)