
import numpy as np

# ─────────────────────────────────────────────────────────────────
#  Fault-Tolerant Quantum Computing Primitives
#
#  1. RepetitionCode    — simplest error-correcting code (1D)
#  2. SteaneCode        — 7-qubit [[7,1,3]] CSS code
#  3. SurfaceCode       — distance-d rotated surface code (simulation)
#  4. LogicalQubit      — high-level interface: encode, apply, measure
#  5. SyndromeDecoder   — majority-vote and lookup-table decoders
# ─────────────────────────────────────────────────────────────────


# ── 1. Repetition Code ───────────────────────────────────────────

class RepetitionCode:
    """
    3-qubit bit-flip repetition code.
    Encodes 1 logical qubit into 3 physical qubits.
    Detects and corrects any single-qubit bit-flip error.

    Usage:
        code = RepetitionCode()
        physical = code.encode(0)           # [0,0,0]
        noisy    = code.add_error(physical, error_qubit=1)
        syndrome = code.measure_syndrome(noisy)
        corrected= code.correct(noisy, syndrome)
        logical  = code.decode(corrected)
        print(logical)  # 0
    """
    def __init__(self, n_rep=3):
        self.n = n_rep  # must be odd

    def encode(self, logical_bit):
        """Encode 0 or 1 into n physical bits."""
        return [int(logical_bit)] * self.n

    def add_error(self, state, error_qubit):
        s = state.copy()
        s[error_qubit] = 1 - s[error_qubit]
        return s

    def add_random_errors(self, state, p_error):
        s = state.copy()
        for i in range(len(s)):
            if np.random.random() < p_error:
                s[i] = 1 - s[i]
        return s

    def measure_syndrome(self, state):
        """Returns parity checks between adjacent qubits."""
        return [state[i] ^ state[i+1] for i in range(self.n-1)]

    def correct(self, state, syndrome):
        """Majority-vote correction."""
        s = state.copy()
        votes = np.array(s)
        logical = int(np.round(np.mean(votes)))
        return [logical] * self.n

    def decode(self, state):
        return int(np.round(np.mean(state)))

    def logical_error_rate(self, p_physical, n_trials=10000):
        """Monte Carlo estimate of logical error rate."""
        errors = 0
        for _ in range(n_trials):
            s = self.encode(0)
            s = self.add_random_errors(s, p_physical)
            syn = self.measure_syndrome(s)
            s = self.correct(s, syn)
            if self.decode(s) != 0:
                errors += 1
        return errors / n_trials


# ── 2. Steane [[7,1,3]] Code ─────────────────────────────────────

class SteaneCode:
    """
    7-qubit [[7,1,3]] Steane code.
    Corrects any single-qubit error (X, Y, or Z).

    The parity check matrix H for this code is:
        H_X = H_Z = [parity checks for the [7,4,3] Hamming code]

    Usage:
        code = SteaneCode()
        enc = code.encode(0)
        err = code.add_x_error(enc, qubit=3)
        syn = code.measure_x_syndrome(err)
        cor = code.correct_x(err, syn)
        print(code.decode(cor))  # should be 0
    """

    # Parity check matrix of [7,4,3] Hamming code
    H = np.array([
        [1,0,1,0,1,0,1],
        [0,1,1,0,0,1,1],
        [0,0,0,1,1,1,1],
    ], dtype=int)

    # Logical |0> and |1> codewords (equal superposition of even/odd weight)
    _CW0 = [0b0000000, 0b1010101, 0b0110011, 0b1100110,
             0b0001111, 0b1011010, 0b0111100, 0b1101001]
    _CW1 = [0b1111111 ^ c for c in _CW0]

    def encode(self, logical_bit):
        """Return 7-bit codeword (as list of ints) for logical 0 or 1."""
        codewords = self._CW0 if logical_bit == 0 else self._CW1
        chosen = codewords[np.random.randint(len(codewords))]
        return [(chosen >> (6-i)) & 1 for i in range(7)]

    def add_x_error(self, state, qubit):
        s = state.copy(); s[qubit] ^= 1; return s

    def add_z_error(self, state, qubit):
        # Z errors are tracked in phase — we represent as separate register
        phase = [0]*7; phase[qubit] = 1; return state, phase

    def measure_x_syndrome(self, state):
        v = np.array(state, dtype=int)
        return list((self.H @ v) % 2)

    def correct_x(self, state, syndrome):
        s = state.copy()
        # Syndrome is binary rep of error position (1-indexed in Hamming)
        err_pos = int("".join(map(str, syndrome)), 2) - 1
        if 0 <= err_pos < 7:
            s[err_pos] ^= 1
        return s

    def decode(self, state):
        v = np.array(state)
        return int(np.round(np.mean(v)))

    def logical_error_rate(self, p_physical, n_trials=5000):
        errors = 0
        for _ in range(n_trials):
            s = self.encode(0)
            # Random X errors
            for i in range(7):
                if np.random.random() < p_physical:
                    s = self.add_x_error(s, i)
            syn = self.measure_x_syndrome(s)
            s = self.correct_x(s, syn)
            if self.decode(s) != 0:
                errors += 1
        return errors / n_trials


# ── 3. Surface Code (distance d) ─────────────────────────────────

