import numpy as np
import time

import qiskit

from qiskit import QuantumCircuit
from qiskit.algorithms.linear_solvers.hhl import HHL as HHL_algo
from qiskit.algorithms.linear_solvers.matrices import TridiagonalToeplitz


def get_hhl_solution(qubits_count: int, main_diag: float, off_diag: float, rhs: np.ndarray, epsilon: float):
    matrix = TridiagonalToeplitz(qubits_count, main_diag, off_diag)

    num_qubits = matrix.num_state_qubits
    vector_circuit = QuantumCircuit(num_qubits)
    vector_circuit.isometry(rhs, list(range(num_qubits)), None)

    hhl = HHL_algo(epsilon=epsilon)
    solution = hhl.solve(matrix, vector_circuit)
    return solution


from qiskit.opflow import Z, I, StateFn, TensoredOp
from qiskit.quantum_info import Statevector, DensityMatrix, partial_trace


def get_quantum_solution(initial_qubits_count, solution):
    total_qubits_count = len(solution.state.qubits)
    zero_op = (I + Z) / 2
    one_op = (I - Z) / 2

    na = initial_qubits_count - 1
    nb = initial_qubits_count
    nl = total_qubits_count - na - nb - 1

    observable = one_op ^ TensoredOp((nl + na) * [zero_op]) ^ (I ^ nb)
    final_state = (observable @ StateFn(solution.state))
    
    final_vector = final_state.eval()
    final_vector /= np.linalg.norm(final_vector.to_matrix())
    
    # It's needed for converting complex amplitudes into real numbers
    norm = np.array(list(map(np.linalg.norm, list(final_vector.to_matrix()))))
    sgn = np.array(list(map(np.sign, list(final_vector.to_matrix()))))
    result_vector = norm * sgn
    
    dmatrix = DensityMatrix(result_vector)
    tr = partial_trace(dmatrix, qargs=range(nb, total_qubits_count))

    print("Partial trace purity:", tr.purity())  # Should be equal to 1
    quantum_solution = tr.to_statevector().data
    return quantum_solution


def get_trigiagonal_matrix(n: int, main_diag: float, off_diag: float) -> np.ndarray:
    A = np.zeros((n, n))
    for i in range(n):
        A[i, i] = main_diag
        if i + 1 < n:
            A[i + 1][i] = off_diag
            A[i][i + 1] = off_diag
    return A


def classic_tridiagonal_system_solver(main_diag: float, off_diag: float, b: np.ndarray) -> np.ndarray:
    n = np.size(b)
    A = get_trigiagonal_matrix(n, main_diag, off_diag)
    solution = np.linalg.solve(A, b)
    return solution


def get_solution_distance(quantum_solution, main_diag, off_diag, rhs):
    classical_solution = classic_tridiagonal_system_solver(main_diag, off_diag, rhs)
    classical_solution = classical_solution / np.linalg.norm(classical_solution)
    solution_distance = min(np.linalg.norm(classical_solution - quantum_solution), np.linalg.norm(classical_solution + quantum_solution))
    return solution_distance


qubits_count = 2
main_diag = 1
off_diag = 0.5

np.random.seed(42)
right_hand_side = np.array([1.0, -2.1, 3.2, -4.3])
rhs = right_hand_side / np.linalg.norm(right_hand_side)

start_eps = 1e-1
eps_counts = 13

all_eps = [start_eps * 2 ** (-i) for i in range(eps_counts)]

precision = []

for eps in all_eps:
    print("Current eps:", eps)
    start_time = time.time()
    solution = get_hhl_solution(qubits_count, main_diag, off_diag, rhs, eps)
    quantum_solution = get_quantum_solution(qubits_count, solution)
    precision.append(get_solution_distance(quantum_solution, main_diag, off_diag, rhs))
    print("Current precision:", precision[-1])
    print("Iteration time:", time.time() - start_time)
    print("")