from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# Load MRR-optimized model
model = SentenceTransformer('ThanhLe0125/ebd-math')

# Example: Find exact chunk with MRR optimization
query = "query: Định nghĩa hàm số đồng biến"
chunks = [
    "passage: Hàm số đồng biến trên khoảng (a;b) là hàm số mà với mọi x1 < x2 thì f(x1) < f(x2)",
    "passage: Ví dụ: Tìm khoảng đồng biến của hàm số y = x^2 - 2x + 1", 
    "passage: Phương trình bậc hai ax^2 + bx + c = 0 có delta = b^2 - 4ac",
    "passage: Tính đạo hàm của hàm số đa thức",
    "passage: Giới hạn của dãy số"
]

# Encode and calculate similarities
query_emb = model.encode([query])
chunk_embs = model.encode(chunks)
similarities = cosine_similarity(query_emb, chunk_embs)[0]

# MRR-optimized ranking
ranked_indices = similarities.argsort()[::-1]

print("MRR-Optimized Rankings:")
for rank, idx in enumerate(ranked_indices, 1):
    chunk_type = "CORRECT" if idx == 0 else ("RELATED" if idx == 1 else "IRRELEVANT")
    print(f"Rank {rank}: {chunk_type} (Score: {similarities[idx]:.4f})")
    print(f"   {chunks[idx][:80]}...")
    print()

# Calculate MRR for this query
correct_rank = None
for rank, idx in enumerate(ranked_indices, 1):
    if idx == 0:  # First chunk is correct
        correct_rank = rank
        break

if correct_rank:
    mrr_score = 1.0 / correct_rank
    print(f"This query MRR: {mrr_score:.4f} (correct chunk at rank #{correct_rank})")
    
    if mrr_score >= 0.5:
        print("✅ Good! Correct chunk in top 2 positions")
    else:
        print("⚠️ Could be better - correct chunk should be higher")

# Efficient inference (common use case)
print("\nEfficient Inference (top 3 only):")
top_3_indices = ranked_indices[:3]
for i, idx in enumerate(top_3_indices, 1):
    print(f"{i}. Score: {similarities[idx]:.4f} - {chunks[idx][:60]}...")