#run.py

'''
참고 논문 :Low Complexity Quantum Matrix Inversion A기gorithm for non-Hermitian Matrices

해당 논문에서, non-Hermitian인 경우에 Hermitian으로 바꾸기 위해서 2*2의 행렬을 4*4 행렬로 바꾸었다.
밑의 주어진 코드를 통해 이해해보도록하자.

A = np.array([[2,-1],[1,4]])의 꼴로 주어진 A는 non-Hermitian의 꼴이다.
그러므로, 이를 def main - #Check if Hermitian 파트에서 다음과 같은 코드로 아래와 같은 Hermitian으로 바꾸게 된다.

(코드) A = np.vstack((np.hstack((np.zeros_like(A),A)),np.hstack((A.T, np.zeros_like(A)))))

이에 맞게 나머지 x와 b의 행렬의 형태도 4*1의 형태로 바꾸어주어야 한다.
즉, x = [ 0, 0, x1, x2]의 꼴로, 그리고 b는 [1 1 0 0]의 꼴로 바꾸어야 한다.


(결과)
[[ 0  0  2 -1]
 [ 0  0  1  4]
 [ 2  1  0  0]
 [-1  4  0  0]]

'''

import numpy as np
from scipy.linalg import expm

from qiskit.extensions import UnitaryGate
from qiskit.circuit.add_control import add_control
from qiskit import Aer

import circuit
import hhl
import tools


def main(A,b,backend,shots,t,n_l,delta):

    #Check if Hermitian
    #위에서 언급한 바대로 non-Hermitian matrice인 A에 대해서 Hermitian으로 바꿔주는 과정
    if np.allclose(A,A.T) == False:
        print("Given A matrice is not Hermitian.")
        print("Given Matrices will be transformed into Hermitian formation.")
        A = np.vstack((np.hstack((np.zeros_like(A),A)),np.hstack((A.T, np.zeros_like(A))))) # Hermitian의 꼴로 바꿈
        #A의 shape와 동일한 zero array를 생성하고, A의 왼쪽에 배치, horizontal 방향도 마찬가지.
        b = np.hstack((b,np.zeros_like((np.shape(A)[0]-np.shape(b)[0],1))))

    i = complex(0,1) #complex(real part, imaginary part)
    U = expm(i*A*t) #여기서 A가 행렬로 주어졌기 때문에, 행렬을 exp에 올리기 위해서는 expm이라는 scipy 패키지가 필요함.
    U_gate = UnitaryGate(U) #위에서 구성한 U라는 행렬로써 Unitary gate를 구성할 수 있음. (4*4) 행렬
    CU = add_control(U_gate,1,ctrl_state=None, label="CU") 
    #CU라는 게이트 이름을 label에 저장
    #control 되는 경우의 state를 지정 -> 해당사항 없음
    #두번째 인자는 컨트롤 큐빗의 개수를 지정함.
    n_b = int(np.log2(U.shape[0])) 
    #Ax =b의 꼴이고, b는 4*1의 shape이므로, A의 행의 개수와 동일함. 따라서, U의 행렬의 행의 개수와 동일함.
    #행의 개수에 log2를 취하면 필요한 n_b의 값을 구할 수 있음.

    #각각 HHL 알고리즘이 구현된 방식으로 결과를 얻는다
    My_HHL_result = hhl.My_HHL(CU,b,n_l,n_b,backend,delta,shots,A,details = True,chevyshev = False)
    print("\n")
    qiskit_result = hhl.qiskit_HHL(A,b)
    print("\n")
    classical_result = hhl.classical_HHL(A,b)
    print("\n")

    #For normalized answer
    print("<Un - normalized Case Comparision>")
    print('Qiskit Error : {0}'.format(np.linalg.norm(qiskit_result[1]-classical_result[1])))
    print('My HHL Error : {0}'.format(np.linalg.norm(My_HHL_result[1]-classical_result[1])))
    print("\n")

    print("<Normalized Case Comparision>")
    print('Qiskit Error : {0}'.format(np.linalg.norm(qiskit_result[0]-classical_result[0])))
    print('My HHL Error : {0}'.format(np.linalg.norm(My_HHL_result[0]-classical_result[0])))


