from qiskit.circuit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.compiler import transpile

from qiskit.circuit.random import random_circuit
from qiskit_aer import AerSimulator

qc = random_circuit(16, 3, measure=True)
backend = AerSimulator()
t_qc = transpile(qc, backend, optimization_level=3, init_method='qubit_reuse')
t_qc_original = transpile(qc, backend, optimization_level=3) # For comparison later

qc.draw('mpl', fold=-1)

t_qc.draw('mpl', fold=-1)

from qiskit.quantum_info import hellinger_fidelity

counts_normal = backend.run(t_qc_original, shots=2**12).result().get_counts()
counts_reduced = backend.run(t_qc, shots=2**12).result().get_counts()

hellinger_fidelity(counts_normal, counts_reduced)

def build_bv_circuit(secret_string):
    input_size = len(secret_string)
    num_qubits = input_size + 1
    qr = QuantumRegister(num_qubits)
    cr = ClassicalRegister(input_size)
    qc = QuantumCircuit(qr, cr, name="main")
    qc.x(qr[input_size])
    qc.h(qr)
    qc.barrier()
    for i_qubit in range(input_size):
        if secret_string[input_size - 1 - i_qubit] == "1":
            qc.cx(qr[i_qubit], qr[input_size])
    qc.barrier()
    qc.h(qr)
    qc.x(qr[input_size])
    qc.barrier()
    qc.measure(qr[:-1], cr)
    return qc

secret = 5871
secret_string = format(secret, f"0b")
qc = build_bv_circuit(secret_string=secret_string)
qc.draw('mpl', fold=-1)

t_qc = transpile(qc, backend, optimization_level=3, init_method='qubit_reuse')
t_qc_original = transpile(qc, backend, optimization_level=3) # For comparison later

t_qc.draw('mpl', fold=-1)

counts_normal = backend.run(t_qc_original, shots=2**12).result().get_counts()
counts_reduced = backend.run(t_qc, shots=2**12).result().get_counts()

print(counts_normal, counts_reduced)
hellinger_fidelity(counts_normal, counts_reduced)

# Matrix Product State generator
from qiskit.circuit.library import RealAmplitudes

def MPS(num_qubits, **kwargs):
    """
    Constructs a Matrix Product State (MPS) quantum circuit.

    Args:
        num_qubits (int): The number of qubits in the circuit.
        **kwargs: Additional keyword arguments to be passed to the 
        RealAmplitudes.

    Returns:
        QuantumCircuit: The constructed MPS quantum circuit.
        
    """
    qc = QuantumCircuit(num_qubits)
    qubits = range(num_qubits)

    print(len(list(zip(qubits[:-1], qubits[1:]))))
    # Iterate over adjacent qubit pairs
    for i, j in zip(qubits[:-1], qubits[1:]):
        qc.compose(RealAmplitudes(num_qubits=2,
                                  parameter_prefix=f'θ_{i},{j}',
                                  **kwargs), [i, j],
                   inplace=True)
        qc.barrier(
        )  # Add a barrier after each block for clarity and separation

    return qc

# Tree tuples generator
def _generate_tree_tuples(n):
    """
    Generate a list of tuples representing the tree structure 
    of consecutive numbers up to n.

    Args:
        n (int): The number up to which the tuples are generated.

    Returns:
        list: A list of tuples representing the tree structure.
    """
    tuples_list = []
    indices = []

    # Generate initial tuples with consecutive numbers up to n
    for i in range(0, n, 2):
        tuples_list.append((i, i + 1))

    indices += [tuples_list]

    # Perform iterations until we reach a single tuple
    while len(tuples_list) > 1:
        new_tuples = []

        # Generate new tuples by combining adjacent larger numbers
        for i in range(0, len(tuples_list), 2):
            new_tuples.append((tuples_list[i][1], tuples_list[i + 1][1]))

        tuples_list = new_tuples
        indices += [tuples_list]

    return indices

# Tree Tensor Network Generator
def TTN(num_qubits, **kwargs):
    """
    Constructs a Tree Tensor Network (TTN) quantum circuit.

    Args:
        num_qubits (int): The number of qubits in the circuit.
        **kwargs: Additional keyword arguments to be passed to the 
        RealAmplitudes.

    Returns:
        QuantumCircuit: The constructed TTN quantum circuit.

    Raises:
        AssertionError: If the number of qubits is not a power of 2 
        or zero.
    """
    qc = QuantumCircuit(num_qubits)
    qubits = range(num_qubits)

    # Compute qubit indices
    assert num_qubits & (
        num_qubits -
        1) == 0 and num_qubits != 0, "Number of qubits must be a power of 2"

    indices = _generate_tree_tuples(num_qubits)

    # Iterate over each layer of TTN indices
    for layer_indices in indices:
        for i, j in layer_indices:
            qc.compose(RealAmplitudes(num_qubits=2,
                                      parameter_prefix=f'θ_{i},{j}',
                                      **kwargs), [i, j],
                       inplace=True)
        qc.barrier(
        )  # Add a barrier after each layer for clarity and separation

    return qc

import numpy as np

qc = TTN(2**3)
qc.measure_all()
params = {param: np.pi/(2**((index%3)+1)) for index, param in enumerate(qc.parameters)}
qc = qc.bind_parameters(params)

qc.draw('mpl', fold=-1)

t_qc = transpile(qc, backend, optimization_level=3, init_method='qubit_reuse')
t_qc_original = transpile(qc, backend, optimization_level=3) # For comparison later

t_qc = transpile(qc, backend, optimization_level=3, init_method='qubit_reuse')
t_qc_original = transpile(qc, backend, optimization_level=3) # For comparison later
t_qc.draw('mpl', fold=-1)


counts_normal = backend.run(t_qc_original, shots=2**12).result().get_counts()
counts_reduced = backend.run(t_qc, shots=2**12).result().get_counts()

# print(counts_normal, counts_reduced)
hellinger_fidelity(counts_normal, counts_reduced)

