#!/usr/bin/env python3
# f3.py — self-contained Python runtime for FrameXlabs/fragment-2 (v3).
#
# Loads the released safetensors checkpoint + BPE tokenizer and answers typed
# questions with calibrated probabilities in a single forward pass.
# Five question types: choice / score / noul / multi / rank.
#
# Usage:
#     from f3 import Fragment
#     m = Fragment.from_pretrained("FrameXlabs/fragment-2")  # hub or local dir
#     res = m.decide(state="...", questions={...})
#
# Dependencies: torch, safetensors.
import json
import math
import os
import re
from typing import Dict

import torch
import torch.nn as nn
import torch.nn.functional as F


class F3Tokenizer:
    # Portable char-level BPE (the fragment build, proven since fragment-1).
    # Non-Latin scripts fall back to single characters.

    def __init__(self, path):
        spec = json.load(open(path))
        self.vocab = spec["vocab"]
        self.id_of = {v: i for i, v in enumerate(self.vocab)}
        self.ranks = {(a, b): i for i, (a, b) in enumerate(spec["merges"])}
        self.added = set(spec.get("added", []))
        self.pad = self.id_of["[PAD]"]
        self.cls = self.id_of["[CLS]"]
        self.sep = self.id_of["[SEP]"]
        self.unk = self.id_of["[UNK]"]
        self._cache = {}

    def _bpe(self, word):
        syms = list(word)
        while len(syms) >= 2:
            best, rank = None, None
            for a, b in zip(syms, syms[1:]):
                r = self.ranks.get((a, b))
                if r is not None and (rank is None or r < rank):
                    best, rank = (a, b), r
            if best is None:
                break
            a, b = best
            out, j = [], 0
            while j < len(syms):
                if j < len(syms) - 1 and syms[j] == a and syms[j + 1] == b:
                    out.append(a + b)
                    j += 2
                else:
                    out.append(syms[j])
                    j += 1
            syms = out
        return syms

    def _word_ids(self, word):
        cached = self._cache.get(word)
        if cached is not None:
            return cached
        ids = [self.id_of.get(s, self.unk) for s in self._bpe(word)]
        if len(self._cache) < 300_000:
            self._cache[word] = ids
        return ids

    def encode(self, text):
        text = text.lower()
        ids = []
        for w in re.findall(r"@?[a-z0-9]+|[^\sa-z0-9]", text):
            if w in self.added:
                ids.append(self.id_of[w])
            else:
                ids.extend(self._word_ids(w))
        return ids

