from qiskit import Aer
from qiskit.circuit.library import QFT
from qiskit.aqua import QuantumInstance, aqua_globals
from qiskit.quantum_info import state_fidelity
from qiskit.aqua.algorithms import HHL, NumPyLSsolver
from qiskit.aqua.components.eigs import EigsQPE
from qiskit.aqua.components.reciprocals import LookupRotation
from qiskit.aqua.operators import MatrixOperator
from qiskit.aqua.components.initial_states import Custom
import numpy as np

def create_eigs(matrix, num_ancillae, num_time_slices, negative_evals):
    """QPE for the eigenvalues estimation.

    Parameters
    -------------------------------------------------------------------
    matrix(np.array): Unitary matrix for the QPE.
    num_ancillae(int): Number of ancillas.
    num_time_slices(float): An optional evolution time which should scale the 
                            eigenvalue onto the range (0,1] (or (−0.5,0.5] for                                  
                            negative eigenvalues).If None is internally computed.
    negative_evals(Boolean): Set True to indicate negative eigenvalues need to be handled

    """
    
    ne_qfts = [None, None]
    
    if negative_evals:
        num_ancillae += 1
        ne_qfts = [QFT(num_ancillae - 1), QFT(num_ancillae - 1).inverse()]

    return EigsQPE(MatrixOperator(matrix=matrix),
                   QFT(num_ancillae).inverse(),
                   num_time_slices=num_time_slices,
                   num_ancillae=num_ancillae,
                   expansion_mode='suzuki',
                   expansion_order=2,
                   evo_time=None,
                   negative_evals=negative_evals,
                   ne_qfts=ne_qfts)

def fidelity(hhl, ref):
    """ Helper function for the fidelity comparing 
    the solution from the HHL and the classical solution.
    """

    solution_hhl_normed = hhl / np.linalg.norm(hhl)
    solution_ref_normed = ref / np.linalg.norm(ref)
    fidelity = state_fidelity(solution_hhl_normed, solution_ref_normed)        
    return fidelity    

A = [[1, 1/3], [1/3, 4]]
b = [1, 4]

orig_size = len(b)
A, b, truncate_powerdim, truncate_hermitian = HHL.matrix_resize(A, b)

# Initialize eigenvalue finding module
eigs = create_eigs(matrix=A, num_ancillae=3, 
                   num_time_slices=50, 
                   negative_evals=False)

num_q, num_a = eigs.get_register_sizes()

# Initialize initial state module
init_state = Custom(num_q, state_vector=b)

# Initialize reciprocal rotation module
reci = LookupRotation(negative_evals=eigs._negative_evals, 
                      evo_time=eigs._evo_time,
                      scale=0.5)

algo = HHL(matrix=A, vector=b, truncate_powerdim=truncate_powerdim, 
           truncate_hermitian=truncate_hermitian, eigs=eigs,
           init_state=init_state, reciprocal=reci, 
           num_q= num_q, num_a=num_a, orig_size=orig_size)

result = algo.run(QuantumInstance(Aer.get_backend('statevector_simulator')))

print(f"solution: {np.round(result['solution'], 5)}")

result_ref = NumPyLSsolver(A, b).run()
print(f"classical solution: {np.round(result_ref['solution'], 5)} ")

print(f"probability: {result['probability_result']}")

fidelity(result['solution'], result_ref['solution'])

orig_size = len(b)
A, b, truncate_powerdim, truncate_hermitian = HHL.matrix_resize(A, b)

# Initialize eigenvalue finding module
eigs = create_eigs(matrix=A, num_ancillae=3, 
                   num_time_slices=50, 
                   negative_evals=False)

num_q, num_a = eigs.get_register_sizes()

# Initialize initial state module
init_state = Custom(num_q, state_vector=b)

# Initialize reciprocal rotation module
reci = LookupRotation(negative_evals=eigs._negative_evals, 
                      evo_time=eigs._evo_time,
                      scale=0.5)

algo = HHL(matrix=A, vector=b, truncate_powerdim=truncate_powerdim, 
           truncate_hermitian=truncate_hermitian, eigs=eigs,
           init_state=init_state, reciprocal=reci, 
           num_q= num_q, num_a=num_a, orig_size=orig_size)

result = algo.run(QuantumInstance(Aer.get_backend('qasm_simulator')))

print(f"solution: {np.round(result['solution'], 5)}")

result_ref = NumPyLSsolver(A, b).run()
print(f"classical solution: {np.round(result_ref['solution'], 5)} ")

print(f"probability: {result['probability_result']}")

fidelity(result['solution'], result_ref['solution'])

from qiskit.tools.jupyter import *
%qiskit_version_table