import os
print(f"Working directory: '{os.getcwd()}'")

from test.python.transpiler.aqc.sample_data import ORIGINAL_CIRCUIT, INITIAL_THETAS
import numpy as np
from scipy.stats import unitary_group
from time import perf_counter
from qiskit import QuantumCircuit
from qiskit.algorithms.optimizers import L_BFGS_B
from qiskit.quantum_info import Operator
from qiskit.transpiler.synthesis.aqc.aqc import AQC
from qiskit.transpiler.synthesis.aqc.cnot_structures import make_cnot_network
from qiskit.transpiler.synthesis.aqc.cnot_unit_circuit import CNOTUnitCircuit
from qiskit.transpiler.synthesis.aqc.cnot_unit_objective import DefaultCNOTUnitObjective
from qiskit.transpiler.synthesis.aqc.fast_gradient.fast_gradient import FastCNOTUnitObjective

def print_misfits(target_mat: np.ndarray, approx_circ: QuantumCircuit, exe_tm: float):
    """
    Calculates and prints out misfit measures.
    """
    dim = target_mat.shape[0]
    approx_mat = Operator(approx_circ).data
    diff = approx_mat - target_mat 
    hs = np.vdot(approx_mat, target_mat)  # HS == Tr(V.H @ U)
    fidelity = (1.0 + np.abs(hs) ** 2 / dim) / (dim + 1)
    fro_err = 0.5 * (np.linalg.norm(diff, "fro") ** 2) / dim
    sin_err = np.linalg.norm(diff, 2)
    print("\nApproximation misfit:")
    print(f"Cost function based on Frobenius norm: {fro_err:0.5f}")
    print(f"Fidelity: {fidelity:0.5f}")
    print(f"Max. singular value of (V - U): {sin_err:0.5f}")
    print(f"Execution time: {exe_tm:0.3f} secs")
    print()

def make_target_matrix(num_qubits: int, target_name: str):
    """Creats a target matrix for testing."""
    if target_name == "mcx":
        # Define the target matrix: multi-control CNOT gate matrix.
        target_matrix = np.eye(2**num_qubits, dtype=np.cfloat)
        target_matrix[-2:, -2:] = [[0, 1], [1, 0]]
    else:
        # Define a random SU target matrix.
        dim = 2 ** num_qubits
        target_matrix = unitary_group.rvs(dim)
        target_matrix /= (np.linalg.det(target_matrix) ** (1.0 / float(dim)))
    return target_matrix

def create_circuit_and_objective(num_qubits: int, cnots: np.ndarray, fast: bool):
    """Instantiates AQC circuit and objective classes."""
    circ = CNOTUnitCircuit(num_qubits, cnots)
    if fast:
        objv = FastCNOTUnitObjective(num_qubits, cnots)
    else:
        objv = DefaultCNOTUnitObjective(num_qubits, cnots)
    return circ, objv

def aqc_example_hardcoded(fast: bool):
    tic = perf_counter()
    print(f"{'-' * 80}\nRunning \"{'accelerated' if fast else 'default'}\" "
          f"AQC implementation on a hardcoded target matrix...\n{'-' * 80}")

    # Define the number of qubits, circuit layout and other parameters.
    seed = 12345
    num_qubits = int(round(np.log2(ORIGINAL_CIRCUIT.shape[0])))
    depth = 0  # depth=0 implies maximum depth
    cnots = make_cnot_network(
        num_qubits=num_qubits, network_layout="spin", connectivity_type="full", depth=depth
    )
    # Define the target matrix.
    target_matrix = ORIGINAL_CIRCUIT

    # Create instances of the optimizer, AQC, circuit and objective classes.
    optimizer = L_BFGS_B(maxiter=200)
    circ, objv = create_circuit_and_objective(num_qubits, cnots, fast=fast)
    aqc = AQC(optimizer=optimizer, seed=seed)

    # Optimize the circuit to approximate the target matrix.
    aqc.compile_unitary(
        target_matrix=target_matrix,
        approximate_circuit=circ,
        approximating_objective=objv,
        initial_point=INITIAL_THETAS,
    )
    toc = perf_counter()
    print_misfits(target_matrix, circ, toc - tic)

aqc_example_hardcoded(fast=False)
aqc_example_hardcoded(fast=True)

def aqc_example(target_name: str, depth: int, fast: bool):
    tic = perf_counter()
    print(f"{'-' * 80}\nRunning {'accelerated' if fast else 'default'} AQC on "
          f"'{target_name}' target matrix with random initial guess...\n{'-' * 80}")

    # Define the number of qubits, circuit layout and other parameters.
    seed = 777
    np.random.seed(seed)
    num_qubits = int(4)
    cnots = make_cnot_network(
        num_qubits=num_qubits, network_layout="spin", connectivity_type="full", depth=depth
    )
    # Define a target matrix:
    target_matrix = make_target_matrix(num_qubits, target_name)

    # Create instances of the optimizer, AQC, circuit and objective classes.
    optimizer = L_BFGS_B(maxiter=1500)
    circ, objv = create_circuit_and_objective(num_qubits, cnots, fast=fast)
    aqc = AQC(optimizer=optimizer, seed=seed)

    # Optimize the circuit to approximate the target matrix.
    aqc.compile_unitary(
        target_matrix=target_matrix,
        approximate_circuit=circ,
        approximating_objective=objv,
        initial_point=2 * np.pi * np.random.rand(objv.num_thetas),
    )
    toc = perf_counter()
    print_misfits(target_matrix, circ, toc - tic)

depth = 64
print(f"\n\n******* Circuit depth = {depth}: *******\n")
aqc_example("random", depth, fast=True)
aqc_example("mcx", depth, fast=True)

depth = 40
print(f"\n\n******* Circuit depth = {depth}: *******\n")
aqc_example("random", depth, fast=True)
aqc_example("mcx", depth, fast=True)
aqc_example("mcx", depth, fast=False)