
"""
PKTron Finance — QAOA Portfolio Optimization
============================================
Uses the Quantum Approximate Optimization Algorithm (QAOA) to find the
optimal binary asset-allocation vector that maximises risk-adjusted return.

Memory profile (12 GB Colab):
  - All statevectors stored as complex64  (8 bytes / element)
  - n_assets <= 20 is safe (2^20 * 8 = 8 MB)
  - Default demo: 6 assets
"""
import numpy as np
from itertools import product as iproduct

_C = np.complex64
_R = np.float32

# ─────────────────────────────────────────────────────────────────────────
#  Low-level statevector helpers (memory-safe, complex64)
# ─────────────────────────────────────────────────────────────────────────

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

def _hadamard_all(n):
    """Return uniform superposition |+>^n."""
    s = np.ones(2**n, dtype=_C) / _R(np.sqrt(2**n))
    return s

def _apply_rzz(state, qi, qj, angle):
    """e^{-i angle/2 Z_i Z_j} diagonal gate."""
    dim = len(state)
    new = state.copy()
    for s in range(dim):
        bi = (s >> qi) & 1
        bj = (s >> qj) & 1
        parity = 1 if bi == bj else -1
        new[s] = state[s] * np.exp(-1j * _R(angle) / 2 * parity, dtype=_C)
    return new

def _apply_rx(state, q, angle):
    """Rx gate on qubit q."""
    c  = _R(np.cos(angle / 2))
    si = _C(-1j * np.sin(angle / 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] + si * state[b]
            new[b] = si * state[a] + _C(c) * state[b]
    return new

# ─────────────────────────────────────────────────────────────────────────
#  QAOA Cost Hamiltonian for Portfolio
# ─────────────────────────────────────────────────────────────────────────

def _portfolio_energy(bitstring, returns, covariance, risk_penalty=1.0):
    """
    E(x) = -r^T x  +  risk_penalty * x^T Sigma x
    x in {0,1}^n  (1 = include asset)
    """
    x = np.array(bitstring, dtype=_R)
    return _R(-returns @ x + risk_penalty * x @ covariance @ x)

def _expectation_cost(state, returns, covariance, risk_penalty=1.0):
    """<state|H_C|state> for the portfolio cost Hamiltonian."""
    n = int(np.log2(len(state)))
    E = _R(0.0)
    for idx in range(len(state)):
        p = float(np.abs(state[idx])**2)
        if p < 1e-12:
            continue
        bits = [(idx >> q) & 1 for q in range(n)]
        E += _R(p) * _portfolio_energy(bits, returns, covariance, risk_penalty)
    return float(E)

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

def build_portfolio_qaoa(returns, covariance, p_layers=2, risk_penalty=1.0):
    """
    Build a QAOA ansatz closure for portfolio optimisation.

    Parameters
    ----------
    returns     : array-like, shape (n,)   expected return per asset
    covariance  : array-like, shape (n,n)  covariance matrix
    p_layers    : int   number of QAOA repetitions (depth)
    risk_penalty: float Lagrange multiplier on variance term

    Returns
    -------
    (n_assets, qaoa_fn)
        qaoa_fn(params) → (energy, state)
        params has length 2*p_layers (gammas + betas)
    """
    returns    = np.array(returns,    dtype=_R)
    covariance = np.array(covariance, dtype=_R)
    n          = len(returns)

    def qaoa_fn(params):
        gammas = params[:p_layers]
        betas  = params[p_layers:]
        state  = _hadamard_all(n)

        for layer in range(p_layers):
            # Cost unitary  e^{-i gamma H_C}
            for i in range(n):
                for j in range(i + 1, n):
                    w = _R(covariance[i, j] * risk_penalty)
                    state = _apply_rzz(state, i, j, float(w * gammas[layer]))
                # Diagonal term from -returns[i]
                # Implemented as Z_i rotation
                for s in range(len(state)):
                    bi = (s >> i) & 1
                    phase = _R(-returns[i] * gammas[layer] * (1 - 2*bi))
                    state[s] = state[s] * np.exp(1j * phase / 2, dtype=_C)

            # Mixer unitary  e^{-i beta B}  where B = sum_i X_i
            for i in range(n):
                state = _apply_rx(state, i, float(2 * betas[layer]))

        energy = _expectation_cost(state, returns, covariance, risk_penalty)
        return energy, state

    return n, qaoa_fn


def optimize_portfolio(returns, covariance, p_layers=2, n_restarts=5,
                       risk_penalty=1.0, seed=42):
    """
    Run classical optimiser over QAOA parameters and decode result.

    Returns
    -------
    dict with keys:
        selected_assets  : list[int]   indices of chosen assets
        expected_return  : float
        portfolio_risk   : float
        qaoa_energy      : float
        optimal_params   : np.ndarray
    """
    from scipy.optimize import minimize as sp_min

    returns    = np.array(returns,    dtype=_R)
    covariance = np.array(covariance, dtype=_R)
    n, qaoa_fn = build_portfolio_qaoa(returns, covariance, p_layers, risk_penalty)

    rng     = np.random.default_rng(seed)
    best_e  = np.inf
    best_p  = None

    for _ in range(n_restarts):
        p0   = rng.uniform(0, 2*np.pi, 2*p_layers).astype(_R)
        res  = sp_min(lambda p: qaoa_fn(p)[0], p0,
                      method='COBYLA',
                      options={'maxiter': 200, 'rhobeg': 0.5})
        if res.fun < best_e:
            best_e = res.fun
            best_p = res.x

    energy, state = qaoa_fn(best_p)
    # Decode: pick most probable bitstring
    probs     = np.abs(state)**2
    best_idx  = int(np.argmax(probs))
    selection = [(best_idx >> i) & 1 for i in range(n)]

    sel_assets = [i for i, v in enumerate(selection) if v]
    exp_ret    = float(returns[sel_assets].sum()) if sel_assets else 0.0
    cov_sub    = covariance[np.ix_(sel_assets, sel_assets)] if sel_assets else np.zeros((1,1))
    port_risk  = float(np.ones(len(sel_assets)) @ cov_sub @ np.ones(len(sel_assets))) if sel_assets else 0.0

    return {
        'selected_assets': sel_assets,
        'expected_return': exp_ret,
        'portfolio_risk' : port_risk,
        'qaoa_energy'    : float(energy),
        'optimal_params' : best_p,
        'probabilities'  : probs,
    }
