
"""
PKTron Finance — VQE Portfolio Optimization
==========================================
Variational Quantum Eigensolver for portfolio cost Hamiltonian.
Parameterised hardware-efficient ansatz: Ry layers + CNOT entanglement.
"""
import numpy as np
from scipy.optimize import minimize as sp_min

_C = np.complex64
_R = np.float32

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

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

def _ry(state, q, theta):
    c, s_ = _R(np.cos(theta/2)), _R(np.sin(theta/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, theta):
    step = 1 << q; dim = len(state); new = state.copy()
    ep = np.exp( 1j*_R(theta/2), dtype=_C)
    em = np.exp(-1j*_R(theta/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

# ── Ansatz ───────────────────────────────────────────────────────────────

def create_ansatz(n_assets, depth=2):
    """
    Hardware-efficient ansatz: depth * [Ry(t) layer + linear CNOT chain].
    Returns (n_params, ansatz_fn).
    ansatz_fn(params) -> statevector complex64
    """
    n_params = n_assets * depth * 2   # Ry + Rz per qubit per layer

    def ansatz_fn(params):
        state = _zero_state(n_assets)
        idx   = 0
        for d in range(depth):
            for q in range(n_assets):
                state = _ry(state, q, float(params[idx]));   idx += 1
                state = _rz(state, q, float(params[idx]));   idx += 1
            # Linear entanglement
            for q in range(n_assets - 1):
                state = _cnot(state, q, q+1)
        return state

    return n_params, ansatz_fn

# ── Cost function ────────────────────────────────────────────────────────

def cost_function(params, returns, covariance, ansatz_fn, risk_penalty=1.0):
    """
    <psi(params)| H_portfolio |psi(params)>
    H = -sum_i r_i Z_i  +  risk_penalty * sum_{ij} Sigma_{ij} Z_i Z_j
    (Z eigenvalue +1 for |0>, -1 for |1>)
    """
    returns    = np.array(returns,    dtype=_R)
    covariance = np.array(covariance, dtype=_R)
    n          = len(returns)
    state      = ansatz_fn(params)
    probs      = np.abs(state)**2
    E          = _R(0.0)
    for s in range(len(state)):
        p = _R(probs[s])
        if p < 1e-12:
            continue
        z = np.array([1 - 2*((s >> q) & 1) for q in range(n)], dtype=_R)
        E_s = -float(returns @ z) + risk_penalty * float(z @ covariance @ z)
        E  += p * _R(E_s)
    return float(E)

# ── Runner ───────────────────────────────────────────────────────────────

def run_vqe_portfolio(returns, covariance, depth=2, risk_penalty=1.0,
                      n_restarts=5, seed=7):
    """
    Minimise portfolio cost Hamiltonian via VQE.

    Returns
    -------
    dict:
        selected_assets, expected_return, portfolio_risk,
        vqe_energy, optimal_params, statevector
    """
    returns    = np.array(returns,    dtype=_R)
    covariance = np.array(covariance, dtype=_R)
    n          = len(returns)
    n_params, ansatz_fn = create_ansatz(n, depth)
    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, n_params).astype(_R)
        res = sp_min(lambda p: cost_function(p, returns, covariance,
                                             ansatz_fn, risk_penalty),
                     p0, method='COBYLA',
                     options={'maxiter': 300, 'rhobeg': 0.5})
        if res.fun < best_e:
            best_e = res.fun; best_p = res.x

    state     = ansatz_fn(best_p)
    probs     = np.abs(state)**2
    best_idx  = int(np.argmax(probs))
    selection = [(best_idx >> i) & 1 for i in range(n)]

    sel = [i for i, v in enumerate(selection) if v]
    er  = float(returns[sel].sum())  if sel else 0.0
    cov_s = covariance[np.ix_(sel, sel)] if sel else np.zeros((1,1))
    risk  = float(np.ones(len(sel)) @ cov_s @ np.ones(len(sel))) if sel else 0.0

    return {'selected_assets': sel, 'expected_return': er,
            'portfolio_risk': risk, 'vqe_energy': float(best_e),
            'optimal_params': best_p, 'statevector': state}
