import numpy as np
import torch
from torch.autograd import Function
import scipy.sparse as sp
import osqp

device = "cuda" if torch.cuda.is_available() else "cpu"

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

def osqp_solve(P_csc, q_np, A_csc, l_np, u_np, eps=1e-6):
    prob = osqp.OSQP()
    prob.setup(P_csc, q_np, A_csc, l_np, u_np, verbose=False,
               eps_abs=eps, eps_rel=eps, eps_prim_inf=eps, eps_dual_inf=eps)
    res = prob.solve()
    if res.x is None:
        raise RuntimeError(res.info.status)
    return res.x.astype(np.float64), res.y.astype(np.float64)

def pack_osqp(P, q, G, h, A, b):
    Pn, qn, Gn, hn, An, bn = [_np(x) for x in [P, q, G, h, A, b]]
    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)
    m, p = Gn.shape[0], An.shape[0]
    if p > 0:
        Aos = sp.csc_matrix(np.vstack([Gn, An]))
        l = np.hstack([-np.inf*np.ones(m), bn])
        u = np.hstack([hn, bn])
    else:
        Aos = sp.csc_matrix(Gn)
        l = -np.inf*np.ones(m)
        u = hn
    return sp.csc_matrix(Pn), qn, Aos, l.astype(np.float64), u.astype(np.float64), m, p, Gn, hn, An

def bpqp_backward(xn, yn, P_csc, Gn, hn, An, gnp, m, p, act_tol, backward_eps):
    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:
        Pd = P_csc.toarray()
        try: z = -np.linalg.solve(Pd, gnp)
        except np.linalg.LinAlgError: z = -np.linalg.solve(Pd + 1e-8*np.eye(n), gnp)
        yb = np.zeros((0,), dtype=np.float64)
        return z, yb, active
    Ab = sp.csc_matrix(np.vstack(rows))
    k = Ab.shape[0]
    z, yb = osqp_solve(P_csc, gnp.astype(np.float64), Ab, np.zeros(k), np.zeros(k), eps=backward_eps)
    return z, yb, active

def BPQPLayer(sign=1, act_tol=1e-6, forward_eps=1e-6, backward_eps=1e-10):
    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
                P_csc, qn, Aos, l, u, m, p, *_ = pack_osqp(Pi, sign * qi, Gi, hi, Ai, bi)
                x, y = osqp_solve(P_csc, qn, Aos, l, u, eps=forward_eps)
                xs.append(torch.from_numpy(x).to(device=Pi.device, dtype=Pi.dtype))
                ys.append(torch.from_numpy(y).to(device=Pi.device, dtype=Pi.dtype))
                ms.append(m); ps.append(p)
            x = torch.stack(xs, 0) if batched else xs[0]
            y = torch.stack(ys, 0) if batched else ys[0]
            ctx.save_for_backward(P, q, G, h, A, b, x, y)
            ctx.meta = (batched, B, sign, act_tol, forward_eps, backward_eps, ms, ps)

            return x

        @staticmethod
        def backward(ctx, grad_output):
            P, q, G, h, A, b, x, y = ctx.saved_tensors
            batched, B, sign, act_tol, forward_eps, backward_eps, 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
                P_csc, _, _, _, _, m, p, Gn, hn, An = pack_osqp(Pi, sign * qi, Gi, hi, Ai, bi)
                z, yb, active = bpqp_backward(_np(xi), _np(yi), P_csc, Gn, hn, An, _np(gi), m, p, act_tol, backward_eps)
                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