"""
dps_hierarchy.py

General DPS (Doherty-Parrilo-Spedalieri) level-k symmetric-extension SDP,
for a bipartite qudit state rho_AB with local dimension d, extending party
B to k Bose-symmetric copies.

Key design point (discussed at length in the chat this was extracted
from): the SDP *variable* sigma is parametrized directly on
A x Sym^k(C^d), dimension d * C(d+k-1,k) -- POLYNOMIAL in k. But the PPT
constraint (sigma^{T_A} >= 0) must be checked on the full, unsymmetrized
embedding A x B_1 x ... x B_k, dimension d^{k+1} -- EXPONENTIAL in k. So
this construction saves on free parameters but NOT on the size of the
PSD cone that actually drives SDP solve time. See the README for measured
timings (k=2: ~27x27 cone, sub-second; k=3: ~81x81 cone, ~2-4 minutes
with SCS in the original sandbox).

Requires: numpy, cvxpy.
"""
import math
from itertools import permutations
import numpy as np
import cvxpy as cp


def sym_isometry(d, k):
    """Isometry W, shape (d**k, dim Sym^k(C^d)), spanning the totally
    symmetric subspace of (C^d)^{tensor k}. Built by brute-force averaging
    over all k! permutations of the k tensor factors -- fine for k up to
    ~6-7; for larger k this construction itself becomes the bottleneck,
    independently of the SDP."""
    n = d ** k
    P = np.zeros((n, n))
    for perm in permutations(range(k)):
        M = np.zeros((n, n))
        for idx in np.ndindex(*([d] * k)):
            new_idx = tuple(idx[perm[i]] for i in range(k))
            row = 0
            col = 0
            for i in range(k):
                row = row * d + new_idx[i]
                col = col * d + idx[i]
            M[row, col] = 1
        P += M
    P /= math.factorial(k)
    eigvals, eigvecs = np.linalg.eigh(P)
    cols = [eigvecs[:, i] for i in range(n) if abs(eigvals[i] - 1) < 1e-9]
    return np.column_stack(cols)


def partial_trace_keep_first_copy(full_expr, d, k):
    """full_expr indexed by (a, b_1, ..., b_k) with combined index
    a*d**k + b_1*d**(k-1) + ... + b_k. Traces out b_2..b_k, keeping (a,b_1)
    -- i.e. returns the marginal on A x (first copy of B)."""
    rest_dim = d ** (k - 1)
    rows = []
    for a in range(d):
        for b1 in range(d):
            row = []
            for ap in range(d):
                for b1p in range(d):
                    terms = [full_expr[(a * d + b1) * rest_dim + r,
                                        (ap * d + b1p) * rest_dim + r]
                             for r in range(rest_dim)]
                    row.append(sum(terms))
            rows.append(row)
    return cp.bmat(rows)


def partial_transpose_first_system(full_expr, d1, d2):
    """Partial transpose on the first (d1-dim) system of a
    (d1*d2) x (d1*d2) matrix. Has the same eigenvalues as transposing the
    second system instead (standard fact: M^{T_A} and M^{T_B} always share
    a spectrum, since M^{T_B} = (M^{T_A})^T)."""
    rows = []
    for i in range(d1):
        for kk in range(d2):
            row = []
            for j in range(d1):
                for l in range(d2):
                    row.append(full_expr[j * d2 + kk, i * d2 + l])
            rows.append(row)
    return cp.bmat(rows)


def build_dps_problem(d, k):
    """Returns (prob, rho_param, sigma) for the level-k DPS feasibility
    SDP. Set rho_param.value = <(d*d)x(d*d) target state>, then call
    dps_feasible(...) or prob.solve(...) directly."""
    W = sym_isometry(d, k)
    dim_sym = W.shape[1]
    Iso = np.kron(np.eye(d), W)  # d**(k+1)  x  (d * dim_sym)
    sigma = cp.Variable((d * dim_sym, d * dim_sym), hermitian=True)
    full = Iso @ sigma @ Iso.conj().T
    ptrace = partial_trace_keep_first_copy(full, d, k)
    pt = partial_transpose_first_system(full, d1=d, d2=d ** k)
    rho_param = cp.Parameter((d * d, d * d), hermitian=True)
    constraints = [sigma >> 0, cp.trace(sigma) == 1,
                   ptrace == rho_param, pt >> 0]
    prob = cp.Problem(cp.Minimize(0), constraints)
    return prob, rho_param, sigma


def dps_feasible(prob, rho_param, rho_target, solver=cp.SCS, **solve_kwargs):
    """Solve the (already-built) DPS problem for a given target state and
    return True iff a valid extension was found (i.e. rho_target is NOT
    certified entangled at this level)."""
    rho_param.value = rho_target
    prob.solve(solver=solver, **solve_kwargs)
    return prob.status in ("optimal", "optimal_inaccurate")