import torch
import sentencepiece as spm
from karModel import DictionaryTransformer, device

# -------------------------
# 1. Load SentencePiece tokenizer
# -------------------------
sp = spm.SentencePieceProcessor()
sp.load("cam/dict_spm.model")
print("✅ Loaded SentencePiece tokenizer")

# Check if eos exists in vocab
eos_id = sp.piece_to_id("<eos>")
if eos_id == -1:
    print("⚠️ No <eos> token in vocab — will remove by string replace")

def encode(s):
    return sp.encode(s, out_type=int)

def decode(ids):
    # Drop EOS ids if tokenizer has them
    if eos_id != -1:
        ids = [i for i in ids if i != eos_id]

    text = sp.decode(ids)

    # Trim at the first <eos> if it somehow appears in raw text
    if "<eos>" in text:
        text = text.split("<eos>")[0]

    return text.strip()


# -------------------------
# 2. Load trained model
# -------------------------
vocab_size = sp.get_piece_size()
model = DictionaryTransformer(vocab_size).to(device)
model.load_state_dict(torch.load("cam/dictionary_transformer.pt", map_location=device))
model.eval()
print("✅ Loaded trained DictionaryTransformer")

# -------------------------
# 3. Prompt + Generation
# -------------------------
def generate_definition(prompt, max_new_tokens=50, top_k=40):
    context = torch.tensor([encode(prompt)], dtype=torch.long, device=device)

    for _ in range(max_new_tokens):
        idx_cond = context[:, -model.position_embedding_table.num_embeddings:]
        logits, _ = model(idx_cond)
        logits = logits[:, -1, :]
        probs = torch.softmax(logits, dim=-1)

        # --- Top-k filtering ---
        if top_k is not None:
            values, indices = torch.topk(probs, k=top_k)
            probs = torch.zeros_like(probs).scatter_(1, indices, values)
            probs = probs / probs.sum(dim=-1, keepdim=True)

        next_id = torch.multinomial(probs, num_samples=1)
        context = torch.cat((context, next_id), dim=1)

        if eos_id != -1 and next_id.item() == eos_id:
            break

    return decode(context[0].tolist())


if __name__ == "__main__":
    while True:
        word = input("Enter a word (or 'quit'): ").strip()
        if word.lower() == "quit":
            break
        if not word.endswith(":"):
            word += ":"
        print("\n📖", generate_definition(word), "\n")