"""DecoyFace client-side protection (DGCM) and server-side recovery (ACM).

All projections act on the channel dimension of a C x H x W feature at every
spatial location. C = 64 (F_c output), r = 32, G = 16 groups of size 4.

The RSSD statistics (P_rec, P_N, mu_rec), the fixed channel-shuffle permutation
and the auxiliary complementary bank D_N are built once (see build_rssd.py /
build_bank.py) and loaded into DecoyProtector.
"""
import torch
import torch.nn.functional as F


def proj(P, z):
    """Apply channel-space projection P (C,C) to feature z (B,C,H,W)."""
    b, c, h, w = z.shape
    out = P @ z.reshape(b, c, h * w)
    return out.reshape(b, c, h, w)


def group_shuffle(z, perm, G):
    """CSG: permute channels by `perm`, then split into G contiguous groups.
    z (B,C,H,W) -> (B,G,c_g,H,W)."""
    b, c, h, w = z.shape
    z = z[:, perm]
    return z.reshape(b, G, c // G, h, w)


def group_unshuffle(zg, inv_perm):
    """Inverse CSG: (B,G,c_g,H,W) -> (B,C,H,W) in the original channel order."""
    b, G, cg, h, w = zg.shape
    z = zg.reshape(b, G * cg, h, w)
    return z[:, inv_perm]


def group_instance_norm(zg, eps=1e-5):
    """InstanceNorm applied to every group's channels over the spatial dims."""
    b, G, cg, h, w = zg.shape
    x = zg.reshape(b * G, cg, h, w)
    x = F.instance_norm(x, eps=eps)
    return x.reshape(b, G, cg, h, w)


class DecoyProtector:
    def __init__(self, P_rec, P_N, mu_rec, aux_bank, perm, cfg, device):
        self.device = device
        self.P_rec = P_rec.to(device)
        self.P_N = P_N.to(device)
        self.mu_rec = mu_rec.to(device)          # (C,H,W)
        self.perm = perm.to(device)              # (C,)
        self.inv_perm = torch.argsort(perm).to(device)
        self.cfg = cfg
        self.G = cfg["G"]
        self.K = cfg["K"]
        self.M = aux_bank.shape[0]
        # Pre-group and pre-normalize the auxiliary complementary bank once.
        # aux_bank holds n_j = P_N F_c(y_j); grouped+IN gives the decoy groups.
        ag = group_shuffle(aux_bank.to(device), self.perm, self.G)
        self.aux_grouped_norm = group_instance_norm(ag)   # (M,G,c_g,H,W)

    def _sample_routing(self, B):
        """Per-sample mask m (B,G) and decoy source indices src (B,G)."""
        G, K, M, cfg = self.G, self.K, self.M, self.cfg
        mask = torch.zeros(B, G, dtype=torch.bool)
        src = torch.zeros(B, G, dtype=torch.long)
        for b in range(B):
            keep = torch.randperm(G)[:K]
            mask[b, keep] = True
            g0 = torch.nonzero(~mask[b], as_tuple=False).flatten()
            sc = int(torch.randint(cfg["sc_lo"], cfg["sc_hi"] + 1, (1,)))
            sc = min(sc, g0.numel())
            L = int(torch.randint(cfg["L_lo"], cfg["L_hi"] + 1, (1,)))
            L = min(L, sc) if sc > 0 else L
            perm_g0 = g0[torch.randperm(g0.numel())]
            coherent = perm_g0[:sc]
            independent = perm_g0[sc:]
            shared_pool = torch.randperm(M)[:L]
            for g in coherent.tolist():
                src[b, g] = shared_pool[torch.randint(L, (1,))]
            for g in independent.tolist():
                src[b, g] = torch.randint(M, (1,))
        return mask.to(self.device), src.to(self.device)

    def protect(self, z_A, z_B):
        """Client-side DGCM. z_A, z_B: (B,C,H,W). Returns transmitted z_t."""
        B = z_A.shape[0]
        r_B = proj(self.P_rec, z_B)
        n_A = proj(self.P_N, z_A)

        uA = group_instance_norm(group_shuffle(n_A, self.perm, self.G))
        mask, src = self._sample_routing(B)
        ar = torch.arange(self.G, device=self.device)
        decoy = self.aux_grouped_norm[src, ar]            # (B,G,c_g,H,W)
        m = mask[:, :, None, None, None]
        mixed = torch.where(m, uA, decoy)
        n_tilde = group_unshuffle(mixed, self.inv_perm)

        z_tilde = n_tilde + r_B
        k = torch.randint(-self.cfg["t"], self.cfg["t"] + 1, (B, 1, 1, 1),
                          device=self.device)
        z_t = (2.0 ** k.float()) * z_tilde
        return z_t

    def server_recover(self, z_t):
        """Server-side ACM. z_t: (B,C,H,W). Returns canonicalized z_can."""
        z_s = proj(self.P_N, z_t)
        zg = group_shuffle(z_s, self.perm, self.G)        # (B,G,c_g,H,W)
        b, G, cg, h, w = zg.shape
        flat = zg.reshape(b, G, cg * h * w).abs()
        q = torch.quantile(flat, self.cfg["q"], dim=2)    # (B,G)
        scale = self.cfg["tau_s"] / (q + self.cfg["eps"])
        zg = zg * scale[:, :, None, None, None]
        n_cal = group_unshuffle(zg, self.inv_perm)
        return self.mu_rec[None] + n_cal