
"""
PKTron Finance — BB84 Quantum Key Distribution
===============================================
Simulates the BB84 protocol for secure banking communication.

Parties: Alice (sender) | Bob (receiver) | Eve (optional eavesdropper)

Steps:
  1. Alice prepares random qubits in Z/X bases
  2. Bob measures in random bases
  3. Basis sifting — keep matching basis positions
  4. Error estimation — detect Eve
  5. Privacy amplification (simple hash) → final secret key

Memory: O(n_bits) — trivially safe for 12 GB.
"""
import numpy as np
import hashlib

_R = np.float32

# ── Single qubit helpers ─────────────────────────────────────────────────

def _prepare_qubit(bit, basis):
    """
    Prepare a qubit state vector.
    basis=0 (Z): |0> or |1>
    basis=1 (X): |+> or |->
    Returns complex64 array of length 2.
    """
    if basis == 0:
        return np.array([1,0], dtype=np.complex64) if bit==0                else np.array([0,1], dtype=np.complex64)
    else:
        f = _R(1/np.sqrt(2))
        return np.array([f, f], dtype=np.complex64) if bit==0                else np.array([f,-f], dtype=np.complex64)

def _measure_qubit(state, basis, rng):
    """
    Measure qubit in given basis (0=Z, 1=X).
    Returns measured bit (0 or 1).
    """
    if basis == 1:
        H = _R(1/np.sqrt(2)) * np.array([[1,1],[1,-1]], dtype=np.complex64)
        state = H @ state
    probs = np.abs(state)**2
    probs = np.clip(probs.real, 0, None); probs /= probs.sum()
    return int(rng.choice(2, p=probs))

# ── Eve intercept ────────────────────────────────────────────────────────

def _eve_intercept(state, rng, intercept_rate=1.0):
    """Eve randomly measures in Z or X basis and re-prepares."""
    if rng.random() > intercept_rate:
        return state
    eve_basis = int(rng.integers(0, 2))
    bit = _measure_qubit(state, eve_basis, rng)
    return _prepare_qubit(bit, eve_basis)

# ── Public API ───────────────────────────────────────────────────────────

def generate_key(n_bits=256, seed=0):
    """
    Alice generates random bits and bases.

    Returns
    -------
    dict: bits (array), bases (array), n_bits
    """
    rng   = np.random.default_rng(seed)
    bits  = rng.integers(0, 2, size=n_bits, dtype=np.uint8)
    bases = rng.integers(0, 2, size=n_bits, dtype=np.uint8)
    return {'bits': bits, 'bases': bases, 'n_bits': n_bits}


def simulate_channel(alice_key, eavesdrop=False,
                     intercept_rate=0.5, channel_error=0.02, seed=1):
    """
    Simulate quantum channel: Alice -> (Eve?) -> Bob.

    Parameters
    ----------
    alice_key      : dict from generate_key()
    eavesdrop      : bool   include Eve in channel
    intercept_rate : float  fraction of qubits Eve intercepts
    channel_error  : float  baseline depolarising error

    Returns
    -------
    dict:
        alice_bits, alice_bases,
        bob_bits, bob_bases,
        eve_present (bool),
        sifted_alice, sifted_bob, sifted_positions
    """
    rng         = np.random.default_rng(seed)
    a_bits      = alice_key['bits']
    a_bases     = alice_key['bases']
    n           = len(a_bits)
    bob_bases   = rng.integers(0, 2, size=n, dtype=np.uint8)
    bob_bits    = np.zeros(n, dtype=np.uint8)

    for i in range(n):
        state = _prepare_qubit(int(a_bits[i]), int(a_bases[i]))

        # Channel noise
        if rng.random() < channel_error:
            state = np.array([rng.random(), rng.random()], dtype=np.complex64)
            state /= np.linalg.norm(state)

        # Eve intercept
        if eavesdrop:
            state = _eve_intercept(state, rng, intercept_rate)

        bob_bits[i] = _measure_qubit(state, int(bob_bases[i]), rng)

    # Basis sifting
    match = (a_bases == bob_bases)
    sift_pos     = np.where(match)[0]
    sift_alice   = a_bits[sift_pos]
    sift_bob     = bob_bits[sift_pos]

    return {'alice_bits'      : a_bits,
            'alice_bases'     : a_bases,
            'bob_bits'        : bob_bits,
            'bob_bases'       : bob_bases,
            'eve_present'     : eavesdrop,
            'sifted_alice'    : sift_alice,
            'sifted_bob'      : sift_bob,
            'sifted_positions': sift_pos}


def detect_eavesdropping(channel_result, sample_fraction=0.25,
                         qber_threshold=0.11):
    """
    Estimate Quantum Bit Error Rate (QBER) and detect Eve.

    A QBER > 11% (BB84 security threshold) strongly indicates eavesdropping.

    Parameters
    ----------
    channel_result   : dict from simulate_channel()
    sample_fraction  : fraction of sifted key used for error estimation
    qber_threshold   : float  standard BB84 threshold (0.11)

    Returns
    -------
    dict:
        qber, eve_detected, secret_key (hex), key_length,
        sifted_key_length
    """
    sa  = channel_result['sifted_alice']
    sb  = channel_result['sifted_bob']
    n_s = len(sa)

    if n_s == 0:
        return {'qber': 1.0, 'eve_detected': True,
                'secret_key': '', 'key_length': 0,
                'sifted_key_length': 0}

    # Sample for QBER estimation
    n_check  = max(1, int(n_s * sample_fraction))
    idx      = np.arange(n_s); np.random.shuffle(idx)
    check    = idx[:n_check]
    keep     = idx[n_check:]

    errors   = int((sa[check] != sb[check]).sum())
    qber     = errors / n_check

    eve_det  = qber > qber_threshold

    # Privacy amplification on remaining sifted bits (simple SHA-256 hash)
    if len(keep) > 0 and not eve_det:
        raw_bits = ''.join(str(b) for b in sa[keep])
        # Pack bits into bytes for hashing
        padded = raw_bits + '0' * ((-len(raw_bits)) % 8)
        raw_bytes = bytes(int(padded[i:i+8], 2)
                          for i in range(0, len(padded), 8))
        secret_key = hashlib.sha256(raw_bytes).hexdigest()
        key_length = 256   # SHA-256 always 256 bits
    else:
        secret_key = ''
        key_length = 0

    return {'qber'              : round(qber, 4),
            'eve_detected'      : eve_det,
            'secret_key'        : secret_key,
            'key_length_bits'   : key_length,
            'sifted_key_length' : n_s,
            'bits_checked'      : n_check,
            'security_status'   : 'COMPROMISED' if eve_det else 'SECURE'}
