"""Tiny random model with logit-lens heads, for local UI development on CPU: the real dssd
engine runs end to end (exact verification, copies, head, star), on gibberish text."""

import torch
from transformers import Qwen3Config, Qwen3ForCausalLM

from dssd import DSSD
from dssd.heads import ExitHead

LAYERS = [2, 4, 6]
_CHARS = [chr(c) for c in range(32, 127)] + ["\n", "\t"]  # 97 ids
_TREE = {"copy": {"m": 11, "n_max": 4, "n_min": 1}, "head_budget": 4,
         "star": {"head": 0, "C": 4, "tau": [0.0, 0.0, 0.0, 0.0], "gamma": 0.0}}
_POLICY = {"type": "threshold", "thresholds": {"0": 0.0, "2": 0.0}, "cum_gamma": {"0": 0.0}, "max_draft": 15,
           "tree": _TREE}


class CharTokenizer:
    eos_token_id = None

    def encode(self, text, add_special_tokens=False):
        return [_CHARS.index(c) if c in _CHARS else _CHARS.index("?") for c in text]

    def decode(self, ids, skip_special_tokens=True):
        return "".join(_CHARS[i % len(_CHARS)] for i in ids)


def tiny_engine() -> DSSD:
    torch.manual_seed(0)
    cfg = Qwen3Config(vocab_size=len(_CHARS), hidden_size=64, intermediate_size=128, num_hidden_layers=8,
                      num_attention_heads=4, num_key_value_heads=2, head_dim=16, max_position_embeddings=2048,
                      tie_word_embeddings=False)
    m = Qwen3ForCausalLM(cfg).eval()
    with torch.no_grad():
        for p in m.parameters():
            if p.dim() > 1:
                p.mul_(4.0)
        for layer in m.model.layers[5:]:  # deep layers do nothing: the heads are often right
            layer.self_attn.o_proj.weight.zero_()
            layer.mlp.down_proj.weight.zero_()
    heads = []
    for _ in LAYERS:
        full = ExitHead(64, cfg.vocab_size, cfg.rms_norm_eps)
        with torch.no_grad():
            full.norm.weight.copy_(m.model.norm.weight)
            full.proj.weight.copy_(m.lm_head.weight)
        heads.append(ExitHead.from_full(full, torch.arange(cfg.vocab_size), reject=True))
    config = {"base_model": "tiny", "default_domain": "web",
              "domains": {"web": {"policy": _POLICY, "stop": None},
                          "gsm8k": {"policy": _POLICY, "stop": "gsm8k"},
                          "code": {"policy": _POLICY, "stop": "humaneval"}}}
    return DSSD.from_components(m, CharTokenizer(), heads, LAYERS, config, max_len=2048)