from __future__ import annotations

import math
from collections import deque
from typing import Dict, List, Tuple

import numpy as np


class SurfaceCode:
    """
    Lightweight stabilizer-style simulator.

    Uses binary error frames (X/Z) plus parity-check matrices to expose a
    syndrome-based control problem while staying within small CPU/RAM limits.
    """

    def __init__(self, distance: int):
        self.distance = int(distance)
        self.n_qubits = 2 * (self.distance ** 2) - 2 * self.distance + 1
        self.grid_cols = math.ceil(math.sqrt(self.n_qubits))
        self.grid_rows = math.ceil(self.n_qubits / self.grid_cols)

        self.x_err = np.zeros(self.n_qubits, dtype=np.uint8)
        self.z_err = np.zeros(self.n_qubits, dtype=np.uint8)

        self.hx, self.hz = self._build_parity_checks()

    def _neighbors(self, idx: int) -> List[int]:
        r = idx // self.grid_cols
        c = idx % self.grid_cols
        out: List[int] = [idx]
        cand = [
            (r - 1, c),
            (r + 1, c),
            (r, c - 1),
            (r, c + 1),
        ]
        for rr, cc in cand:
            if 0 <= rr < self.grid_rows and 0 <= cc < self.grid_cols:
                j = rr * self.grid_cols + cc
                if j < self.n_qubits:
                    out.append(j)
        return out

    def _build_parity_checks(self) -> Tuple[np.ndarray, np.ndarray]:
        x_rows: List[np.ndarray] = []
        z_rows: List[np.ndarray] = []
        for i in range(self.n_qubits - 1):
            row_x = np.zeros(self.n_qubits, dtype=np.uint8)
            row_z = np.zeros(self.n_qubits, dtype=np.uint8)
            for j in self._neighbors(i):
                row_x[j] = 1
            for j in self._neighbors((i + self.distance) % self.n_qubits):
                row_z[j] = 1
            x_rows.append(row_x)
            z_rows.append(row_z)
        return np.vstack(x_rows), np.vstack(z_rows)

    def apply_error(self, qubit_idx: int, error_type: str) -> None:
        if qubit_idx < 0 or qubit_idx >= self.n_qubits:
            return
        if error_type == "X":
            self.x_err[qubit_idx] ^= 1
        elif error_type == "Z":
            self.z_err[qubit_idx] ^= 1
        elif error_type == "Y":
            self.x_err[qubit_idx] ^= 1
            self.z_err[qubit_idx] ^= 1

    def measure_syndromes(self) -> List[int]:
        x_syn = (self.hx @ self.z_err) % 2
        z_syn = (self.hz @ self.x_err) % 2
        return np.concatenate([x_syn, z_syn]).astype(int).tolist()

    def apply_correction(self, corrections: Dict[int, str]) -> None:
        for q, op in corrections.items():
            self.apply_error(int(q), op)

    def _spans_vertical(self, active: np.ndarray) -> bool:
        top = [i for i in range(min(self.grid_cols, self.n_qubits)) if active[i] == 1]
        if not top:
            return False
        last_row_start = (self.grid_rows - 1) * self.grid_cols
        bottom = {
            i
            for i in range(last_row_start, min(last_row_start + self.grid_cols, self.n_qubits))
            if active[i] == 1
        }
        if not bottom:
            return False

        seen = set(top)
        queue = deque(top)
        while queue:
            cur = queue.popleft()
            if cur in bottom:
                return True
            for nxt in self._neighbors(cur):
                if active[nxt] == 1 and nxt not in seen:
                    seen.add(nxt)
                    queue.append(nxt)
        return False

    def _spans_horizontal(self, active: np.ndarray) -> bool:
        left = [i for i in range(0, self.n_qubits, self.grid_cols) if active[i] == 1]
        if not left:
            return False
        right = {
            i
            for i in range(self.grid_cols - 1, self.n_qubits, self.grid_cols)
            if active[i] == 1
        }
        if not right:
            return False

        seen = set(left)
        queue = deque(left)
        while queue:
            cur = queue.popleft()
            if cur in right:
                return True
            for nxt in self._neighbors(cur):
                if active[nxt] == 1 and nxt not in seen:
                    seen.add(nxt)
                    queue.append(nxt)
        return False

    def check_logical_failure(self) -> bool:
        return self._spans_vertical(self.z_err) or self._spans_horizontal(self.x_err)

    def get_fidelity_estimate(self) -> float:
        weight = int(self.x_err.sum() + self.z_err.sum())
        ratio = weight / max(1, 2 * self.n_qubits)
        return max(0.0, 1.0 - ratio)

    @property
    def error_frame(self) -> Dict[str, List[int]]:
        return {
            "x": self.x_err.astype(int).tolist(),
            "z": self.z_err.astype(int).tolist(),
        }