"""Run from the downloaded release: python example.py [--device cuda:0]."""

import argparse
import sys
from pathlib import Path

import numpy as np

# HF snapshots use symlinks; retain the snapshot directory for sibling imports.
sys.path.insert(0, str(Path(__file__).absolute().parent))
from qime import QIMEModel


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", default=None)
    parser.add_argument("--local-files-only", action="store_true")
    args = parser.parse_args()
    model = QIMEModel.from_pretrained(
        Path(__file__).parent, device=args.device, local_files_only=args.local_files_only,
    )
    texts = [
        "Patient presents with severe chest pain.",
        "Treatment involves daily insulin injections.",
    ]
    embeddings = model.encode(texts)
    assert embeddings.shape == (2, 8855)
    assert np.all(embeddings.sum(axis=1) == 256)
    assert np.isin(embeddings, [0, 1]).all()
    print("Embedding shape:", embeddings.shape)
    print("Active dimensions per text:", embeddings.sum(axis=1))
    for text, embedding in zip(texts, embeddings):
        print("\nText:", text)
        print("First 10 active coordinates (coordinate order, not relevance order):")
        for index in np.flatnonzero(embedding)[:10]:
            print(f"[{index}] {model.questions[index]}")


if __name__ == "__main__":
    main()