if __name__ == "__main__":
    
    #setups
    A = np.array([[2,-1],[1,4]]) #non-Hermitian인 경우의 행렬에 대한 저장
    b = np.array([1,1]) 
    backend = Aer.get_backend('aer_simulator')
    shots = 8192

    t = np.pi*2/16
    n_l = 3 #QPE 상에서 n_ㅣ는 하다마드로 초기화 되는 부분 
    delta = 1/16*(2**(n_l-1))

    main(A,b,backend,shots,t,n_l,delta)

#my_hhl.py

from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
import numpy as np

import circuit as circ
import tools 

def My_HHL(CU,b,n_l,n_b,backend,delta,shots,A,details = True,chevyshev = False):

    #circuit initialization
    #qubit의 개수들을 정의한다.
    n_f = 1
    nb = int(np.log2(b.shape))

    #정의된 qubit의 개수들을 기반으로 양자 레지스터들을 생성한다.
    nl_rg = QuantumRegister(n_l, "l")
    nb_rg = QuantumRegister(n_b, "b")
    na_rg = QuantumRegister(n_l, "a")
    nf_rg = QuantumRegister(n_f, "f")
    #고전 레지스터들을 생성한다.
    cf = ClassicalRegister(n_f, "classical_f")
    cb = ClassicalRegister(n_b, "classical_b")
    #생성된 레지스터들을 기반으로 비어있는 양자회로를 설치한다.
    qc = QuantumCircuit(nf_rg,nl_rg, nb_rg, na_rg, cf, cb)
    #isometry 패키지를 이용해서 normalized 된 b 행렬을 양자회로에 인코딩한다.
    #즉. isometry라는 메서드는 어떤 행렬을 양자회로에 업로드 해주는 역할을 한다.
    qc.isometry(b/np.linalg.norm(b), list(range(nb)), None)
    qc.barrier(nf_rg,nl_rg,nb_rg)

    #details는 양자 회로의 어떠한 instruction들을 세부적으로 표현할 것인지 아닌지에 대한 설정이다.
    if details == True:
        #compose를 이용해서 위에서 생성한 레지스터에 붙이는 경우, instruction 내의 모든 회로가 다 보임.
        #또한, 더 많은 register들이 있는 회로에 적은 register들이 있는 회로를 가져다 붙이기 위해서는 compose에 인자로써 register가 적용될 순서를 명시해주어야함.

        #구성한 circ 패키지 내에 있는 회로들을 compose를 이용해 가져다 붙이는 과정
        #QPE -> Eigenvalue_inversion -> QPE_dagger 순으로 회로를 가져다 붙임 
        qc = qc.compose(circ.QPE(n_l,n_b,CU),nl_rg[:]+nb_rg[:]) 
        qc = qc.compose(circ.Eigenvalue_inversion(n_l,delta,chevyshev),[nl_rg[2]]+[nl_rg[1]]+[nl_rg[0]]+nf_rg[:])
        qc = qc.compose(circ.QPE_dagger(n_l,n_b,CU),nl_rg[:]+nb_rg[:])

    else:
        qc.append(circ.QPE(n_l,n_b,CU),nl_rg[:]+nb_rg[:])
        qc.append(circ.Eigenvalue_inversion(n_l),[nl_rg[2]]+[nl_rg[1]]+[nl_rg[0]]+nf_rg[:])
        qc.append(circ.QPE_dagger(n_l,n_b,CU),nl_rg[:]+nb_rg[:])
        
    qc.barrier(nf_rg[:]+nl_rg[:]+nb_rg[:]) #레지스터의 이름과 그 안에 있는 큐비트들의 순서들을 조합하여 한번에 barrier를 세울 수 있음.
    #측정 실시 파트
    qc.measure(nf_rg,cf)
    qc.measure(nb_rg,cb)
    answer = circ.measurement(qc,n_l,n_b,CU,backend,shots)
    qc.draw(output = 'mpl').savefig('./outputs/qc_HHL')

    #Obtaining Normalized answer
    normalized_result = tools.normalize_vector(answer, n_b)
    
    #Obtaining Real Answer
    
    '''
    즉, normalized된 x값을 집어넣고, Ax = b 꼴에서 좌변을 계산한 후에 b와 몇배 차이인가를 알아내는 과정
    여기서 몇배 차이를 나타내는게 constant임.
    constant는 b와 같은 shape를 갖게 되고, ensemble로 얻어지는 양자회로의 특성상 explicit한 결과가 주어지지 않는다.
    따라서, constant 벡터의 모든 요소들의 평균으로서 몇배 차이인지를 나타낸다.
    '''

    constant = b/(A @ normalized_result)
    constant = (constant[0]+constant[1])/2
    constant = np.mean(constant)

    #결과 출력 파트
    print('<My_HHL>')
    print('Normalized Answer : {0}'.format(normalized_result)) 
    print('Un-normalized Answer : {0}'.format(normalized_result * constant))
    print('Normalize Constant: ' ,constant)

    return [normalized_result,normalized_result * constant]

