# -*- coding: utf-8 -*-
"""Self-contained Kenga instruct encoder. Byte UTF-8, bidirectional.

Usage (retrieval, Giga-style prefix on the query only):

    from modeling_kenga_embed import KengaEmbed
    m = KengaEmbed.from_pretrained("GermannM/kenga-embed-prophet-instruct")
    q = m.encode_queries(["Where is the capital of Russia?"])
    p = m.encode_documents(["Moscow is the capital of Russia."])

STS (prefix on both sentences):

    a = m.encode_sts(["The cat sits on the mat."])
    b = m.encode_sts(["A cat is sitting on a mat."])
"""
import json
import os
import torch
import torch.nn as nn
import torch.nn.functional as F

INSTR_RETR = "Instruct: Given a query, retrieve relevant passages\nQuery: "
INSTR_STS = "Instruct: Retrieve semantically similar text.\nQuery: "


class BL(nn.Module):
    def __init__(self, in_f, out_f, k):
        super().__init__()
        kk = min(k, in_f, out_f)
        self.U = nn.Parameter(torch.zeros(out_f, kk))
        self.V = nn.Parameter(torch.zeros(in_f, kk))
        self.S = nn.Parameter(torch.ones(kk))
        self.b = nn.Parameter(torch.zeros(out_f))
        self.register_buffer("mask", torch.ones(kk))

    def set_active(self, k_t):
        m = torch.zeros_like(self.mask)
        m[: min(int(k_t), self.mask.numel())] = 1
        self.mask.copy_(m)

    def forward(self, x):
        return ((x @ self.V) * (self.S * self.mask)) @ self.U.T + self.b


class BAttn(nn.Module):
    def __init__(self, d, heads, k_attn):
        super().__init__()
        self.h = heads
        self.dh = d // heads
        self.scale = self.dh ** -0.5
        self.q = BL(d, d, k_attn)
        self.k_ = BL(d, d, k_attn)
        self.v = BL(d, d, k_attn)
        self.o = BL(d, d, k_attn)

    def set_rank(self, ka):
        for bl in (self.q, self.k_, self.v, self.o):
            bl.set_active(ka)

    def forward(self, x, pad):
        B, T, _ = x.shape
        Q = self.q(x).view(B, T, self.h, self.dh).transpose(1, 2)
        K = self.k_(x).view(B, T, self.h, self.dh).transpose(1, 2)
        V = self.v(x).view(B, T, self.h, self.dh).transpose(1, 2)
        att = (Q.float() @ K.float().transpose(-2, -1)) * self.scale
        att = att.masked_fill(pad, float("-inf")).softmax(-1)
        return self.o((att @ V).transpose(1, 2).reshape(B, T, x.size(-1)))


class BBlock(nn.Module):
    def __init__(self, d, dff, heads, k_attn, k_ff):
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.ln2 = nn.LayerNorm(d)
        self.attn = BAttn(d, heads, k_attn)
        self.fc1 = BL(d, dff, k_ff)
        self.fc2 = BL(dff, d, k_ff)

    def set_rank(self, ka, kf):
        self.attn.set_rank(ka)
        self.fc1.set_active(kf)
        self.fc2.set_active(kf)

    def forward(self, x, pad):
        x = x + self.attn(self.ln1(x), pad)
        return x + self.fc2(F.gelu(self.fc1(self.ln2(x))))


class KengaEmbedEnc(nn.Module):
    def __init__(self, d=768, layers=8, heads=12, dff=3072, seq=160,
                 v=256, k_attn=192, k_ff=512):
        super().__init__()
        self.d = d
        self.seq = seq
        self.tok = nn.Embedding(v, d)
        self.pos = nn.Parameter(torch.zeros(1, seq, d))
        self.blocks = nn.ModuleList(
            [BBlock(d, dff, heads, k_attn, k_ff) for _ in range(layers)]
        )
        self.ln = nn.LayerNorm(d)
        self.proj = nn.Linear(d, d)

    def set_rank(self, ka, kf):
        for b in self.blocks:
            b.set_rank(ka, kf)

    def forward(self, idx, lengths=None):
        B, T = idx.shape
        x = self.tok(idx) + self.pos[:, :T]
        pad = idx.eq(0)
        mask = pad[:, None, None, :].expand(B, 1, T, T)
        for blk in self.blocks:
            x = blk(x, mask)
        x = self.ln(x)
        m = (~pad).unsqueeze(-1).float()
        pooled = (x * m).sum(1) / m.sum(1).clamp(min=1)
        return F.normalize(self.proj(pooled), dim=-1)


class KengaEmbed:
    def __init__(self, model, seq, instruct=True, device="cpu"):
        self.model = model.to(device).eval()
        self.seq = seq
        self.instruct = instruct
        self.device = device

    @classmethod
    def from_pretrained(cls, path, device="cpu"):
        cfg = json.load(open(os.path.join(path, "config.json"), encoding="utf-8"))
        m = KengaEmbedEnc(
            d=cfg.get("d", 768),
            layers=cfg.get("layers", 8),
            heads=cfg.get("heads", 12),
            dff=cfg.get("dff", 3072),
            seq=cfg.get("seq", 160),
            v=cfg.get("v", 256),
            k_attn=cfg.get("k_attn", 192),
            k_ff=cfg.get("k_ff", 512),
        )
        blob = torch.load(os.path.join(path, "pytorch_model.bin"),
                          map_location="cpu", weights_only=False)
        sd = blob["model"] if isinstance(blob, dict) and "model" in blob else blob
        m.load_state_dict(sd, strict=True)
        rank = cfg.get("rank") or blob.get("rank")
        if rank:
            m.set_rank(*list(rank))
        return cls(m, seq=cfg.get("seq", 160),
                   instruct=bool(cfg.get("instruct", True)), device=device)

    def _ids(self, text):
        raw = (text or "").encode("utf-8", errors="ignore")[: self.seq]
        ids = [b if b else 1 for b in raw][: self.seq]
        n = len(ids)
        ids = ids + [0] * (self.seq - n)
        return ids, n

    def _prefix(self, text, kind):
        if not self.instruct or not text:
            return text
        if text.startswith("Instruct:"):
            return text
        p = INSTR_STS if kind == "sts" else INSTR_RETR
        return p + text

    def encode(self, texts, kind="doc", bs=32):
        vecs = []
        with torch.no_grad():
            for i in range(0, len(texts), bs):
                chunk = texts[i : i + bs]
                ids, ln = [], []
                for t in chunk:
                    s = t if kind == "doc" else self._prefix(t, kind)
                    ii, n = self._ids(s)
                    ids.append(ii)
                    ln.append(n)
                x = torch.tensor(ids, device=self.device)
                l = torch.tensor(ln, device=self.device)
                vecs.append(self.model(x, l).cpu())
        return torch.cat(vecs, 0)

    def encode_queries(self, texts, **kw):
        return self.encode(texts, kind="retr", **kw)

    def encode_documents(self, texts, **kw):
        return self.encode(texts, kind="doc", **kw)

    def encode_sts(self, texts, **kw):
        return self.encode(texts, kind="sts", **kw)