
"""
PKTron Finance — VQC Fraud Detection
=====================================
Variational Quantum Classifier (VQC) for detecting fraudulent transactions.

Workflow:
  1. Encode transaction features as qubit rotation angles (angle encoding)
  2. Apply parameterised hardware-efficient ansatz
  3. Measure expectation value of Z_0 as class score
  4. Train with binary cross-entropy via COBYLA

Memory safe: float32; n_features <= 8 qubits recommended for 12 GB Colab.
"""
import numpy as np
from scipy.optimize import minimize as sp_min

_R = np.float32
_C = np.complex64

# ── Gate helpers ─────────────────────────────────────────────────────────

def _zero_state(n):
    s = np.zeros(2**n, dtype=_C); s[0] = 1.0; return s

def _ry(state, q, t):
    c, s_ = _R(np.cos(t/2)), _R(np.sin(t/2))
    dim = len(state); step = 1<<q; new = state.copy()
    for i in range(0, dim, 2*step):
        for j in range(step):
            a, b = i+j, i+j+step
            new[a] = _C(c)*state[a] - _C(s_)*state[b]
            new[b] = _C(s_)*state[a] + _C(c)*state[b]
    return new

def _rz(state, q, t):
    step=1<<q; dim=len(state); new=state.copy()
    ep=np.exp( 1j*_R(t/2), dtype=_C)
    em=np.exp(-1j*_R(t/2), dtype=_C)
    for i in range(0, dim, 2*step):
        for j in range(step):
            new[i+j]=em*state[i+j]; new[i+j+step]=ep*state[i+j+step]
    return new

def _cnot(state, ctrl, tgt):
    dim=len(state); new=state.copy()
    for s in range(dim):
        if (s>>ctrl)&1: new[s]=state[s^(1<<tgt)]
    return new

def _expect_z0(state):
    """<Z_0> = P(qubit0=0) - P(qubit0=1)"""
    probs = np.abs(state)**2
    p0 = float(probs[0::2].sum())
    p1 = float(probs[1::2].sum())
    return _R(p0 - p1)

# ── VQC Components ───────────────────────────────────────────────────────

def build_vqc(n_features, depth=2):
    """
    Build VQC components for fraud detection.

    Parameters
    ----------
    n_features : int   number of transaction features (= number of qubits)
    depth      : int   ansatz depth (layers of Ry+Rz+CNOT)

    Returns
    -------
    dict: n_params, circuit_fn, n_qubits
    circuit_fn(x_encoded, params) -> float  (Z_0 expectation in [-1,+1])
    """
    n = n_features
    n_params = n * depth * 2   # Ry + Rz per qubit per layer

    def encode(x):
        """Angle encoding: Ry(pi * x_i) on each qubit."""
        state = _zero_state(n)
        for q in range(n):
            state = _ry(state, q, float(np.pi * x[q]))
        return state

    def ansatz(state, params):
        idx = 0
        for d in range(depth):
            for q in range(n):
                state = _ry(state, q, float(params[idx])); idx += 1
                state = _rz(state, q, float(params[idx])); idx += 1
            for q in range(n-1):
                state = _cnot(state, q, q+1)
        return state

    def circuit_fn(x, params):
        state = encode(x)
        state = ansatz(state, params)
        return float(_expect_z0(state))

    return {'n_params': n_params, 'circuit_fn': circuit_fn,
            'n_qubits': n, 'depth': depth}


def _sigmoid(x):
    return _R(1.0 / (1.0 + np.exp(-x)))

def _bce(y_pred, y_true, eps=1e-7):
    y_pred = np.clip(y_pred, eps, 1-eps)
    return -float(np.mean(y_true * np.log(y_pred)
                          + (1-y_true) * np.log(1-y_pred)))


def train_vqc(data, labels, n_features=None, depth=2,
              max_iter=200, seed=0):
    """
    Train VQC on transaction data.

    Parameters
    ----------
    data       : array (n_samples, n_features), values in [0,1]
    labels     : array (n_samples,), 0=legit 1=fraud
    n_features : int or None (inferred from data)
    depth      : int  ansatz depth
    max_iter   : int  COBYLA iterations

    Returns
    -------
    dict: params, train_accuracy, vqc_info
    """
    data   = np.array(data,   dtype=_R)
    labels = np.array(labels, dtype=_R)
    n_f    = data.shape[1] if n_features is None else n_features
    vqc    = build_vqc(n_f, depth)
    cfn    = vqc['circuit_fn']
    rng    = np.random.default_rng(seed)
    p0     = rng.uniform(0, 2*np.pi, vqc['n_params']).astype(_R)

    def loss(params):
        scores = np.array([cfn(x, params) for x in data], dtype=_R)
        probs  = _sigmoid(scores)        # map [-1,1] -> (0,1)
        return _bce(probs, labels)

    res    = sp_min(loss, p0, method='COBYLA',
                    options={'maxiter': max_iter, 'rhobeg': 0.5})
    params = res.x

    # Compute train accuracy
    scores   = np.array([cfn(x, params) for x in data], dtype=_R)
    preds    = (scores > 0).astype(int)
    accuracy = float((preds == labels.astype(int)).mean())

    return {'params': params, 'train_accuracy': accuracy,
            'loss': float(res.fun), 'vqc_info': vqc}


def predict_transaction(transaction, trained_result):
    """
    Classify a single transaction as fraudulent (1) or legitimate (0).

    Parameters
    ----------
    transaction    : array-like, shape (n_features,), values in [0,1]
    trained_result : dict returned by train_vqc()

    Returns
    -------
    dict: label (0/1), score (raw Z_0), probability
    """
    x      = np.array(transaction, dtype=_R)
    cfn    = trained_result['vqc_info']['circuit_fn']
    params = trained_result['params']
    score  = float(cfn(x, params))
    prob   = float(_sigmoid(_R(score)))
    label  = 1 if score > 0 else 0
    return {'label': label, 'score': score,
            'fraud_probability': prob,
            'verdict': 'FRAUD' if label == 1 else 'LEGITIMATE'}
