#!/usr/bin/env python3
"""Re-encode every card in the local Scryfall bulk file into mechanics embeddings.

Mirrors minimaxir/mtg-embeddings: same model, same "prettified JSON of the card's
mechanical fields" input, same unit-normalized output -- but built from OUR
oracle-cards.jsonl, so it is current and includes new commanders.

    python3 tools/embed.py            # -> data/embeddings.npz
"""
import json, os, sys, numpy as np
from sentence_transformers import SentenceTransformer

ROOT  = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MODEL = "Alibaba-NLP/gte-modernbert-base"
SRC   = os.path.join(ROOT, "data", "oracle-cards.jsonl")
OUT   = os.path.join(ROOT, "data", "embeddings.npz")

def card_text(c):
    """The mechanical identity of a card, as prettified JSON. No flavor, no art."""
    faces = c.get("card_faces") or []
    # NOTE: name is deliberately excluded -- including it makes the model match on
    # shared words in titles ("Reaper", "Massacre") instead of on mechanics.
    d = {
        "mana_cost": c.get("mana_cost") or " // ".join(f.get("mana_cost","") for f in faces),
        "type": c.get("type_line"),
        "text": c.get("oracle_text") or "\n//\n".join(f.get("oracle_text","") for f in faces),
        "power": c.get("power"), "toughness": c.get("toughness"),
        "loyalty": c.get("loyalty"),
    }
    return json.dumps({k: v for k, v in d.items() if v}, indent=2)

def main():
    cards = [json.loads(l) for l in open(SRC)]
    # skip tokens/art series -- they have no mechanics worth matching
    cards = [c for c in cards if c.get("layout") not in ("token","double_faced_token","art_series","emblem")]
    names = [c["name"] for c in cards]
    texts = [card_text(c) for c in cards]
    print(f"encoding {len(texts)} cards with {MODEL}", file=sys.stderr)

    m = SentenceTransformer(MODEL, device="cuda", trust_remote_code=True)
    emb = m.encode(texts, batch_size=64, normalize_embeddings=True,
                   show_progress_bar=True, convert_to_numpy=True).astype("float32")

    np.savez_compressed(OUT, names=np.array(names, dtype=object), embedding=emb)
    print(f"wrote {OUT}  shape={emb.shape}", file=sys.stderr)

if __name__ == "__main__":
    main()