# artificial_quotom_chip_toy.py
# A tiny educational quantum "chip" simulator (statevector) — for learning.
import numpy as np
from typing import List, Tuple

SQRT2_INV = 1 / np.sqrt(2)

class ArtificialQuotomChip:
    def __init__(self, n_qubits: int):
        self.n = n_qubits
        self.state = np.zeros(2**n, dtype=complex)
        self.state[0] = 1.0  # |00...0>
    
    def _apply_unitary(self, U: np.ndarray, targets: List[int]):
        """Apply an n-qubit unitary on specified target qubits (by building full matrix)."""
        # Build full operator by tensoring identities and U at target positions.
        # Note: simple but exponential; fine for small n (<= 16 practically).
        ops = []
        tset = set(targets)
        k = 0
        for i in range(self.n):
            if i in tset:
                ops.append(U if len(targets) == 1 else None)  # We'll handle multi-target separately
                k += 1
            else:
                ops.append(np.eye(2, dtype=complex))
        # If single-target, tensor directly
        if len(targets) == 1:
            full = ops[0]
            for op in ops[1:]:
                if op is None:
                    # should not happen here
                    op = np.eye(2, dtype=complex)
                full = np.kron(full, op)
        else:
            # For CNOT (2-qubit) quick path: generate full operator by acting on basis
            full = np.eye(2**self.n, dtype=complex)
            # We'll implement CNOT by permuting basis amplitudes (more efficient than building big matrices)
            return self._apply_custom_on_basis(targets, self._cnot_action)
        self.state = full @ self.state

    def _apply_custom_on_basis(self, targets: List[int], action_fn):
        """Apply a basis-level action function that maps basis index -> new basis index/value."""
        new = np.zeros_like(self.state)
        for idx, amp in enumerate(self.state):
            if amp == 0:
                continue
            new_idx, scale = action_fn(idx, targets)
            new[new_idx] += amp * scale
        self.state = new

    def _cnot_action(self, idx: int, targets: List[int]) -> Tuple[int, complex]:
        # targets: [control, target] (qubit indices with 0 = MSB if we constructed that way)
        control, target = targets
        # Convert index to bitstring array (LSB = last qubit). We'll treat qubit-0 as leftmost (MSB)
        bits = [(idx >> (self.n - 1 - i)) & 1 for i in range(self.n)]
        if bits[control] == 1:
            bits[target] ^= 1
        # convert bits back to index
        new_idx = 0
        for b in bits:
            new_idx = (new_idx << 1) | b
        return new_idx, 1.0

    # gates
    def H(self, q: int):
        H = np.array([[SQRT2_INV, SQRT2_INV], [SQRT2_INV, -SQRT2_INV]], dtype=complex)
        self._apply_unitary(H, [q])

    def X(self, q: int):
        X = np.array([[0,1],[1,0]], dtype=complex)
        self._apply_unitary(X, [q])

    def CNOT(self, control: int, target: int):
        # implement via basis mapping
        self._apply_custom_on_basis([control, target], self._cnot_action)

    def measure(self, q: int) -> int:
        """Measure qubit q (collapses state). Returns 0/1."""
        zero_mask = []
        one_mask = []
        for basis in range(2**self.n):
            # extract bit at position q
            b = (basis >> (self.n - 1 - q)) & 1
            if b == 0:
                zero_mask.append(basis)
            else:
                one_mask.append(basis)
        p0 = np.sum(np.abs(self.state[zero_mask])**2)
        if np.random.rand() < p0:
            # collapse to 0
            self.state[one_mask] = 0
            self.state /= np.sqrt(p0) if p0>0 else 1
            return 0
        else:
            p1 = 1 - p0
            self.state[zero_mask] = 0
            self.state /= np.sqrt(p1) if p1>0 else 1
            return 1

    def probs(self):
        return np.abs(self.state)**2

    def statevector(self):
        return self.state.copy()

# Example: create Bell pair on 2 qubits
if __name__ == "__main__":
    chip = ArtificialQuotomChip(2)
    chip.H(0)
    chip.CNOT(0,1)
    print("Statevector:", chip.statevector())
    print("Probs:", chip.probs())
    # Measure both
    m0 = chip.measure(0)
    m1 = chip.measure(1)
    print("Measurements:", m0, m1)