class SurfaceCode:
    """
    Rotated surface code of distance d.
    Physical qubits: d^2 data + (d^2-1) ancilla
    Corrects up to floor((d-1)/2) errors.

    This is a classical simulation of the syndrome extraction
    and decoding — not a full statevector simulation (that would
    require 2^(d^2) dimensions, impractical for d>4).

    Usage:
        sc = SurfaceCode(distance=3)
        state = sc.init_logical_zero()
        state = sc.apply_noise(state, p=0.01)
        syndrome = sc.measure_syndrome(state)
        corrections = sc.decode(syndrome)
        log_err = sc.logical_error(state, corrections)
    """
    def __init__(self, distance=3):
        self.d = distance
        self.n_data = distance**2
        self.n_ancilla = distance**2 - 1

    def init_logical_zero(self):
        """All data qubits in |0>."""
        return np.zeros(self.n_data, dtype=int)

    def init_logical_one(self):
        """Logical |1> — apply X on a vertical line of qubits."""
        state = np.zeros(self.n_data, dtype=int)
        for i in range(self.d):
            state[i * self.d] = 1  # left column
        return state

    def apply_noise(self, state, p=0.01):
        """Apply independent depolarizing noise to each data qubit."""
        s = state.copy()
        for i in range(self.n_data):
            r = np.random.random()
            if r < p:           s[i] ^= 1          # X error
            elif r < 2*p:       pass                # Z error (not tracked here)
        return s

    def _x_stabilizers(self):
        """Return list of (ancilla_idx, [data_qubit_indices]) for X stabilizers."""
        d = self.d; stabs = []
        idx = 0
        for row in range(d-1):
            for col in range(d-1):
                q0 = row*d + col
                stabs.append((idx, [q0, q0+1, q0+d, q0+d+1]))
                idx += 1
        return stabs

    def measure_syndrome(self, state):
        """Measure X-type stabilizers (detect Z errors via bit-flip proxy)."""
        stabs = self._x_stabilizers()
        syndrome = np.zeros(len(stabs), dtype=int)
        for i, (_, qubits) in enumerate(stabs):
            syndrome[i] = int(np.sum(state[qubits]) % 2)
        return syndrome

    def decode(self, syndrome):
        """
        Minimum-weight perfect matching (simplified: greedy nearest-neighbour).
        Returns list of correction positions.
        """
        corrections = []
        stabs = self._x_stabilizers()
        flipped = [stabs[i][1] for i, s in enumerate(syndrome) if s == 1]
        for qlist in flipped:
            # Correct the first qubit in each flipped stabilizer
            corrections.append(qlist[0])
        return corrections

    def apply_corrections(self, state, corrections):
        s = state.copy()
        for q in corrections:
            s[q] ^= 1
        return s

    def logical_error(self, state, corrections):
        """Check if a logical error occurred (horizontal string of errors)."""
        s = self.apply_corrections(state, corrections)
        # Logical X = top row
        top_row = s[:self.d]
        return bool(np.sum(top_row) % 2 == 1)

    def logical_error_rate(self, p_physical, n_trials=3000):
        errors = 0
        for _ in range(n_trials):
            s = self.init_logical_zero()
            s = self.apply_noise(s, p_physical)
            syn = self.measure_syndrome(s)
            cor = self.decode(syn)
            if self.logical_error(s, cor):
                errors += 1
        return errors / n_trials

    def properties(self):
        return {"distance": self.d, "data_qubits": self.n_data,
                "ancilla_qubits": self.n_ancilla,
                "max_correctable_errors": (self.d-1)//2}


# ── 4. LogicalQubit (high-level interface) ────────────────────────

class LogicalQubit:
    """
    High-level logical qubit backed by a surface code.
    Abstracts away all the physical qubit details.

    Usage:
        lq = LogicalQubit(distance=3)
        lq.initialize(0)
        lq.apply_noise(p=0.005)
        result = lq.measure()
        print(f"Logical measurement: {result}")
    """
    def __init__(self, distance=3):
        self.code = SurfaceCode(distance)
        self._state = None
        self._logical_value = 0

    def initialize(self, logical_bit=0):
        self._logical_value = logical_bit
        if logical_bit == 0:
            self._state = self.code.init_logical_zero()
        else:
            self._state = self.code.init_logical_one()

    def apply_noise(self, p=0.01):
        if self._state is not None:
            self._state = self.code.apply_noise(self._state, p)

    def error_correct(self):
        if self._state is None: return
        syn = self.code.measure_syndrome(self._state)
        cor = self.code.decode(syn)
        self._state = self.code.apply_corrections(self._state, cor)

    def measure(self):
        if self._state is None: return None
        self.error_correct()
        err = self.code.logical_error(self._state, [])
        return self._logical_value ^ int(err)

    def fidelity(self, n_samples=200, p_noise=0.005):
        """Estimate logical fidelity by repeated trials."""
        correct = 0
        for _ in range(n_samples):
            self.initialize(self._logical_value)
            self.apply_noise(p_noise)
            if self.measure() == self._logical_value:
                correct += 1
        return correct / n_samples

    def __repr__(self):
        d = self.code.d
        return f"LogicalQubit(distance={d}, logical={self._logical_value})"
