import numpy as np
import torch
from torch.autograd import Function
import cvxpy as cp

def _np(x): return x.detach().cpu().numpy()
def _sym(P): return 0.5 * (P + P.T)

CACHE = True
_QP_CACHE = {}
_EQ_CACHE = {}

def _qp_cvx_osqp(Pn, qn, Gn, hn, An, bn, eps_abs=1e-5, eps_rel=1e-5, max_iter=10000):
    Pn = _sym(Pn).astype(np.float64)
    qn = qn.reshape(-1).astype(np.float64)
    Gn = Gn.astype(np.float64); hn = hn.reshape(-1).astype(np.float64)
    An = An.astype(np.float64); bn = bn.reshape(-1).astype(np.float64)

    n = Pn.shape[0]
    m = Gn.shape[0]
    p = An.shape[0]
    key = (n, m, p)

    if (not CACHE) or (key not in _QP_CACHE):
        x = cp.Variable(n)
        Pp = cp.Parameter((n, n), symmetric=True)
        qp = cp.Parameter(n)
        cons = []
        ineq = eq = None
        Gp = hp = Ap = bp = None
        if m > 0:
            Gp = cp.Parameter((m, n))
            hp = cp.Parameter(m)
            ineq = (Gp @ x <= hp)
            cons.append(ineq)
        if p > 0:
            Ap = cp.Parameter((p, n))
            bp = cp.Parameter(p)
            eq = (Ap @ x == bp)
            cons.append(eq)
        obj = cp.Minimize(0.5 * cp.quad_form(x, cp.psd_wrap(Pp)) + qp @ x)
        prob = cp.Problem(obj, cons)
        bundle = {"x": x, "Pp": Pp, "qp": qp, "Gp": Gp, "hp": hp, "Ap": Ap, "bp": bp, "ineq": ineq, "eq": eq, "prob": prob}
        if CACHE:
            _QP_CACHE[key] = bundle
    else:
        bundle = _QP_CACHE[key]

    bundle["Pp"].value = Pn
    bundle["qp"].value = qn
    if m > 0:
        bundle["Gp"].value = Gn
        bundle["hp"].value = hn
    if p > 0:
        bundle["Ap"].value = An
        bundle["bp"].value = bn

    bundle["prob"].solve(
        solver=cp.OSQP,
        warm_start=True,
        verbose=False,
        eps_abs=eps_abs,
        eps_rel=eps_rel,
        max_iter=max_iter,
        polish=True,
    )
    st = bundle["prob"].status
    if st not in (cp.OPTIMAL, cp.OPTIMAL_INACCURATE):
        raise RuntimeError(f"CVXPY/OSQP status: {st}")

    xval = np.asarray(bundle["x"].value, dtype=np.float64).reshape(-1)
    lam = np.asarray(bundle["ineq"].dual_value, dtype=np.float64).reshape(-1) if m > 0 else np.zeros((0,), dtype=np.float64)
    nu = np.asarray(bundle["eq"].dual_value, dtype=np.float64).reshape(-1) if p > 0 else np.zeros((0,), dtype=np.float64)
    y = np.concatenate([lam, nu], axis=0)
    return xval, y

def _eq_qp_cvx_osqp(Pn, gnp, Ab, eps_abs=1e-5, eps_rel=1e-5, max_iter=10000):
    Pn = _sym(Pn).astype(np.float64)
    gnp = gnp.reshape(-1).astype(np.float64)
    Ab = Ab.astype(np.float64)

    k, n = Ab.shape
    key = (n, k)

    if (not CACHE) or (key not in _EQ_CACHE):
        z = cp.Variable(n)
        Pp = cp.Parameter((n, n), symmetric=True)
        gp = cp.Parameter(n)
        Ap = cp.Parameter((k, n))
        con = (Ap @ z == 0)
        obj = cp.Minimize(0.5 * cp.quad_form(z, cp.psd_wrap(Pp)) + gp @ z)
        prob = cp.Problem(obj, [con])
        bundle = {"z": z, "Pp": Pp, "gp": gp, "Ap": Ap, "con": con, "prob": prob}
        if CACHE:
            _EQ_CACHE[key] = bundle
    else:
        bundle = _EQ_CACHE[key]

    bundle["Pp"].value = Pn
    bundle["gp"].value = gnp
    bundle["Ap"].value = Ab

    bundle["prob"].solve(
        solver=cp.OSQP,
        warm_start=True,
        verbose=False,
        eps_abs=eps_abs,
        eps_rel=eps_rel,
        max_iter=max_iter,
        polish=True,
    )
    st = bundle["prob"].status
    if st not in (cp.OPTIMAL, cp.OPTIMAL_INACCURATE):
        raise RuntimeError(f"CVXPY/OSQP status: {st}")

    zval = np.asarray(bundle["z"].value, dtype=np.float64).reshape(-1)
    yb = np.asarray(bundle["con"].dual_value, dtype=np.float64).reshape(-1)
    return zval, yb

def bpqp_backward(xn, yn, Pn, Gn, hn, An, gnp, m, p, act_tol):
    lam = yn[:m] if m > 0 else np.zeros((0,), dtype=np.float64)
    if m > 0:
        resid = Gn @ xn - hn
        active = np.where((resid > -act_tol) | (lam > act_tol))[0].astype(np.int64)
    else:
        active = np.zeros((0,), dtype=np.int64)

    rows = []
    if active.size > 0:
        rows.append(Gn[active, :])
    if p > 0:
        rows.append(An)

    n = xn.size
    if len(rows) == 0:
        try:
            z = -np.linalg.solve(Pn, gnp)
        except np.linalg.LinAlgError:
            z = -np.linalg.solve(Pn + 1e-8 * np.eye(n), gnp)
        return z, np.zeros((0,), dtype=np.float64), active

    Ab = np.vstack(rows)
    z, yb = _eq_qp_cvx_osqp(Pn, gnp.astype(np.float64), Ab)
    return z, yb, active

