"""Vision connector: SigLIP2 NaFlex patch features [B,576,768] -> soft tokens [B,144 (or 64),1536] for the frozen Jais-3-0.9B.
pixel-shuffle 2x2 over the 24x24 grid (576x768 -> 144x3072) -> Linear(3072,1536) -> GELU -> Linear(1536,1536) -> per-token L2 normalise * learnable scalar
(initialised to the Jais embedding mean row norm 0.896, so soft tokens start at the scale of real text embeddings).  7.08M parameters.
low-token mode (64 tokens): the 12x12 shuffled grid is adaptively average-pooled to 8x8 (same MLP, so one connector serves both modes)."""
import torch, torch.nn as nn, torch.nn.functional as F

EMBED_NORM = 0.896          # mean row norm of Jais embed_tokens (design doc, section 2.2)


def pixel_shuffle(x, grid=24, r=2):
    """[B, grid*grid, D] -> [B, (grid/r)^2, D*r*r]; each output token concatenates a rxr block of neighbouring patches"""
    B, N, D = x.shape
    assert N == grid * grid, f"expected {grid*grid} patches, got {N}"
    g = grid // r
    x = x.view(B, g, r, g, r, D).permute(0, 1, 3, 2, 4, 5)
    return x.reshape(B, g * g, r * r * D)


class Connector(nn.Module):
    def __init__(self, in_dim=768, hidden=1536, out_dim=1536, grid=24, r=2, target_norm=EMBED_NORM, low_grid=8):
        super().__init__()
        self.grid, self.r, self.low_grid = grid, r, low_grid
        self.fc1 = nn.Linear(in_dim * r * r, hidden)
        self.fc2 = nn.Linear(hidden, out_dim)
        self.scale = nn.Parameter(torch.tensor(float(target_norm)))     # learnable output norm

    def tokens(self, feats, low=False):
        """feats [B,576,768] -> shuffled tokens [B,144,3072] (or [B,64,3072] in low mode)"""
        x = pixel_shuffle(feats, self.grid, self.r)
        if low:
            g = self.grid // self.r
            x = x.transpose(1, 2).reshape(x.shape[0], -1, g, g)
            x = F.adaptive_avg_pool2d(x, self.low_grid).flatten(2).transpose(1, 2)
        return x

    def forward(self, feats, low=False):
        x = self.tokens(feats.to(self.fc1.weight.dtype), low)
        y = self.fc2(F.gelu(self.fc1(x)))
        # normalise in fp32 for stability under bf16 autocast, then rescale to the embedding norm
        return (F.normalize(y.float(), dim=-1) * self.scale.float()).to(y.dtype)

    # ---------------------------------------------------------------- PORTAL-style warm start (stage-1 ablation; default OFF)
    @torch.no_grad()
    def portal_init(self, feats_sample, embed_matrix, k=512, min_norm=0.3):
        """Bures-Wasserstein (Gaussian optimal transport) initialisation after PORTAL (arXiv 2610.18598), as summarised in docs/research/vision_design.md s.3:
        W1 = V_t diag(sigma_t/sigma_s) V_s^T, b1 = mu_t - W1 mu_s, W2 = identity.  Source = shuffled patch tokens (3072-d) of an unlabeled image pool
        (feats_sample [N,576,768], 1-2k images are enough), target = LLM input-embedding rows (rows with norm < min_norm = reserved/unused are dropped).
        CAVEATS: written from the design-doc description, NOT from the paper's Appendix F code, and not validated against it; the GELU after layer 1 is kept
        (as in the summary), so layer-1 output is only approximately in the embedding distribution.  The 2-layer MLP has hidden=out=1536 so W2 = I fits.
        Use: conn.portal_init(feats[:1024].float(), model.get_input_embeddings().weight)"""
        dev = self.fc1.weight.device
        src = self.tokens(feats_sample.float().to(dev)).reshape(-1, self.fc1.in_features)
        tgt = embed_matrix.float().to(dev); tgt = tgt[tgt.norm(dim=1) > min_norm]

        def stats(X, k):
            mu = X.mean(0); Xc = X - mu
            k = min(k, X.shape[1], X.shape[0] - 1)
            U, S, V = torch.pca_lowrank(Xc, q=k, center=False, niter=3)
            sigma = (Xc @ V).std(0)                       # per-axis spread along each principal direction
            return mu, V, sigma
        mu_s, Vs, sg_s = stats(src, k); mu_t, Vt, sg_t = stats(tgt, k)
        kk = min(Vs.shape[1], Vt.shape[1])
        W1 = Vt[:, :kk] @ torch.diag(sg_t[:kk] / sg_s[:kk].clamp_min(1e-6)) @ Vs[:, :kk].T          # [1536, 3072]
        self.fc1.weight.copy_(W1); self.fc1.bias.copy_(mu_t - W1 @ mu_s)
        self.fc2.weight.copy_(torch.eye(self.fc2.out_features, self.fc2.in_features)); self.fc2.bias.zero_()
        return {"k": kk, "src_tokens": src.shape[0], "tgt_rows": tgt.shape[0]}


def n_params(m): return sum(p.numel() for p in m.parameters())


if __name__ == "__main__":
    c = Connector(); print("params", n_params(c))
    x = torch.randn(2, 576, 768); print(c(x).shape, c(x, low=True).shape, c(x).norm(dim=-1).mean().item())