"""
QED-Base-v2 inference library.

Usage:
    from qed_infer import load_model, load_tokenizer, run

    model = load_model("QED-Base-v2.pt")
    tokenizer = load_tokenizer("tok.model")

    text = run("Once upon a time", model, tokenizer, max_new_tokens=100)

    # streaming / batched:
    for texts in generate_stream(model, tokenizer, ["prompt A", "prompt B"]):
        ...  # texts[i] is the completion-so-far for prompt i
"""

from __future__ import annotations

import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Iterator, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
import sentencepiece as spm


DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32


@dataclass
class Config:
    vocab_size: int = 48000
    hidden_size: int = 768
    num_layers: int = 12
    num_heads: int = 12
    num_kv_heads: int = 4
    intermediate_size: int = 1792
    max_seq_len: int = 2048
    rope_theta: float = 10000.0
    rms_eps: float = 1e-6


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

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


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, max_seq_len, theta):
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
        positions = torch.arange(max_seq_len).float()
        freqs = torch.outer(positions, inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        self.register_buffer("cos", emb.cos(), persistent=False)
        self.register_buffer("sin", emb.sin(), persistent=False)

    def forward(self, q, k, offset: int):
        q_len = q.shape[-2]
        k_len = k.shape[-2]

        cos_q = self.cos[offset:offset + q_len][None, None, :, :].to(q.dtype)
        sin_q = self.sin[offset:offset + q_len][None, None, :, :].to(q.dtype)

        k_offset = offset + q_len - k_len
        cos_k = self.cos[k_offset:k_offset + k_len][None, None, :, :].to(k.dtype)
        sin_k = self.sin[k_offset:k_offset + k_len][None, None, :, :].to(k.dtype)

        return (
            q * cos_q + rotate_half(q) * sin_q,
            k * cos_k + rotate_half(k) * sin_k,
        )


class SwiGLU(nn.Module):
    def __init__(self, hidden, intermediate):
        super().__init__()
        self.gate_proj = nn.Linear(hidden, intermediate, bias=False)
        self.up_proj = nn.Linear(hidden, intermediate, bias=False)
        self.down_proj = nn.Linear(intermediate, hidden, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class GQAttention(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.num_heads = cfg.num_heads
        self.num_kv_heads = cfg.num_kv_heads
        self.head_dim = cfg.hidden_size // cfg.num_heads

        self.q_proj = nn.Linear(cfg.hidden_size, cfg.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(cfg.hidden_size, cfg.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(cfg.hidden_size, cfg.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=False)

        self.rope = RotaryEmbedding(self.head_dim, cfg.max_seq_len, cfg.rope_theta)

    def forward(self, x, offset: int, past_kv: Optional[tuple] = None):
        B, T, C = x.shape

        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)

        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)

        q, k = self.rope(q, k, offset)
        present = (k, v)

        repeat = self.num_heads // self.num_kv_heads
        k_rep = k.repeat_interleave(repeat, dim=1)
        v_rep = v.repeat_interleave(repeat, dim=1)

        y = F.scaled_dot_product_attention(q, k_rep, v_rep, is_causal=T > 1)

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.o_proj(y), present


class QEDBlock(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.hidden_size, cfg.rms_eps)
        self.attention = GQAttention(cfg)
        self.ffn_norm = RMSNorm(cfg.hidden_size, cfg.rms_eps)
        self.ffn = SwiGLU(cfg.hidden_size, cfg.intermediate_size)

    def forward(self, x, offset: int, past_kv=None):
        attn_out, present = self.attention(self.attn_norm(x), offset, past_kv)
        x = x + attn_out
        x = x + self.ffn(self.ffn_norm(x))
        return x, present


class QEDBaseV2(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
        self.layers = nn.ModuleList([QEDBlock(cfg) for _ in range(cfg.num_layers)])
        self.final_norm = RMSNorm(cfg.hidden_size, cfg.rms_eps)
        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.embed_tokens.weight

    def forward(self, input_ids, offset: int = 0, past_key_values: Optional[list] = None):
        x = self.embed_tokens(input_ids)
        new_past = []
        for i, layer in enumerate(self.layers):
            past_kv = past_key_values[i] if past_key_values is not None else None
            x, present = layer(x, offset, past_kv)
            new_past.append(present)
        x = self.final_norm(x)
        return self.lm_head(x), new_past


def load_model(model_path: str, cfg: Config = Config()) -> QEDBaseV2:
    path = Path(model_path)
    if not path.exists():
        raise FileNotFoundError(f"Checkpoint not found: {model_path}")

    package = torch.load(path, map_location="cpu", weights_only=True)
    state_dict = package["state_dict"] if "state_dict" in package else package

    model = QEDBaseV2(cfg)
    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    if missing:
        print(f"[warn] missing keys: {missing}", file=sys.stderr)
    if unexpected:
        print(f"[warn] unexpected keys: {unexpected}", file=sys.stderr)

    model.to(DEVICE, dtype=DTYPE)
    model.eval()

    name = package.get("Name", "QED-Base-v2") if isinstance(package, dict) else "QED-Base-v2"
    author = package.get("Author", "unknown") if isinstance(package, dict) else "unknown"
    print(f"Loaded {name} by {author} on {DEVICE} ({DTYPE})")
    return model


def load_tokenizer(tokenizer_path: str) -> spm.SentencePieceProcessor:
    if not Path(tokenizer_path).exists():
        raise FileNotFoundError(f"Tokenizer not found: {tokenizer_path}")
    tok = spm.SentencePieceProcessor()
    tok.load(tokenizer_path)
    return tok


def _apply_repetition_penalty(logits: torch.Tensor, generated: torch.Tensor, penalty: float):
    if penalty == 1.0:
        return logits
    for b in range(logits.shape[0]):
        seen = torch.unique(generated[b])
        vals = logits[b, seen]
        logits[b, seen] = torch.where(vals > 0, vals / penalty, vals * penalty)
    return logits


def _top_k_top_p_filter(logits: torch.Tensor, top_k: int, top_p: float):
    if top_k > 0:
        top_k = min(top_k, logits.size(-1))
        kth_val = torch.topk(logits, top_k, dim=-1).values[..., -1, None]
        logits = torch.where(logits < kth_val, torch.full_like(logits, float("-inf")), logits)

    if top_p < 1.0:
        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
        probs = F.softmax(sorted_logits, dim=-1)
        cum_probs = torch.cumsum(probs, dim=-1)

        remove = cum_probs > top_p
        remove[..., 1:] = remove[..., :-1].clone()
        remove[..., 0] = False

        sorted_logits[remove] = float("-inf")
        logits = torch.full_like(logits, float("-inf")).scatter(-1, sorted_idx, sorted_logits)

    return logits


@torch.no_grad()
def generate_stream(
    model: QEDBaseV2,
    tokenizer: spm.SentencePieceProcessor,
    prompts: list[str],
    max_new_tokens: int = 200,
    temperature: float = 0.8,
    top_k: int = 50,
    top_p: float = 0.95,
    repetition_penalty: float = 1.15,
    eos_id: Optional[int] = None,
) -> Iterator[list[str]]:
    if eos_id is None:
        eos_id = tokenizer.eos_id() if tokenizer.eos_id() >= 0 else None

    encoded = [tokenizer.encode(p) for p in prompts]
    max_len = max(len(e) for e in encoded)
    pad_id = tokenizer.pad_id() if tokenizer.pad_id() >= 0 else 0

    B = len(prompts)
    input_ids = torch.full((B, max_len), pad_id, dtype=torch.long, device=DEVICE)
    for i, e in enumerate(encoded):
        input_ids[i, max_len - len(e):] = torch.tensor(e, dtype=torch.long, device=DEVICE)

    generated = input_ids.clone()
    finished = torch.zeros(B, dtype=torch.bool, device=DEVICE)
    text_so_far = ["" for _ in range(B)]

    logits, past = model(input_ids, offset=0)
    offset = input_ids.shape[1]

    for _ in range(max_new_tokens):
        next_logits = logits[:, -1, :].float()
        next_logits = _apply_repetition_penalty(next_logits, generated, repetition_penalty)

        if temperature <= 0:
            next_token = next_logits.argmax(dim=-1, keepdim=True)
        else:
            next_logits = next_logits / temperature
            next_logits = _top_k_top_p_filter(next_logits, top_k, top_p)
            probs = F.softmax(next_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)

        next_token = torch.where(
            finished.unsqueeze(-1), torch.full_like(next_token, pad_id), next_token
        )
        generated = torch.cat([generated, next_token], dim=1)

        if eos_id is not None:
            finished |= next_token.squeeze(-1) == eos_id

        for i in range(B):
            if not finished[i]:
                text_so_far[i] = tokenizer.decode(generated[i].tolist())

        yield list(text_so_far)

        if bool(finished.all()):
            break

        logits, past = model(next_token, offset=offset, past_key_values=past)
        offset += 1


@torch.no_grad()
def run(
    prompt: str,
    model: QEDBaseV2,
    tokenizer: spm.SentencePieceProcessor,
    max_new_tokens: int = 200,
    temperature: float = 0.7,
    top_k: int = 40,
    top_p: float = 0.95,
    repetition_penalty: float = 1.15,
) -> str:
    final = ""
    for texts in generate_stream(
        model, tokenizer, [prompt],
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
    ):
        final = texts[0]
    return final