!pip install -qq -U "classiq[qsp]"
!pip install -qq "classiq[chemistry]"

import time

import matplotlib.pyplot as plt
import numpy as np
from banded_be import get_banded_diags_be
from cheb_utils import *
from classical_functions_be import get_projected_state_vector, get_svd_range
from pauli_be import get_pauli_be
from scipy import sparse

from classiq import *
from classiq.applications.qsp import qsvt_phases

np.random.seed(53)

PAULI_TRIM_REL_TOL = 0.1

def qsvt_solver(
    mat_raw_scr,
    b_raw,
    poly_degree,
    be_method="banded",
    cheb_approx_type="numpy_interpolated",
    preferences=Preferences(),
    constraints=Constraints(),
):

    SCALE = 0.5
    b_norm = np.linalg.norm(b_raw)  # b normalization
    b_normalized = b_raw / b_norm

    # Define block encoding
    if be_method == "pauli":
        data_size, block_size, be_scaling_factor, be_qfunc = get_pauli_be(
            mat_raw_scr, PAULI_TRIM_REL_TOL
        )
        print(
            f"Pauli block encoding with block size {block_size} and scaling factor {be_scaling_factor}"
        )

    elif be_method == "banded":
        data_size, block_size, be_scaling_factor, be_qfunc = get_banded_diags_be(
            mat_raw_scr
        )
        print(
            f"Banded diagonal block encoding with block size {block_size} and scaling factor {be_scaling_factor}"
        )

    class BlockEncodedState(QStruct):
        data: QNum[data_size]
        block: QNum[block_size]

    # Get SVD range
    w_min, w_max = get_svd_range(mat_raw_scr / be_scaling_factor)
    # Get Chebyshev polynomial and the corresponding QSVT angles
    pcoefs, poly_scale = get_cheb_coeff(
        w_min, poly_degree, w_max, scale=SCALE, method=cheb_approx_type, epsilon=0.01
    )
    inv_phases = qsvt_phases(pcoefs, tol=0.001)

    # Define QSVT projector
    @qfunc
    def projector(be: BlockEncodedState, res: QBit):
        res ^= be.block == 0

    @qfunc
    def main(
        qsvt_aux: Output[QBit],
        data: Output[QNum[data_size]],
        block: Output[QNum[block_size]],
    ):
        allocate(qsvt_aux)
        allocate(block)
        prepare_amplitudes(b_normalized.tolist(), 0, data)

        be_state = BlockEncodedState()

        within_apply(
            lambda: bind([data, block], be_state),
            lambda: qsvt_inversion(
                inv_phases,
                lambda aux: projector(be_state, aux),
                lambda: be_qfunc(be_state.block, be_state.data),
                qsvt_aux,
            ),
        )

    start_time_syn = time.time()
    qprog = synthesize(main, preferences=preferences, constraints=constraints)
    print("time to syn:", time.time() - start_time_syn)

    execution_preferences = ExecutionPreferences(
        num_shots=1,
        backend_preferences=ClassiqBackendPreferences(
            backend_name=ClassiqSimulatorBackendNames.SIMULATOR_STATEVECTOR
        ),
    )

    start_time_exe = time.time()
    with ExecutionSession(qprog, execution_preferences) as es:
        es.set_measured_state_filter("block", lambda state: state == 0.0)
        es.set_measured_state_filter("qsvt_aux", lambda state: state == 0.0)
        results = es.sample()
    print("time to exe:", time.time() - start_time_exe)

    resulting_state = get_projected_state_vector(results, "data")

    normalization_factor = (be_scaling_factor * poly_scale) / b_norm

    return resulting_state / normalization_factor, qprog

# For faster results use the commented out preferences


prefs = Preferences()

# prefs = Preferences(
#     transpilation_option="none", optimization_level=0, debug_mode=False, qasm3=True
# )
poly_degree = 101

import pathlib

path = (
    pathlib.Path(__file__).parent.resolve()
    if "__file__" in locals()
    else pathlib.Path(".")
)

mat_name = "nozzle_small_scr"
matfile = "matrices/" + mat_name + ".npz"
mat_raw_scr = sparse.load_npz(path / matfile)

b_raw = np.load(path / "matrices/b_nozzle_small.npy")
raw_mat_qubits = len(b_raw).bit_length() - 1  # matrix raw size

print(f"Raw matrix size: {raw_mat_qubits} qubits.")

qsol_banded, qprog_banded = qsvt_solver(
    mat_raw_scr,
    b_raw,
    poly_degree,
    be_method="banded",
    preferences=prefs,
    constraints=Constraints(optimization_parameter="width"),
)
show(qprog_banded)

qsol_pauli, qprog_pauli = qsvt_solver(
    mat_raw_scr,
    b_raw,
    poly_degree,
    be_method="pauli",
    preferences=prefs,
    constraints=Constraints(optimization_parameter="width"),
)
show(qprog_pauli)

mat_raw = mat_raw_scr.toarray()
expected_sol = np.linalg.solve(mat_raw, b_raw)
plt.plot(expected_sol, "o", label="Classical")
ext_idx = np.argmax(np.abs(expected_sol))
correct_sign = np.sign(expected_sol[ext_idx]) / np.sign(qsol_pauli[ext_idx])
qsol_pauli *= correct_sign
plt.plot(qsol_pauli, ".", label=f"QSVT-inv; Pauli BE; degree {poly_degree}")

correct_sign = np.sign(expected_sol[ext_idx]) / np.sign(qsol_banded[ext_idx])
qsol_banded *= correct_sign
plt.plot(qsol_banded, ".", label=f"QSVT-inv; Banded BE; degree {poly_degree}")
plt.legend()

assert np.linalg.norm(qsol_banded - expected_sol) < 0.2
assert np.linalg.norm(qsol_pauli - expected_sol) < 0.2