
import numpy as np
import gc

# ─────────────────────────────────────────────────────────
#  Basis & bit constants
# ─────────────────────────────────────────────────────────
RECTILINEAR = 0   # { |0>, |1> }
DIAGONAL    = 1   # { |+>, |-> }


# ─────────────────────────────────────────────────────────
#  Low-level qubit helpers  (no external deps, RAM-safe)
# ─────────────────────────────────────────────────────────
def _prepare(bit: int, basis: int) -> np.ndarray:
    """
    Return a 2-element statevector for the given bit/basis.
      Rectilinear : 0 → |0>, 1 → |1>
      Diagonal    : 0 → |+>, 1 → |->
    """
    if basis == RECTILINEAR:
        sv = np.array([1.0, 0.0]) if bit == 0 else np.array([0.0, 1.0])
    else:
        sv = np.array([1.0,  1.0]) / np.sqrt(2) if bit == 0         else np.array([1.0, -1.0]) / np.sqrt(2)
    return sv


def _measure(sv: np.ndarray, basis: int) -> int:
    """
    Measure statevector sv in the chosen basis.
    Returns 0 or 1.
    """
    if basis == RECTILINEAR:
        prob1 = abs(sv[1]) ** 2
    else:
        # project onto |+> / |->
        plus  = np.array([1.0,  1.0]) / np.sqrt(2)
        minus = np.array([1.0, -1.0]) / np.sqrt(2)
        p_plus = abs(np.dot(plus, sv)) ** 2
        prob1  = 1.0 - p_plus   # prob of |-> = bit 1
    return int(np.random.random() < prob1)


def _intercept(sv: np.ndarray, eve_basis: int) -> np.ndarray:
    """
    Eve measures sv in eve_basis and re-prepares.
    Returns the (possibly disturbed) new statevector.
    """
    bit = _measure(sv, eve_basis)
    return _prepare(bit, eve_basis)


