#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import qiskit
from qiskit_aer import AerSimulator


def initialize_backend(backend_config):
    backend_options = backend_config["backend_options"]
    simulator_type = backend_options["simulator_type"]
    shots = backend_options.get("shots", 1)

    # Map legacy simulator types to AerSimulator methods
    simulator_methods = {
        "aer_simulator": "automatic",
        "statevector_simulator": "statevector",
        "qasm_simulator": "automatic",
        "unitary_simulator": "unitary",
    }

    if simulator_type in simulator_methods:
        backend = AerSimulator(method=simulator_methods[simulator_type])
    else:
        backend = AerSimulator(method=simulator_type)

    backend.shots = shots
    return backend


def create_empty_circuit(num_qubits: int | None = None):
    if num_qubits is not None:
        return qiskit.QuantumCircuit(num_qubits)
    else:
        return qiskit.QuantumCircuit()


def apply_not_gate(circuit, qubit_index):
    # Apply a NOT gate (X gate) on the specified qubit
    circuit.x(qubit_index)


def apply_hadamard_gate(circuit, qubit_index):
    # Apply a Hadamard gate on the specified qubit
    circuit.h(qubit_index)


def apply_cnot_gate(circuit, control_qubit_index, target_qubit_index):
    # Apply a CNOT gate (controlled-X gate) with the specified control and
    # target qubits
    circuit.cx(control_qubit_index, target_qubit_index)


def apply_toffoli_gate(
    circuit, control_qubit_index1, control_qubit_index2, target_qubit_index
):
    # Apply a Toffoli gate (controlled-controlled-X gate) with the
    # specified control and target qubits
    circuit.ccx(control_qubit_index1, control_qubit_index2, target_qubit_index)


def apply_swap_gate(circuit, qubit_index1, qubit_index2):
    # Apply a SWAP gate to exchange the states of two qubits
    circuit.swap(qubit_index1, qubit_index2)


def apply_cswap_gate(
    circuit, control_qubit_index, target_qubit_index1, target_qubit_index2
):
    # Apply a controlled-SWAP (Fredkin) gate with the specified control and target qubits
    circuit.cswap(control_qubit_index, target_qubit_index1, target_qubit_index2)


def apply_pauli_x_gate(circuit, qubit_index):
    # Apply a Pauli X gate on the specified qubit
    circuit.x(qubit_index)


def apply_pauli_y_gate(circuit, qubit_index):
    # Apply a Pauli Y gate on the specified qubit
    circuit.y(qubit_index)


def apply_pauli_z_gate(circuit, qubit_index):
    # Apply a Pauli Z gate on the specified qubit
    circuit.z(qubit_index)


def apply_t_gate(circuit, qubit_index):
    # Apply a T gate (π/8 gate) on the specified qubit
    circuit.t(qubit_index)


def execute_circuit(circuit, backend, backend_config):
    working_circuit = circuit.copy()

    # Add measurements if they are not already present
    # Check if circuit already has measurement operations
    has_measurements = any(
        isinstance(inst.operation, qiskit.circuit.Measure)
        for inst in working_circuit.data
    )
    if not has_measurements:
        working_circuit.measure_all()

    # Ensure the circuit is parameterized properly
    if working_circuit.parameters:
        # Parse the global parameter configuration
        parameter_bindings = {
            param: backend_config["parameter_values"][str(param)]
            for param in working_circuit.parameters
        }
        transpiled_circuit = qiskit.transpile(working_circuit, backend)
        bound_circuit = transpiled_circuit.assign_parameters(parameter_bindings)
        job = backend.run(
            bound_circuit, shots=backend_config["backend_options"].get("shots", 1)
        )
        result = job.result()
        return result.get_counts()
    else:
        transpiled_circuit = qiskit.transpile(working_circuit, backend)
        job = backend.run(
            transpiled_circuit, shots=backend_config["backend_options"].get("shots", 1)
        )
        result = job.result()
        return result.get_counts()


# placeholder method for use in the testing suite
def get_final_state_vector(circuit, backend, backend_config):
    working_circuit = circuit.copy()

    simulator = AerSimulator(method="statevector")

    # Add save_statevector instruction
    working_circuit.save_statevector()

    # Bind parameters if present
    if working_circuit.parameters:
        parameter_values = backend_config.get("parameter_values", {})
        parameter_bindings = {
            param: parameter_values[str(param)] for param in working_circuit.parameters
        }
        working_circuit = working_circuit.assign_parameters(parameter_bindings)

    # Simulate the circuit
    transpiled_circuit = qiskit.transpile(working_circuit, simulator)
    job = simulator.run(transpiled_circuit)
    result = job.result()

    return result.get_statevector()


def draw_circuit(circuit):
    # Use Qiskit's built-in drawing function
    return circuit.draw()


def apply_rx_gate(circuit, qubit_index, angle):
    param = qiskit.circuit.Parameter(angle) if isinstance(angle, str) else angle
    circuit.rx(param, qubit_index)


def apply_ry_gate(circuit, qubit_index, angle):
    param = qiskit.circuit.Parameter(angle) if isinstance(angle, str) else angle
    circuit.ry(param, qubit_index)


def apply_rz_gate(circuit, qubit_index, angle):
    param = qiskit.circuit.Parameter(angle) if isinstance(angle, str) else angle
    circuit.rz(param, qubit_index)


def apply_u_gate(circuit, qubit_index, theta, phi, lambd):
    # Apply the U gate directly with specified parameters
    circuit.u(theta, phi, lambd, qubit_index)


def calculate_prob_zero(results, ancilla_qubit, num_qubits):
    """
    Calculate the probability of measuring the ancilla qubit in |0> state.

    Qiskit uses little-endian qubit ordering with string format results,
    where the rightmost bit corresponds to qubit 0.

    Args:
        results: Measurement results from execute_circuit() (dict with string keys)
        ancilla_qubit: Index of the ancilla qubit
        num_qubits: Total number of qubits in the circuit

    Returns:
        float: Probability of measuring ancilla in |0> state
    """
    # Handle different result formats from different backends
    if isinstance(results, list):
        results = results[0]

    total_shots = sum(results.values())
    count_zero = 0

    for state, count in results.items():
        # Qiskit: little-endian, rightmost bit is qubit 0
        if len(state) > ancilla_qubit and state[-(ancilla_qubit + 1)] == "0":
            count_zero += count

    return count_zero / total_shots if total_shots > 0 else 0.0