#classical_hhl.py
from qiskit.algorithms.linear_solvers.numpy_linear_solver import NumPyLinearSolver
import numpy as np

def classical_HHL(A,b):
    
    #NumpyLinearSolver를 이용해서 주어진 행렬방정식을 풀어낸다.
    
    '''
    참고로, 여기선 un-normalized된 해로써 결과값이 나온다는 사실은 자명하다.
    하지만 여기서 나온 해를 normalize를 통해서 my_hhl과 qiskit_hhl에서 내놓은 normalized된 해와의 차이를 분석하기 위해서 
    normalized Answer 또한 작성했다. 
    '''
    
    sol = NumPyLinearSolver().solve(A, b)
    sol_state = sol.state
    norm_state = sol_state/np.linalg.norm(sol_state)

    print('<Classical case using Numpy>')

    if np.shape(b)[0] == 2:
        sol_state = np.pad(sol_state,(2,0))
        norm_state = np.pad(norm_state,(2,0))

    print('Un-normalized Classical Numpy answer : {0}'.format(sol_state,(2,0)))
    print('Normalized Classical Numpy answer : {0}'.format(norm_state,(2,0)))
    
    return [norm_state,sol_state]


#qiskit_hhl.py

from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, Aer
from qiskit.quantum_info import Statevector
from qiskit.algorithms.linear_solvers.hhl import HHL

import numpy as np

def qiskit_HHL(A,b):
    #해당 내용은 https://qiskit.org/textbook/ch-applications/hhl_tutorial.html의 solving 과정을 참고하였음

    backend = Aer.get_backend('aer_simulator')
    #qiskit HHL 코드를 불러옴
    hhl = HHL(quantum_instance=backend)
    #A, b에 대해서 HHL 회로를 구성
    solution = hhl.solve(A, b)
    #만들어진 회로를 그림으로 저장
    solution.state.draw("mpl").savefig("./outputs/HHL_circuit_qiskit.png")
    #연산된 상태를 상태 벡터의 형태로 결과를 얻음
    naive_sv = Statevector(solution.state).data
    #qubit수를 확인
    num_qubit = solution.state.num_qubits
    #상태 벡터에서 필요한 상태만을 골라서 저장함
    naive_full_vector = np.array([naive_sv[2**(num_qubit-1)+i] for i in range(len(b))])
    #실수 부분만 취함
    naive_full_vector = np.real(naive_full_vector)
    #얻어진 벡터를 normalize하여 반환
    normalized_result = naive_full_vector/np.linalg.norm(naive_full_vector)

    #마찬가지로 주어진 해를 통해 constant를 구함.
    constant = b/(A @ normalized_result)
    constant = (constant[0]+constant[1])/2
    constant = np.mean(constant)

    print('<Qiskit_HHL>')
    print('Normalized Qiskit Answer : {0}'.format(normalized_result))
    print('Un-normalized Qiskit Answer : {0}'.format(normalized_result * constant))
    print('Normalize Constant: ' ,constant)

    return [normalized_result,normalized_result * constant]