# ─────────────────────────────────────────────────────────
#  BB84Protocol
# ─────────────────────────────────────────────────────────
class BB84Protocol:
    """
    Full BB84 quantum key distribution simulation.

    Features
    --------
    - Alice prepares random bits in random bases
    - Bob measures in random bases
    - Optional Eve intercepts with a configurable strategy
    - Sifting: keep only matching-basis bits
    - Error rate estimation and eavesdropping detection
    - Final secret key extraction

    Usage
    -----
        bb84 = BB84Protocol(num_bits=200)
        result = bb84.run(eavesdrop=False)
        bb84.report(result)

        result_eve = bb84.run(eavesdrop=True, eve_intercept_rate=1.0)
        bb84.report(result_eve)
    """

    def __init__(self, num_bits: int = 200, seed: int = None):
        if num_bits < 10 or num_bits > 100_000:
            raise ValueError("num_bits must be between 10 and 100 000")
        self.num_bits = num_bits
        if seed is not None:
            np.random.seed(seed)

    # ── main run ──────────────────────────────────────────
    def run(self,
            eavesdrop: bool = False,
            eve_intercept_rate: float = 1.0,
            error_sample_size: int = 20) -> dict:
        """
        Simulate the full BB84 protocol.

        Parameters
        ----------
        eavesdrop          : Whether Eve is present.
        eve_intercept_rate : Fraction of qubits Eve intercepts (0–1).
        error_sample_size  : How many sifted bits to sacrifice for QBER check.

        Returns
        -------
        dict with all protocol details.
        """
        n = self.num_bits

        # ── Step 1: Alice prepares ────────────────────────
        alice_bits   = np.random.randint(0, 2, n)
        alice_bases  = np.random.randint(0, 2, n)

        # ── Step 2: (Optional) Eve intercepts ────────────
        intercepted  = np.zeros(n, dtype=bool)
        transmitted  = []   # statevectors after Eve

        for i in range(n):
            sv = _prepare(int(alice_bits[i]), int(alice_bases[i]))
            if eavesdrop and np.random.random() < eve_intercept_rate:
                eve_basis = np.random.randint(0, 2)
                sv = _intercept(sv, eve_basis)
                intercepted[i] = True
            transmitted.append(sv)

        # ── Step 3: Bob measures ──────────────────────────
        bob_bases  = np.random.randint(0, 2, n)
        bob_bits   = np.array([
            _measure(transmitted[i], int(bob_bases[i])) for i in range(n)
        ])

        # free statevectors immediately
        del transmitted
        gc.collect()

        # ── Step 4: Basis sifting ─────────────────────────
        matching = (alice_bases == bob_bases)
        sifted_alice = alice_bits[matching]
        sifted_bob   = bob_bits[matching]
        sifted_idx   = np.where(matching)[0]
        sifted_len   = len(sifted_alice)

        # ── Step 5: QBER estimation ───────────────────────
        sample_n = min(error_sample_size, sifted_len // 2)
        if sample_n > 0:
            sample_idx  = np.random.choice(sifted_len, sample_n, replace=False)
            errors      = np.sum(sifted_alice[sample_idx] != sifted_bob[sample_idx])
            qber        = errors / sample_n
            # remove sacrificed bits from final key
            keep_mask   = np.ones(sifted_len, dtype=bool)
            keep_mask[sample_idx] = False
            key_alice   = sifted_alice[keep_mask]
            key_bob     = sifted_bob[keep_mask]
        else:
            qber      = 0.0
            key_alice = sifted_alice
            key_bob   = sifted_bob

        # ── Step 6: Eve detection decision ───────────────
        # BB84 theory: random intercept → ~25% QBER
        eve_detected = bool(qber > 0.10)   # threshold: 10%

        # ── Step 7: Final key (only if safe) ─────────────
        if not eve_detected:
            secret_key = key_alice.tolist()
        else:
            secret_key = []   # abort — channel compromised

        return {
            "num_bits"        : n,
            "sifted_length"   : sifted_len,
            "key_length"      : len(secret_key),
            "qber"            : round(float(qber), 4),
            "eve_present"     : eavesdrop,
            "eve_detected"    : eve_detected,
            "intercepted_pct" : round(float(intercepted.sum()) / n * 100, 1),
            "secret_key"      : secret_key,
            "alice_bases"     : alice_bases.tolist(),
            "bob_bases"       : bob_bases.tolist(),
            "matching_bases"  : int(matching.sum()),
        }

    # ── pretty report ─────────────────────────────────────
    @staticmethod
    def report(result: dict):
        eve_sym  = "🔴 YES — DETECTED" if result["eve_detected"] else "🟢 No"
        key_info = (
            f"{result['key_length']} bits"
            if result["key_length"] > 0
            else "❌ ABORTED (channel unsafe)"
        )
        print("=" * 56)
        print("  BB84 QKD Protocol — Report")
        print("=" * 56)
        print(f"  Qubits sent          : {result['num_bits']}")
        print(f"  Matching bases       : {result['matching_bases']}")
        print(f"  Sifted key length    : {result['sifted_length']}")
        print(f"  QBER                 : {result['qber']*100:.2f}%")
        print(f"  Eve present          : {result['eve_present']}")
        print(f"  Intercepted          : {result['intercepted_pct']:.1f}%")
        print(f"  Eve detected         : {eve_sym}")
        print(f"  Final secret key     : {key_info}")
        if result["key_length"] > 0:
            preview = result["secret_key"][:16]
            print(f"  Key preview (16b)    : {''.join(map(str, preview))}...")
        print("=" * 56)

    # ── key as hex string ─────────────────────────────────
    @staticmethod
    def key_to_hex(result: dict) -> str:
        """Convert the secret key bit list to a hex string."""
        bits = result.get("secret_key", [])
        if not bits:
            return ""
        # pad to multiple of 8
        pad  = (-len(bits)) % 8
        bits = bits + [0] * pad
        hexstr = ""
        for i in range(0, len(bits), 8):
            byte = int("".join(map(str, bits[i:i+8])), 2)
            hexstr += f"{byte:02x}"
        return hexstr


# ─────────────────────────────────────────────────────────
#  Eavesdropping Detector  (standalone utility)
# ─────────────────────────────────────────────────────────
class EavesdroppingDetector:
    """
    Analyse a QBER value and give a confidence verdict.

    Usage
    -----
        det = EavesdroppingDetector()
        det.analyse(qber=0.22)
    """

    THRESHOLDS = [
        (0.00, 0.05, "✅ SAFE",     "QBER < 5% — channel is clean."),
        (0.05, 0.10, "⚠️  WARNING",  "QBER 5–10% — possible noise or weak eavesdrop."),
        (0.10, 0.25, "🔴 DANGER",   "QBER 10–25% — likely eavesdropping. Abort."),
        (0.25, 1.01, "☠️  CRITICAL", "QBER > 25% — full intercept-resend attack detected."),
    ]

    def analyse(self, qber: float) -> dict:
        for lo, hi, label, msg in self.THRESHOLDS:
            if lo <= qber < hi:
                verdict, detail = label, msg
                break
        else:
            verdict, detail = "UNKNOWN", "QBER out of range."

        print(f"  QBER = {qber*100:.2f}%  →  {verdict}")
        print(f"  {detail}")
        return {"qber": qber, "verdict": verdict, "detail": detail}
