# Create the matrix that diagonalizes A
diagonalizer = qt.Qobj(np.array([eigenvecs[i].full().T.flatten()
                                 for i in range(len(eigenvals))]))
b = diagonalizer * b
A = diagonalizer.dag() * A * diagonalizer

"""
Explanation: STEP 1: STATE PREPARATION OF b
[HHL: page 2, column 2, center]
Write b in the eigenbasis of A
End of explanation
"""


T = prec
t0 = κ / ϵ    # It should be O(κ/ϵ), whatever that means
ψ0 = qt.Qobj([[math.sqrt(2 / T) * math.sin(π * (τ + 0.5) / T)] 
                for τ in range(T)])

# Order is b, τ, and then ancilla
evo = qt.tensor(qt.identity(A.shape[0]), qt.ket2dm(qt.basis(T, 0))) 

for τ in range(1, T):
    evo += qt.tensor((1j * A * τ * t0 / T).expm(), qt.ket2dm(qt.basis(T, τ)))

ψev = evo * qt.tensor(b, ψ0)

ftrans = qt.tensor(qt.identity(b.shape[0]), qft(T))

ψfourier = ftrans * ψev    # This is Eq. 3 in HHL

"""
Explanation: STEP 2: QUANTUM SIMULATION AND QUANTUM PHASE ESTIMATION
[HHL: page 2, column 2, bottom]
End of explanation
"""


# w = (ψfourier[:T] / b[0]).argmax()
# prj = qt.ket2dm(qt.basis(T, w))

# ψfourier = qt.tensor(qt.identity(b.shape[0]), prj) * ψfourier

"""
Explanation: STEP 2-1: Making true the assumption of perfect phase estimation
End of explanation
"""


total_state = qt.tensor(ψfourier, qt.basis(2, 0))   # Add ancilla for swapping

C = 1 / κ    # Constant, should be O(1/κ)

# Do conditional rotation only on τ and ancilla
rotation = qt.tensor(qt.ket2dm(qt.basis(T, 0)), qt.identity(2)) 

for τ in range(1, T):
    rotation += qt.tensor(qt.ket2dm(qt.basis(T, τ)), rot(τ, t0, C))
    
final_state = qt.tensor(qt.identity(b.shape[0]), rotation) * total_state

"""
Explanation: STEP 3-1: Conditional Rotation of Ancilla
End of explanation
"""


projector = qt.tensor(qt.identity(b.shape[0]), qt.identity(T), qt.ket2dm(qt.basis(2, 1)))
postsel = projector * final_state
prob1 = qt.expect(projector, final_state)

# Trace out ancilla and τ registers, leaving only the b register
finalstate = qt.ket2dm(postsel).ptrace([0]) / prob1
finalstate.eigenenergies()

"""
Explanation: STEP 3-2: Post-selection on $\left|1\right\rangle$ on the ancilla register
We perform the post-selection by projecting onto the desired ancilla state and later tracing out the ancilla and the $\left|\psi_0\right\rangle$ registers.
End of explanation
"""


fsevls, fsevcs = finalstate.eigenstates()
x = math.sqrt(fsevls.max()) * fsevcs[fsevls.argmax()]
x

"""
Explanation: $\left|finalstate\right\rangle$ is essentially a pure state (it should be if all the process was perfect), so now we just isolate that main part, that is our $\left|x\right\rangle$ (after a transformation if our original A was not Hermitian) in the basis that diagonalizes $A$.
End of explanation
"""


