"""
Talk to Tallow, a 15.7M-parameter crow.

    pip install torch sentencepiece safetensors
    python tallow.py

Type to him. Empty line = quit. /new = start over as a stranger.
He sees the whole conversation every turn, so he remembers how you treated him (within ~1,900 tokens).

Everything needed is in this one file: the model, and the two text tricks his tokenizer expects
(a <nl> token for line breaks, and a space after every OPENING "*" or "[").
"""
import json
import math
import os
from dataclasses import dataclass

import sentencepiece as spm
import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file

HERE = os.path.dirname(os.path.abspath(__file__))
DEV = "cuda" if torch.cuda.is_available() else "cpu"


# ---------------------------------------------------------------- the model (Llama-style, written out by hand)
@dataclass
class Config:
    vocab: int = 4096
    ctx: int = 2048
    d: int = 384
    layers: int = 8
    heads: int = 6
    hidden: int = 1024


class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        xf = x.float()
        return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight).type_as(x)


def rope_tables(head_dim, ctx, base=10000.0):
    inv = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
    ang = torch.outer(torch.arange(ctx).float(), inv)
    ang = torch.cat([ang, ang], dim=-1)
    return ang.cos(), ang.sin()


def rotate(x, cos, sin):
    x1, x2 = x.chunk(2, dim=-1)
    return x * cos + torch.cat([-x2, x1], dim=-1) * sin


class Attention(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.heads, self.hd = c.heads, c.d // c.heads
        self.qkv = nn.Linear(c.d, 3 * c.d, bias=False)
        self.out = nn.Linear(c.d, c.d, bias=False)

    def forward(self, x, cos, sin):
        B, T, D = x.shape
        q, k, v = self.qkv(x).split(D, dim=2)
        q, k, v = (t.view(B, T, self.heads, self.hd).transpose(1, 2) for t in (q, k, v))
        q, k = rotate(q, cos, sin), rotate(k, cos, sin)
        a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.out(a.transpose(1, 2).reshape(B, T, D))


class SwiGLU(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.gate = nn.Linear(c.d, c.hidden, bias=False)
        self.up = nn.Linear(c.d, c.hidden, bias=False)
        self.down = nn.Linear(c.hidden, c.d, bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))


class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.norm1, self.attn = RMSNorm(c.d), Attention(c)
        self.norm2, self.mlp = RMSNorm(c.d), SwiGLU(c)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.norm1(x), cos, sin)
        return x + self.mlp(self.norm2(x))


class Tallow(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.c = c
        self.embed = nn.Embedding(c.vocab, c.d)
        self.blocks = nn.ModuleList(Block(c) for _ in range(c.layers))
        self.norm = RMSNorm(c.d)
        cos, sin = rope_tables(c.d // c.heads, c.ctx)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)

    def forward(self, idx):
        T = idx.shape[1]
        x = self.embed(idx)
        cos, sin = self.cos[:T].to(x.dtype), self.sin[:T].to(x.dtype)
        for b in self.blocks:
            x = b(x, cos, sin)
        return self.norm(x) @ self.embed.weight.T      # tied embeddings


# ---------------------------------------------------------------- the tokenizer and its two text tricks
NL = "<nl>"


def to_sp(text):
    """Line breaks -> <nl>; a space after every OPENING '*' (1st, 3rd... on a line) and every '['."""
    lines = []
    for line in text.split("\n"):
        out, stars = [], 0
        for i, ch in enumerate(line):
            out.append(ch)
            stars += ch == "*"
            opener = ch == "[" or (ch == "*" and stars % 2 == 1)
            if opener and i + 1 < len(line) and not line[i + 1].isspace():
                out.append(" ")
        lines.append("".join(out))
    return (NL + " ").join(lines)


def from_sp(text):
    text = text.replace(NL + " ", "\n").replace(NL, "\n")
    lines = []
    for line in text.split("\n"):
        out, stars, skip = [], 0, False
        for i, ch in enumerate(line):
            if skip:
                skip = False
                continue
            out.append(ch)
            stars += ch == "*"
            opener = ch == "[" or (ch == "*" and stars % 2 == 1)
            if opener and line[i + 1:i + 2] == " " and i + 2 < len(line) and not line[i + 2].isspace():
                skip = True
        lines.append("".join(out))
    return "\n".join(lines)


def render(sp, history):
    """you: ...<nl>tallow: ...<nl>  (the exact format he was trained on)"""
    nl = sp.piece_to_id(NL)
    ids = []
    for turn in history:
        if turn["role"] == "user":
            ids += sp.encode(to_sp(f"you: {turn['text']}")) + [nl]
        else:
            ids += sp.encode(to_sp("tallow:")) + sp.encode(to_sp(turn["text"])) + [nl]
    return ids


@torch.no_grad()
def reply(model, sp, history, temperature=0.8, top_k=50, max_new=80):
    while len(history) > 1 and len(render(sp, history)) > 1900:   # forget the oldest lines if too long
        history = history[2:]
    ids = render(sp, history) + sp.encode(to_sp("tallow:"))
    start, stop = len(ids), {sp.piece_to_id(NL), sp.eos_id()}
    for _ in range(max_new):
        logits = model(torch.tensor([ids[-model.c.ctx:]], device=DEV))[0, -1].float() / temperature
        logits[logits < torch.topk(logits, top_k).values[-1]] = -float("inf")
        nxt = int(torch.multinomial(F.softmax(logits, -1), 1))
        if nxt in stop:
            break
        ids.append(nxt)
    return from_sp(sp.decode(ids[start:])).strip()


def load():
    with open(os.path.join(HERE, "config.json"), encoding="utf-8") as f:
        model = Tallow(Config(**json.load(f)))
    model.load_state_dict(load_file(os.path.join(HERE, "model.safetensors")))
    model.to(DEV).eval()
    return model, spm.SentencePieceProcessor(model_file=os.path.join(HERE, "tallow_tok.model"))


if __name__ == "__main__":
    model, sp = load()
    history = []
    print("Tallow is on the power line. (empty line = quit, /new = start over)\n")
    while True:
        you = input("you: ").strip()
        if not you:
            break
        if you == "/new":
            history = []
            print("(he doesn't know you anymore)\n")
            continue
        history.append({"role": "user", "text": you})
        line = reply(model, sp, history)
        history.append({"role": "tallow", "text": line})
        print(f"tallow: {line}\n")