"""Trit sentence embedder: 3 post-norm BERT layers, width 384, BitNet b1.58 weights, int8 activations.

Reference inference code for the packed release format (config "format": "trit-packed-v1"):
- every linear layer inside the encoder is stored as 2-bit ternary codes plus one bf16 scale per output row;
- the word and position embedding tables are int8 with one fp32 scale per row;
- LayerNorms, biases, the token-type embedding and the 384x384 output projection are fp32.

    from trit_embed import TritEmbed
    model = TritEmbed.from_pretrained(".")
    vectors = model.encode(["how do index funds work", "Index funds track a market index."])  # (2, 384), unit norm

PyTorch has no ternary matmul, so this decodes the layers to fp32 once at load and emulates the int8 activation
rounding in float. The NPU encoder (phlegm, open_kernels/designs/trit_embed) runs the 2-bit codes natively.
"""
import json
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

GROUP = 256  # channels per activation scale in the NPU kernels


def ste(x: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
    return x + (q - x).detach()


def quant_act(x: torch.Tensor) -> torch.Tensor:
    """int8 per token (per 256-channel block when the width divides), power-of-two scale."""
    shape = x.shape
    if shape[-1] > GROUP and shape[-1] % GROUP == 0:
        x = x.reshape(*shape[:-1], shape[-1] // GROUP, GROUP)
    m = x.detach().float().abs().amax(-1, keepdim=True).clamp(min=1e-5)
    f, k = torch.frexp(m)
    s = torch.ldexp(torch.ones_like(m), k - 7 + (2 * f > 127 / 64).to(k.dtype))
    q = ((x.detach().float() / s).round().clamp(-128, 127) * s).to(x.dtype)
    return ste(x, q).reshape(shape)


def pack_codes(codes: torch.Tensor) -> torch.Tensor:
    """Ternary codes [out, in] -> uint8 [out, in/4]: code + 1 in {0, 1, 2}, column 4b + j at bits 2j..2j+1."""
    c = (codes.to(torch.int16) + 1).to(torch.uint8).reshape(codes.shape[0], -1, 4)
    return c[..., 0] | (c[..., 1] << 2) | (c[..., 2] << 4) | (c[..., 3] << 6)


def unpack_codes(packed: torch.Tensor) -> torch.Tensor:
    """uint8 [out, in/4] -> int8 ternary codes [out, in]."""
    shifts = torch.tensor([0, 2, 4, 6], dtype=torch.uint8)
    c = (packed.unsqueeze(-1) >> shifts) & 3
    if (c == 3).any():
        raise ValueError("Reserved 2-bit ternary code 3")
    return (c.to(torch.int8) - 1).reshape(packed.shape[0], -1)


class TernLinear(nn.Module):
    """y = int8(x) @ (codes * scale).T + bias. codes in {-1, 0, 1}; scale is one bf16 value per output row."""

    def __init__(self, n_in: int, n_out: int):
        super().__init__()
        self.register_buffer("codes", torch.zeros(n_out, n_in, dtype=torch.int8))
        self.register_buffer("scale", torch.zeros(n_out, dtype=torch.bfloat16))
        self.register_buffer("weight", torch.zeros(n_out, n_in), persistent=False)
        self.bias = nn.Parameter(torch.zeros(n_out))

    def decode(self) -> None:
        self.weight = self.codes.float() * self.scale.float()[:, None]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.linear(quant_act(x), self.weight, self.bias)


class Int8Embedding(nn.Module):
    """Lookup table stored as int8 rows with one fp32 scale per row."""

    def __init__(self, rows: int, dim: int):
        super().__init__()
        self.register_buffer("qweight", torch.zeros(rows, dim, dtype=torch.int8))
        self.register_buffer("scale", torch.zeros(rows))

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        return self.qweight[ids].float() * self.scale[ids].unsqueeze(-1)


class Layer(nn.Module):
    def __init__(self, d: int, heads: int, ffn: int, eps: float, gelu: str):
        super().__init__()
        self.heads, self.gelu = heads, gelu
        self.q, self.k, self.v, self.o = (TernLinear(d, d) for _ in range(4))
        self.ln1 = nn.LayerNorm(d, eps=eps)
        self.fc1, self.fc2 = TernLinear(d, ffn), TernLinear(ffn, d)
        self.ln2 = nn.LayerNorm(d, eps=eps)

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        b, t, d = x.shape
        q, k, v = (m(x).view(b, t, self.heads, -1).transpose(1, 2) for m in (self.q, self.k, self.v))
        attention = F.scaled_dot_product_attention(q, k, v, attn_mask=mask[:, None, None, :])
        x = self.ln1(x + self.o(attention.transpose(1, 2).reshape(b, t, d)))
        return self.ln2(x + self.fc2(F.gelu(self.fc1(x), approximate=self.gelu)))


class TritEmbed(nn.Module):
    def __init__(self, vocab, emb, d, heads, ffn, layers, max_pos, eps, out_dim, loops=1, pool="mean", gelu="tanh"):
        super().__init__()
        assert emb == d and loops == 1 and pool == "mean", "this release is a single-pass, mean-pooled model"
        self.max_pos = max_pos
        self.word = Int8Embedding(vocab, emb)
        self.pos = Int8Embedding(max_pos, emb)
        self.tok_type = nn.Embedding(2, emb)
        self.ln_e = nn.LayerNorm(emb, eps=eps)
        self.layers = nn.ModuleList(Layer(d, heads, ffn, eps, gelu) for _ in range(layers))
        self.out = nn.Linear(d, out_dim, bias=False)
        self.tokenizer = None

    def embed(self, ids: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
        """Input states: word + position + token-type embeddings, then LayerNorm."""
        return self.ln_e(self.word(ids) + self.pos(positions) + self.tok_type.weight[0])

    def forward(self, ids: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        x = self.embed(ids, torch.arange(ids.size(1), device=ids.device)[None])
        for layer in self.layers:
            x = layer(x, mask)
        m = mask.unsqueeze(-1).to(x.dtype)
        return F.normalize(self.out((x * m).sum(1) / m.sum(1).clamp(min=1)), dim=-1)

    def ternary_layers(self):
        """(name, module) for every ternary linear, in encoder order."""
        return [(n, m) for n, m in self.named_modules() if isinstance(m, TernLinear)]

    @classmethod
    def from_pretrained(cls, path) -> "TritEmbed":
        from safetensors.torch import load_file
        from transformers import AutoTokenizer
        path = Path(path)
        config = json.loads((path / "config.json").read_text())
        if config.get("format") != "trit-packed-v1":
            raise ValueError(f"Unknown weight format {config.get('format')!r}")
        model = cls(**config["dims"])
        state = load_file(path / "model.safetensors")
        for name, _ in model.ternary_layers():
            state[f"{name}.codes"] = unpack_codes(state[f"{name}.codes"])
        model.load_state_dict(state)
        for _, m in model.ternary_layers():
            m.decode()
        model.tokenizer = AutoTokenizer.from_pretrained(path)
        return model.eval()

    @torch.inference_mode()
    def encode(self, texts: list[str], batch: int = 128) -> np.ndarray:
        """Unit-norm float32 vectors; inputs are truncated to 128 tokens. Queries need no prefix.
        Runs on the model's device: model.to("cuda") for a GPU."""
        device = self.out.weight.device
        tokens = self.tokenizer(texts, truncation=True, max_length=self.max_pos)["input_ids"]
        order = sorted(range(len(tokens)), key=lambda i: len(tokens[i]))
        result = np.empty((len(tokens), self.out.out_features), dtype=np.float32)
        for start in range(0, len(order), batch):
            indexes = order[start:start + batch]
            enc = self.tokenizer.pad({"input_ids": [tokens[i] for i in indexes]}, padding=True, return_tensors="pt")
            vectors = self(enc["input_ids"].to(device), enc["attention_mask"].bool().to(device))
            result[indexes] = vectors.float().cpu().numpy()
        return result