
"""
PKTron Finance — Grover Search for Fraud Patterns
==================================================
Uses Grover's algorithm to amplify quantum states corresponding to
fraudulent transaction patterns within a database.

Memory safe: complex64; n_items <= 2^20 = 1M entries safe for 12 GB.
"""
import numpy as np

_C = np.complex64
_R = np.float32

# ── Oracle & Diffusion ───────────────────────────────────────────────────

def build_oracle(n_qubits, marked_indices):
    """
    Build a phase-flip oracle that marks specific database entries.

    Parameters
    ----------
    n_qubits       : int        log2 of database size
    marked_indices : list[int]  indices of fraudulent patterns

    Returns
    -------
    oracle_fn : function  oracle_fn(state) -> state with phase flips
    """
    marked = set(int(i) for i in marked_indices)

    def oracle_fn(state):
        new = state.copy()
        for idx in marked:
            if 0 <= idx < len(new):
                new[idx] = -new[idx]
        return new

    return oracle_fn


def _diffusion(state):
    """Grover diffusion operator: 2|s><s| - I  where |s> = uniform."""
    n   = len(state)
    mu  = _C(state.mean())
    return (2 * mu - state).astype(_C)


def _optimal_iterations(n_qubits, n_marked):
    """Optimal Grover iterations: floor(pi/4 * sqrt(N/M))."""
    N = 2**n_qubits
    M = max(1, n_marked)
    k = int(np.pi / 4 * np.sqrt(N / M))
    return max(1, k)


def run_grover_search(n_qubits, marked_indices, n_iter=None,
                      n_shots=1024, seed=42):
    """
    Run Grover's algorithm to find fraudulent patterns.

    Parameters
    ----------
    n_qubits       : int         number of qubits (database size = 2^n)
    marked_indices : list[int]   fraudulent pattern indices
    n_iter         : int or None optimal iterations if None
    n_shots        : int         measurement shots

    Returns
    -------
    dict:
        found_indices   : list[int]   top detected fraudulent indices
        probabilities   : array       final probability distribution
        success_prob    : float       total probability on marked items
        n_iterations    : int
    """
    N        = 2**n_qubits
    n_marked = len(marked_indices)
    k        = n_iter if n_iter is not None else _optimal_iterations(n_qubits, n_marked)
    oracle   = build_oracle(n_qubits, marked_indices)

    # Initial uniform superposition
    state = np.ones(N, dtype=_C) / _R(np.sqrt(N))

    for _ in range(k):
        state = oracle(state)
        state = _diffusion(state)

    probs        = np.abs(state)**2
    success_prob = float(probs[list(marked_indices)].sum())

    # Simulate measurement
    rng      = np.random.default_rng(seed)
    meas     = rng.choice(N, size=n_shots, p=probs / probs.sum())
    counts   = np.bincount(meas, minlength=N)

    # Top detected indices
    top_k    = min(10, n_marked * 3)
    top_idx  = list(np.argsort(counts)[::-1][:top_k])
    found    = [i for i in top_idx if i in set(marked_indices)]

    return {'found_indices'  : found,
            'top_candidates' : top_idx,
            'probabilities'  : probs,
            'success_prob'   : success_prob,
            'n_iterations'   : k,
            'total_db_size'  : N}