#circuit_parts.py

from qiskit.circuit.library.arithmetic.exact_reciprocal import ExactReciprocal
from qiskit.circuit.library.arithmetic.piecewise_chebyshev import PiecewiseChebyshev
from qiskit import QuantumCircuit, QuantumRegister,Aer
import numpy as np

def qft_dagger(n_l):

# <qft를 구현하는 과정에 있어서 SWAP gate에 대한 참고사항> - issue 해결에 대한 report 2022-09-02

# SWAP 게이트를 걸어주는 목적은 qiskit은 qubit을 반대방향으로 읽기 때문임.
# 하지만, SWAP 게이트를 위와 같은 이유로 걸어주게 된다고 하면, 
# HHL 알고리즘 상에서 Eigeninversion 단계에서 문제가 생기게 됨. 
# 즉, Eigeninversion에서는 SWAP이 된 상태를 인지하지 못하고 연산을 실시하여 잘못된 연산이 나오게 됨.

    """n-qubit QFTdagger the first n qubits in circ"""
    nl_rg = QuantumRegister(n_l, "l")
    qc = QuantumCircuit(nl_rg)
    # Don't forget the Swaps!
    #QFT의 역연산은 곧 QFT_dagger임을 기억하자.
        
    for j in reversed(range(n_l)):
        qc.h(j)
        for m in reversed(range(j)):
                qc.cp(-np.pi/float(2**(j-m)), m, j)
    qc.name = "QFT†"
    #display(qc.draw(output = 'mpl'))
    return qc
    
def QPE(n_l,n_b,CU):
    #circuit initialization for HHL
    nl_rg = QuantumRegister(n_l, "l")
    nb_rg = QuantumRegister(n_b, "b")
    #QuantumRegister(size=None, name=None, bits=None) 
    qc = QuantumCircuit(nl_rg,nb_rg)
    #display(qc.draw(output = 'mpl'))
    qc.h(nl_rg)
    qc.barrier(nl_rg[:]+nb_rg[:])
    for l in range(n_l):
        for power in range(2**(l)):
            qc.append(CU, [nl_rg[l],nb_rg[0],nb_rg[1]]) 
            #첫번째 큐비트는 2^0번, 이후 2^n꼴로 돌아가게 설계됨.
            #https://qiskit.org/documentation/stubs/qiskit.circuit.ControlledGate.html append의 예제.
            #즉, append의 첫번째 인자는 gate, 두번쨰 인자의 첫번째 요소는 control qubit, 이후 인자의 요소는 target qubit.

    qc.barrier(nl_rg[:]+nb_rg[:])
    qc.append(qft_dagger(n_l),nl_rg[:])
    qc.barrier(nl_rg[:]+nb_rg[:])
    qc.name = "QPE"
    #display(qc.draw(output = 'mpl'))
    return qc
    
def QPE_dagger(n_l,n_b,CU):
    qc = QPE(n_l,n_b,CU)
    qc = qc.inverse()
    #여기서 inverse함수는 모든 rotation 각도까지도 반대로 입력해줌을 확인하였음.
    #QPE dagger는 그저, QPE의 역과정이라고 생각하면 된다. 단, 각도는 반대방향이어야 함.
    #따라서 여기서 inverse함수를 이용하여 QPE의 역과정, 즉, QPE dagger를 실시하였음
    qc.name = 'QPE†'
    return qc

