from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, execute, Aer, IBMQ
from qiskit.circuit.library import ExactReciprocal
from qiskit.extensions import HamiltonianGate
from qiskit.quantum_info import random_hermitian
import qiskit_tools as qt
import numpy as np
import matplotlib.pyplot as plt
import matplotlib

matplotlib.rcParams['mathtext.fontset'] = 'stix'
matplotlib.rcParams['font.family'] = 'STIXGeneral'
width=0.75
color='black'
fontsize=28
ticksize=22
figsize=(10,8)

def HHL(circ, qb, qanc, qtarg, A, t=2.*np.pi, scaling=None, wrap=False, inverse=False, phase=True, label='HHL'):
        
    n = len(qb)
    nanc = len(qanc)

    if inverse:
        wrap = True

    if wrap:
        qb = QuantumRegister(n, 'b')
        qanc = QuantumRegister(nanc, 'anc')
        qtarg = QuantumRegister(1, 'targ')
        circ = QuantumCircuit(qb, qanc, qtarg)

    # Define unitary operator given matrix A and evolution time t
    A_gate = HamiltonianGate(A, t)
    
    # Apply phase estimation
    qe_gate = qt.PhaseEst(circ, qb, qanc, A_gate, wrap=True, do_swaps=False, reverse_bits=True)
    circ.append(qe_gate, [*qb, *qanc]);
        
    # Apply ExactReciprocal to rotate the target qubit proportionally to the given scaling and the inverse of the
    # values stored in the computational basis of the ancillary register
    rec_gate = ExactReciprocal(nanc, scaling=scaling, neg_vals=phase).to_gate()
    circ.append(rec_gate, [*qanc[::-1], qtarg]);

    # Apply the inverse phase estimation
    qe_gate_inv = qt.PhaseEst(circ, qb, qanc, A_gate, wrap=True, inverse=True, do_swaps=False, reverse_bits=True)
    circ.append(qe_gate_inv, [*qb, *qanc]);
    
    if wrap:
        circ = circ.to_gate()
        circ.label = label

    if inverse:
        circ = circ.inverse()
        circ.label = label+'\dag'

    return circ

n = 2

A = random_hermitian(2**n).to_matrix()
print(A)

lambda_min = np.min(np.abs(np.linalg.eigvals(A)))
lambda_max = np.max(np.abs(np.linalg.eigvals(A)))
kappa = np.abs(lambda_max/lambda_min)

print('Minimum eigenvalue, maximum eigenvalue and conditional number of A:', lambda_min, lambda_max, kappa)

phase = True

# Choose a success probability of determining eigenvalue estimation
p_suc = 0.99
# Determine the number of ancillary qubits to ensure the given success probability
nanc = n + int(np.ceil(np.log2(2. + 1./(2.*(1-p_suc)))))

# Calculate the new scale of the relative magnitudes of maximum and minimum eigenvalues
rescale = int((2**(nanc-phase)-1)/kappa)
# Make sure this scaling is not vanishingly small
if rescale < 1.e-7:
    rescale = 1
    
# determine the number of integer qubits required
nint = qt.get_nint(rescale)
# Assert that the number of integer qubits is equal to or greater than the number of ancillary qubits assigned
if nint < nanc - phase:
    nint = nanc - phase
    
# We can determine the scaling of the rotation of step 3 (variable c)
scaling = qt.bin_to_dec(qt.my_binary_repr(rescale, nint=nint, n=nint, phase=False), nint=0, phase=False)
t0 = scaling / (lambda_min * (2**phase))
t = 2 * np.pi * t0

qb = QuantumRegister(n, 'b')
circ = QuantumCircuit(qb)

circ.h(qb);

backend = Aer.get_backend('statevector_simulator')
job = execute(circ, backend)
result = job.result()
in_vector = np.asarray(result.get_statevector())

print('The state amplitudes of b are:',in_vector)

targ = np.matmul(np.linalg.inv(A),in_vector)
print('The target state amplitudes of x are:',targ)

qanc = QuantumRegister(nanc, 'anc')
qtarg = QuantumRegister(1, 'targ')

circ.add_register(qanc, qtarg)

circ = HHL(circ, qb, qanc, qtarg, A, t=t, scaling=scaling)

circ.draw('latex')

backend = Aer.get_backend('statevector_simulator')
job = execute(circ, backend)
result = job.result()
state_vector = np.asarray(result.get_statevector())

state_v = np.asarray(state_vector).reshape((2**1,2**nanc,2**n))

# 1, nanc, n
print(state_v.shape)

norm_ = np.sqrt(np.sum(np.abs(state_v[1])**2))
print('The target qubit is in the 1 state with probability:',norm_)

state_v = state_v[1,0]

print(state_v)

state_v = np.array(state_v/np.sqrt(np.sum(np.abs(state_v)**2)))

norm = norm_/lambda_min
out = state_v*norm

print('The amplitudes of the resulting states are:', out)

print('Ratio of real part of output to target:',out.real/targ.real)
print('Ratio of imaginary part of output to target:',out.imag/targ.imag)