class RMSNorm(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.w = nn.Parameter(torch.ones(d))

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * self.w


class RoPE(nn.Module):
    # Rotary position embeddings (theta=10000). Parameter-free.
    def __init__(self, d, max_len, base=10000.0):
        super().__init__()
        inv = 1.0 / (base ** (torch.arange(0, d, 2).float() / d))
        t = torch.arange(max_len).float()
        freqs = torch.outer(t, inv)                     # (L, d/2)
        self.register_buffer("cos", freqs.cos(), persistent=False)
        self.register_buffer("sin", freqs.sin(), persistent=False)

    def forward(self, x):                              # x: (B, H, L, D)
        L = x.shape[-2]
        cos, sin = self.cos[:L], self.sin[:L]           # (L, D/2)
        x1, x2 = x[..., ::2], x[..., 1::2]              # even/odd split
        cos = cos[None, None, :, :]
        sin = sin[None, None, :, :]
        out = torch.empty_like(x)
        out[..., ::2] = x1 * cos - x2 * sin
        out[..., 1::2] = x1 * sin + x2 * cos
        return out


class Block(nn.Module):
    # Pre-norm transformer block: MHA (RoPE) + SwiGLU FFN. No bias.
    def __init__(self, d, n_heads, ff):
        super().__init__()
        self.n1 = RMSNorm(d)
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.o = nn.Linear(d, d, bias=False)
        self.n2 = RMSNorm(d)
        self.gate = nn.Linear(d, ff, bias=False)
        self.up = nn.Linear(d, ff, bias=False)
        self.down = nn.Linear(ff, d, bias=False)
        self.d, self.h = d, n_heads
        self.rope = None

    def forward(self, x, att):
        # att: (B, L) True = real token (used as a broadcastable key mask)
        B, L, _ = x.shape
        h = self.n1(x)
        q, k, v = self.qkv(h).chunk(3, -1)
        q = q.view(B, L, self.h, -1).transpose(1, 2)    # (B,H,L,dh)
        k = k.view(B, L, self.h, -1).transpose(1, 2)
        v = v.view(B, L, self.h, -1).transpose(1, 2)
        q, k = self.rope(q), self.rope(k)
        a = F.scaled_dot_product_attention(q, k, v,
                                           attn_mask=att[:, None, None, :])
        x = x + self.o(a.transpose(1, 2).reshape(B, L, -1))
        h = self.n2(x)
        return x + self.down(F.silu(self.gate(h)) * self.up(h))


class F3Net(nn.Module):
    # Fragment-2 encoder: the original F2Net, scaled and modernized.
    # 10 layers, d=384, 6 heads, SwiGLU ff=1536, RoPE, RMSNorm, max_len 256.
    # Per-position scalar logit head (the proven fragment-1 design).
    def __init__(self, vocab=32768, d=384, max_len=256, n_layers=10,
                 n_heads=6, ff=1536):
        super().__init__()
        self.d, self.max_len = d, max_len
        self.scale = math.sqrt(d)
        self.tok = nn.Embedding(vocab, d, padding_idx=0)
        self.blocks = nn.ModuleList([Block(d, n_heads, ff) for _ in range(n_layers)])
        self.ln_f = RMSNorm(d)
        self.h1 = nn.Linear(d, d)
        self.h2 = nn.Linear(d, 1)
        self.rope = None
        self.cfg = {"vocab": vocab, "d": d, "max_len": max_len,
                    "n_layers": n_layers, "n_heads": n_heads, "ff": ff}
        self.apply(self._init)
        nn.init.normal_(self.tok.weight, std=0.02)
        nn.init.constant_(self.tok.weight[0], 0.0)

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)

    def _init_rope(self, device):
        # RoPE is parameter-free; created lazily and shared across blocks.
        # NOTE: rotations act on head_dim (d / n_heads), not model dim.
        if self.rope is None or self.rope.cos.device != device:
            r = RoPE(self.d // self.cfg["n_heads"], self.max_len).to(device)
            for blk in self.blocks:
                blk.rope = r
            self.rope = r

    def forward(self, ids):
        self._init_rope(ids.device)
        att = ids != 0
        B, L = ids.shape
        h = self.tok(ids) * self.scale
        for blk in self.blocks:
            h = blk(h, att)
        h = self.ln_f(h)
        return self.h2(F.gelu(self.h1(h))).squeeze(-1).float()


MAX_LEN = 256


def _confidence(p):
    k = len(p)
    if k < 2:
        return 1.0
    ent = -sum(x * math.log(max(x, 1e-12)) for x in p)
    return min(1.0, max(0.0, 1.0 - ent / math.log(k)))


class Fragment:
    def __init__(self, root):
        self.cfg = json.load(open(os.path.join(root, "f3_config.json")))
        self.tok = F3Tokenizer(os.path.join(root, "tokenizer.json"))
        from safetensors.torch import load_file
        sd = load_file(os.path.join(root, "model.safetensors"))
        sd = {k: (v.float() if v.dtype is not torch.float32 else v)
              for k, v in sd.items()}
        arch = self.cfg["architecture"]
        self.net = F3Net(arch["vocab"], arch["d"], arch["max_len"],
                         arch["n_layers"], arch["n_heads"], arch["ff"])
        self.net.load_state_dict(sd, strict=True)
        self.net.eval()
        self.version = self.cfg.get("version", "3.0")
        self.temps = self.cfg.get("temperatures", {})

    @classmethod
    def from_pretrained(cls, src="FrameXlabs/fragment-2"):
        if os.path.isdir(src):
            return cls(src)
        from huggingface_hub import snapshot_download
        return cls(snapshot_download(repo_id=src))

    def _options(self, q):
        t = q["type"]
        if t == "choice":
            crit = q.get("criteria") or {}
            return [f"@{i} {k}: {v}" if v else f"@{i} {k}"
                    for i, (k, v) in enumerate(crit.items())], list(crit.keys())
        if t in ("score", "multi", "rank"):
            crit = q.get("criteria") or []
            return [f"@{i} {c}" for i, c in enumerate(crit)], None
        return ["@no: no, the statement does not hold",
                "@yes: yes, the statement holds"], None

    @torch.no_grad()
    def decide(self, state, questions):
        import time as _t
        t0 = _t.time()
        answers = {}
        for qid, q in questions.items():
            opts, keys = self._options(q)
            pre = self.tok.encode(f"{q['instructions']} input:")
            opt_ids = self.tok.encode("options: " + "; ".join(opts))
            markers = [o.split()[0].rstrip(":") for o in opts]
            marker_ids = [self.tok.id_of[m] for m in markers]
            positions, seen = [], set()
            for idx, t in enumerate(opt_ids):
                if t in marker_ids and t not in seen:
                    seen.add(t)
                    positions.append(idx)
            budget = MAX_LEN - 2 - len(pre) - len(opt_ids)
            body = self.tok.encode(" " + state)[: max(8, budget)]
            ids = [self.tok.cls] + pre + body + opt_ids + [self.tok.sep]
            ids = ids[:MAX_LEN]
            shift = 1 + len(pre) + len(body)
            abs_pos = [p + shift for p in positions]
            abs_pos = [p for p in abs_pos if p < len(ids) - 1]
            t = torch.tensor([ids])
            logits = self.net(t)[0, abs_pos]
            temp = self.temps.get(q["type"], 1.0) or 1.0
            names = [(keys[i] if keys else f"option_{i}")
                     for i in range(len(positions))]
            base = {"type": q["type"], "question": q["instructions"]}
            if q["type"] == "choice":
                p = torch.softmax(logits / temp, -1).tolist()
                base["probabilities"] = {names[i]: round(p[i], 4)
                                         for i in range(len(p))}
                base["choice"] = names[int(max(range(len(p)),
                                               key=lambda i: p[i]))]
                base["confidence"] = round(_confidence(p), 4)
            elif q["type"] == "score":
                p = torch.softmax(logits / temp, -1).tolist()
                exp = sum(i * p[i] for i in range(len(p)))
                base["probabilities"] = {str(i): round(p[i], 4)
                                         for i in range(len(p))}
                base["score"] = round(exp, 4)
            elif q["type"] == "multi":
                p = torch.sigmoid(logits / temp).tolist()
                base["probabilities"] = {names[i]: round(p[i], 4)
                                         for i in range(len(p))}
                base["selected"] = [names[i] for i in range(len(p))
                                    if p[i] > 0.5]
            elif q["type"] == "rank":
                z = (logits / temp).tolist()
                order = sorted(range(len(z)), key=lambda i: -z[i])
                p = torch.softmax(logits / temp, -1).tolist()
                base["probabilities"] = {names[i]: round(p[i], 4)
                                         for i in range(len(p))}
                base["order"] = [names[i] for i in order]
            else:  # noul
                p = torch.softmax(logits / temp, -1).tolist()
                base["probabilities"] = {"false": round(p[0], 4),
                                         "true": round(p[1], 4)}
                base["noul"] = round(p[1], 4)
            answers[qid] = base
        return {"model": "fragment-2", "version": self.version,
                "answers": answers,
                "usage": {"latency_ms": round((_t.time() - t0) * 1000)}}