def Eigenvalue_inversion(n_l,delta,chevyshev = False):

    #Chevyshev 근사를 이용한 풀이방법.
    #Qiskit에서 제공한 HHL 알고리즘 상에서는 Chevyshev 근사를 이용한 부분이 있었다.
    #일단 Chevyshev 근사를 이용하는 경우, 기존 Taylor 근사보다 훨씬 빠르게 급수에 수렴한다는 장점이 있다.
    #참고 문헌 : https://freshrimpsushi.github.io/posts/chebyshev-expansion/
    #여기서는 위의 표현한 cos(theta)에 대한 표현을 Chevyshev근사를 이용해 theta값을 알아내겠다는 접근방법이다.
    #하지만, 근사결과가 좋지 못하다는 점 때문에 Chevyshev 근사를 이용하는 대신에 직접 exact한 theta값을 알아내는 ExactReciprocal을 이용하였다.

    if chevyshev == True:
        print("Maybe using Chevyshev approximation is not accurate.")
        #Using Chebychev Approx. (not recommended!)
        nl_rg = QuantumRegister(n_l, "l")
        na_rg = QuantumRegister(n_l, "a")
        nf_rg = QuantumRegister(1, "f")
        qc = QuantumCircuit(nl_rg, na_rg, nf_rg)

        f_x, degree, breakpoints, num_state_qubits = lambda x: np.arcsin(1 / x), 2, [1,2,3,4], n_l
        #degree : 함수를 polynomial로 근사할 떄, 최고차항 정의
        #breakpoints는 구간을 나누는 느낌. : 근사를 할 떄, 다항식을 어떤 구간에서 나눠서 사용할 지
        #l : eigenvalue를 표현
        #f : rotation
        #a : ancila
        pw_approximation = PiecewiseChebyshev(f_x, degree, breakpoints, num_state_qubits)
        pw_approximation._build()
        qc.append(pw_approximation,nl_rg[:]+[nf_rg[0]]+na_rg[:]) #range(nl*2+1))
        qc.name = 'Chevyshev_inversion'
        return qc

    else:
        qc = ExactReciprocal(n_l, delta, neg_vals = True)
        qc.name = 'Reciprocal_inversion'
        return qc

#measurement.py

'''
주어진 회로에 대해 backend를 가지고서 시뮬레이션을 하기 위한 모듈
'''

from qiskit import QuantumCircuit, transpile, assemble
from qiskit.visualization import plot_histogram

def measurement(qc,n_l,n_b,CU,backend,shots):
    
    t = transpile(qc, backend)
    qobj = assemble(t, shots=shots)
    results = backend.run(qobj).result()
    answer = results.get_counts()    
    plot_histogram(answer, title="Output Histogram").savefig('./outputs/output_histogram.png',facecolor='#eeeeee')

    return answer

#normalization.py

import numpy as np

#양자 회로를 통해서 얻어진 결과(dictionary)를 통해서 normalize된 결과 벡터 x를 구하는 함수
def normalize_vector(answer, nb):
    #nb register에서 얻어질 수 있는 상태들을 dictionary의 key의 형태로 만들어 저장한다.
    possible_states = [] #가능한 모든 상태들을 저장하기 위한 list
    for s in range(2**(nb)): 
        possible_states.append(format(s, "b").zfill(nb)) #nb만큼의 자릿수에 대해서 binary의 형태로 모든 경우의 수를 생성.
    #print(answer)
    #flag register를 측정한 결과가 1이 나온 경우에 대해서 nb register의 결과를 순서대로 추가한다.
    available_result = []
    for i in possible_states:
        for key in answer.keys(): #정답의 key 즉, 상태들을 받아옴.
        
            if key[0:2] == i: #얻은 상태들이 가능한 상태에 존재한다면,
                if int(key[-1]) == 1: #그리고 마지막 자릿수, f_register가 1의 값을 갖는다면,
                    available_result.append(answer[key]) #avaliable_result에 f가 1인 상태들이 나온 횟수를 append 하자.
                else:
                    pass
            else:
                pass
    #확률 분포를 상태 벡터의 형식으로 바꾸기 위해서 제곱근을 취한다.
    available_result = np.sqrt(np.array(available_result)) #확률 진폭은 계수의 제곱의 형태로 나오기 때문.
    #벡터의 크기가 1이 되도록 normalize해준다.
    normalized_result = available_result/np.linalg.norm(available_result)
    return normalized_result