import ollama
import numpy as np

MODEL = "nomic-embed-text"

pairs = [
    # Similar pairs
    ("Dairy-free pasta with chicken", "Creamy chicken pasta without milk", "similar"),
    ("Peanut butter cookies", "Sunflower butter cookies", "similar"),
    ("Replace milk with oat milk", "Use almond milk instead of regular milk", "similar"),
    # Dissimilar pairs
    ("Dairy-free mac and cheese", "Cheesy mac and cheese with cream", "dissimilar"),
    ("Gluten-free banana bread", "Spicy shrimp stir fry", "dissimilar"),
    ("Allergy-safe chocolate cake", "Peanut chicken curry", "dissimilar"),
    # Edge cases
    ("Vegan chocolate cake", "Vegan chocolate cake with almond milk and hazelnuts", "edge"),
    ("Use almond milk instead of regular milk", "Use oat milk instead of regular milk", "edge"),
    ("Gluten-free soy sauce chicken", "Gluten-free tamari chicken", "edge"),
]


def cosine_similarity(a, b):
    a, b = np.array(a), np.array(b)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))


print(f"{'Type':<12} {'Score':>6}  Pair")
print("-" * 80)

for text_a, text_b, pair_type in pairs:
    emb_a = ollama.embeddings(model=MODEL, prompt=text_a).embedding
    emb_b = ollama.embeddings(model=MODEL, prompt=text_b).embedding
    score = cosine_similarity(emb_a, emb_b)
    label = f"{text_a[:35]!r} vs {text_b[:35]!r}"
    print(f"{pair_type:<12} {score:>6.2f}  {label}")

print()
print("Key failure cases:")
print("  'almond milk vs oat milk' scored ~0.58 — dangerously close for a nut allergy user")
print("  'vegan cake vs vegan cake with hazelnuts' scored ~0.65 — hidden allergen not captured")