def BPQPLayer(sign=1, act_tol=1e-6):
    class _Layer(Function):
        @staticmethod
        def forward(ctx, P, q, G, h, A, b):
            batched = (P.dim() == 3)
            B = P.shape[0] if batched else 1
            xs, ys, ms, ps = [], [], [], []

            for i in range(B):
                Pi = P[i] if batched else P
                qi = q[i] if q.dim() == 2 else q
                Gi = G[i] if G.dim() == 3 else G
                hi = h[i] if h.dim() == 2 else h
                Ai = A[i] if A.dim() == 3 else A
                bi = b[i] if b.dim() == 2 else b

                Pn = _sym(_np(Pi)).astype(np.float64)
                qn = (sign * _np(qi)).reshape(-1).astype(np.float64)
                Gn = _np(Gi).astype(np.float64)
                hn = _np(hi).reshape(-1).astype(np.float64)
                An = _np(Ai).astype(np.float64)
                bn = _np(bi).reshape(-1).astype(np.float64)

                m, p = Gn.shape[0], An.shape[0]
                x_np, y_np = _qp_cvx_osqp(Pn, qn, Gn, hn, An, bn)

                xs.append(torch.from_numpy(x_np).to(device=Pi.device, dtype=Pi.dtype))
                ys.append(torch.from_numpy(y_np).to(device=Pi.device, dtype=Pi.dtype))
                ms.append(m); ps.append(p)

            x_out = torch.stack(xs, 0) if batched else xs[0]
            y_out = torch.stack(ys, 0) if batched else ys[0]

            ctx.save_for_backward(P, q, G, h, A, b, x_out, y_out)
            ctx.meta = (batched, B, sign, act_tol, ms, ps)
            return x_out

        @staticmethod
        def backward(ctx, grad_output):
            P, q, G, h, A, b, x, y = ctx.saved_tensors
            batched, B, sign, act_tol, ms, ps = ctx.meta

            gP = torch.zeros_like(P); gq = torch.zeros_like(q); gG = torch.zeros_like(G)
            gh = torch.zeros_like(h); gA = torch.zeros_like(A); gb = torch.zeros_like(b)

            for i in range(B):
                Pi = P[i] if batched else P
                qi = q[i] if q.dim() == 2 else q
                hi = h[i] if h.dim() == 2 else h
                Gi = G[i] if G.dim() == 3 else G
                Ai = A[i] if A.dim() == 3 else A
                bi = b[i] if b.dim() == 2 else b

                xi = x[i] if batched else x
                yi = y[i] if batched else y
                gi = grad_output[i] if batched else grad_output

                Pn = _sym(_np(Pi)).astype(np.float64)
                Gn = _np(Gi).astype(np.float64)
                hn = _np(hi).reshape(-1).astype(np.float64)
                An = _np(Ai).astype(np.float64)
                m, p = Gn.shape[0], An.shape[0]

                z, yb, active = bpqp_backward(
                    _np(xi), _np(yi), Pn, Gn, hn, An, _np(gi), m, p, act_tol
                )

                zt = torch.from_numpy(z).to(device=Pi.device, dtype=Pi.dtype)
                gq_i = sign * zt
                gP_i = 0.5 * (torch.outer(zt, xi) + torch.outer(xi, zt))

                lam = yi[:m] if m > 0 else torch.empty((0,), device=Pi.device, dtype=Pi.dtype)
                nu = yi[m:m+p] if p > 0 else torch.empty((0,), device=Pi.device, dtype=Pi.dtype)

                k = int(active.size)
                mu = torch.from_numpy(yb[:k]).to(device=Pi.device, dtype=Pi.dtype) if k > 0 else torch.empty((0,), device=Pi.device, dtype=Pi.dtype)
                eta = torch.from_numpy(yb[k:k+p]).to(device=Pi.device, dtype=Pi.dtype) if p > 0 else torch.empty((0,), device=Pi.device, dtype=Pi.dtype)

                gG_i = torch.zeros_like(Gi); gh_i = torch.zeros_like(hi)
                if m > 0 and k > 0:
                    at = torch.tensor(active, device=Pi.device, dtype=torch.long)
                    lam_act = lam.index_select(0, at)
                    block = mu[:, None] * xi[None, :] + lam_act[:, None] * zt[None, :]
                    gG_i.index_copy_(0, at, block)
                    gh_i.index_copy_(0, at, -mu)

                if p > 0:
                    gb_i = -eta
                    gA_i = eta[:, None] * xi[None, :] + nu[:, None] * zt[None, :]
                else:
                    gb_i = torch.zeros_like(bi)
                    gA_i = torch.zeros_like(Ai)

                if batched: gP[i] = gP_i
                else: gP = gP + gP_i
                if q.dim() == 2: gq[i] = gq_i
                else: gq = gq + gq_i
                if G.dim() == 3: gG[i] = gG_i
                else: gG = gG + gG_i
                if h.dim() == 2: gh[i] = gh_i
                else: gh = gh + gh_i
                if A.dim() == 3: gA[i] = gA_i
                else: gA = gA + gA_i
                if b.dim() == 2: gb[i] = gb_i
                else: gb = gb + gb_i

            return gP, gq, gG, gh, gA, gb

    return _Layer.apply