# --- 1. Installation and Setup ---
# Install necessary libraries silently
!pip install qiskit qiskit_aer pylatexenc -q

from qiskit import QuantumCircuit
from qiskit_aer import AerSimulator

# --- 2. DeutschSolver Class ---

class DeutschSolver:
    """
    A class to implement and run Deutsch's algorithm for a single-bit function.
    The goal is to determine if f: {0, 1} -> {0, 1} is constant or balanced
    with a single quantum query.
    """

    def __init__(self):
        # Initialize the simulator for execution
        self.simulator = AerSimulator()

    def create_oracle(self, case: int) -> QuantumCircuit:
        """
        Generates the quantum oracle (Uf) circuit for one of the 4 possible functions.
        """
        if case not in [1, 2, 3, 4]:
            raise ValueError("`case` must be 1, 2, 3, or 4.")

        # 2 qubits: Qubit 0 (input), Qubit 1 (auxiliary/output)
        f = QuantumCircuit(2, name=f"Uf{case}")

        # Functions f2 and f3 use CNOT
        if case in [2, 3]:
            f.cx(0, 1)

        # Functions f3 and f4 use an X-gate offset
        if case in [3, 4]:
            f.x(1)

        return f

    def build_circuit(self, oracle: QuantumCircuit) -> QuantumCircuit:
        """
        Compiles the full Deutsch algorithm circuit around the given oracle.
        """

        n = oracle.num_qubits - 1 # n=1 for the input register
        qc = QuantumCircuit(n + 1, n)

        # --- Stage 1: Initial Preparation ---

        # Prepare auxiliary qubit (Qubit n=1) in |1>
        qc.x(n)
        # Apply H to both qubits to create |+>|- >
        qc.h(range(n + 1))

        # --- Stage 2: Quantum Query ---
        qc.barrier()
        # Insert the oracle (Uf)
        qc.compose(oracle, inplace=True)
        qc.barrier()

        # --- Stage 3: Interference and Measurement ---

        # Apply final H gate to the input qubit (Qubit 0)
        qc.h(range(n))
        # Measure input qubit (Qubit 0)
        qc.measure(range(n), range(n))

        return qc

    def run(self, oracle: QuantumCircuit) -> Literal["constant", "balanced"]:
        """
        Executes the algorithm and determines if the function is constant or balanced.
        """

        # 1. Build the full circuit
        full_circuit = self.build_circuit(oracle)

        # 2. Run the simulation
        result = self.simulator.run(full_circuit, shots=1, memory=True).result()
        measurements = result.get_memory()

        # 3. Interpret the result
        # '0' -> Constant, '1' -> Balanced
        if measurements[0] == "0":
            return "constant"
        else:
            return "balanced"